Add partial block reduce cuda kernel (#68122)

BiynXu · web-flow · commit 745a8b93a272 · 2024-09-12T10:44:50.000+08:00
diff --git a/paddle/cinn/hlir/pe/reduction.cc b/paddle/cinn/hlir/pe/reduction.cc
@@ -1144,26 +1144,26 @@ std::string CrossThreadReduceExternalFuncName(const ir::Expr& op,
       ::common::errors::InvalidArgument("Tensor is not a tensor!"));
   if (op.As<ir::Add>()) {
     if (tensor.as_tensor()->type().is_bool()) {
-      return "cinn_block_reduce_any_internal_shm";
+      return "cinn_partial_block_reduce_any_internal_shm";
     }
-    return "cinn_block_reduce_sum" +
+    return "cinn_partial_block_reduce_sum" +
            Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
   } else if (op.As<ir::Mul>()) {
     if (tensor.as_tensor()->type().is_bool()) {
-      return "cinn_block_reduce_all_internal_shm";
+      return "cinn_partial_block_reduce_all_internal_shm";
     }
-    return "cinn_block_reduce_prod" +
+    return "cinn_partial_block_reduce_prod" +
            Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
   } else if (op.As<ir::Max>()) {
-    return "cinn_block_reduce_max" +
+    return "cinn_partial_block_reduce_max" +
            Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
   } else if (op.As<ir::Min>()) {
-    return "cinn_block_reduce_min" +
+    return "cinn_partial_block_reduce_min" +
            Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
   } else if (op.As<ir::And>()) {
-    return "cinn_block_reduce_all_internal_shm";
+    return "cinn_partial_block_reduce_all_internal_shm";
   } else if (op.As<ir::Or>()) {
-    return "cinn_block_reduce_any_internal_shm";
+    return "cinn_partial_block_reduce_any_internal_shm";
   } else {
     std::stringstream ss;
     ss << "Reduce type: " << op << " Not supported yet!";
diff --git a/paddle/cinn/optim/replace_cross_thread_reduction_test.cc b/paddle/cinn/optim/replace_cross_thread_reduction_test.cc
@@ -71,7 +71,7 @@ TEST(CrossThreadReductionReplacer, basic) {
         ScheduleBlock(B)
         {
           i0_0, i1 = axis.bind(i, reduce_j)
-          B[i0_0] = cinn_block_reduce_sum_fp32_internal_shm(A[i0_0, i1], _Buffer_<cinn_buffer_t*: 32>(shm32__fp32_reduce), false)
+          B[i0_0] = cinn_partial_block_reduce_sum_fp32_internal_shm(A[i0_0, i1], _Buffer_<cinn_buffer_t*: 32>(shm32__fp32_reduce), false)
         }
       }
     }
diff --git a/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh b/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh
@@ -651,6 +651,47 @@ EXPAND_REDUCE_FP16_MACRO(CINN_DISCRETE_REDUCE_INTERNAL_SHM_MACRO)
 #undef CINN_DISCRETE_REDUCE_INTERNAL_SHM_IMPL
 #undef CINN_DISCRETE_REDUCE_INTERNAL_SHM_MACRO
 
+#define CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_IMPL(TYPE, value, init_value, cinn_warp_shuffle_internal) \
+  int tid = threadIdx.y * blockDim.x + threadIdx.x;                                                  \
+  int warp_id = tid >> 5;                                                                            \
+  int row_dim =  (blockDim.x + 31) >> 5;                                                             \
+  TYPE tmp_val = cinn_warp_shuffle_internal(value);                                                  \
+  if (blockDim.x <= 32) {                                                                            \
+    return tmp_val;                                                                                  \
+  }                                                                                                  \
+  __syncthreads();                                                                                   \
+  if ((tid & 31) == 0) {                                                                             \
+    shm[warp_id] = tmp_val;                                                                          \
+  }                                                                                                  \
+  __syncthreads();                                                                                   \
+  if (threadIdx.x < 32) {                                                                            \
+    tmp_val = (threadIdx.x < row_dim) ? shm[threadIdx.y * row_dim + threadIdx.x] : init_value;       \
+    shm[warp_id] = cinn_warp_shuffle_internal(tmp_val);                                              \
+  }                                                                                                  \
+  __syncthreads();                                                                                   \
+  return shm[threadIdx.y * row_dim];
+
+#define CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_MACRO(REDUCE_TYPE, INITIAL_VALUE, DTYPE)                                                                \
+  __device__ inline DTYPE cinn_partial_block_reduce_##REDUCE_TYPE##_internal_shm(const DTYPE value, DTYPE* shm, bool return_warp = false) {            \
+    CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_IMPL(DTYPE, value, (DTYPE)(INITIAL_VALUE), cinn_warp_shuffle_##REDUCE_TYPE##_internal);                      \
+  }
+
+EXPAND_REDUCE_INT32_MARCO(CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+EXPAND_REDUCE_INT64_MARCO(CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+EXPAND_REDUCE_FP32_MACRO(CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+EXPAND_REDUCE_FP64_MACRO(CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+EXPAND_REDUCE_BOOL_MACRO(CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+
+#ifdef CINN_CUDA_BF16
+EXPAND_REDUCE_BF16_MACRO(CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+#endif
+
+#ifdef CINN_CUDA_FP16
+EXPAND_REDUCE_FP16_MACRO(CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+#endif
+
+#undef CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_IMPL
+#undef CINN_PARTIAL_BLOCK_REDUCE_INTERNAL_SHM_MACRO
 
 #define CINN_BLOCK_REDUCE_IMPL(REDUCE_TYPE, INITIAL_VALUE, DTYPE)                                     \
   __device__ inline DTYPE cinn_block_reduce_##REDUCE_TYPE(const DTYPE *buf, int offset, int extend) { \
diff --git a/paddle/cinn/runtime/cuda/cuda_intrinsics_reduce.cc b/paddle/cinn/runtime/cuda/cuda_intrinsics_reduce.cc
@@ -127,6 +127,24 @@ CINN_REGISTER_HELPER(cuda_intrinsics_reduce) {
 
 #undef REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL
 
+#define REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL(REDUCE_TYPE, DTYPE)  \
+  REGISTER_FACKED_EXTERN_FUNC_HELPER(                                 \
+      cinn_partial_block_reduce_##REDUCE_TYPE##_internal_shm, target) \
+      .SetRetType<DTYPE>()                                            \
+      .AddInputType<DTYPE>()                                          \
+      .AddInputType<cinn_buffer_t *>()                                \
+      .End();
+
+  EXPAND_REDUCE_INT32_REGISTER_MARCO(REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL)
+  EXPAND_REDUCE_INT64_REGISTER_MARCO(REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL)
+  EXPAND_REDUCE_BF16_REGISTER_MACRO(REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL)
+  EXPAND_REDUCE_FP16_REGISTER_MACRO(REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL)
+  EXPAND_REDUCE_FP32_REGISTER_MACRO(REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL)
+  EXPAND_REDUCE_FP64_REGISTER_MACRO(REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL)
+  EXPAND_REDUCE_BOOL_REGISTER_MACRO(REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL)
+
+#undef REGISTER_BLOCK_REDUCE_INTERNAL_FUNC_IMPL
+
 #define REGISTER_DISCRETE_REDUCE_INTERNAL_FUNC_IMPL(REDUCE_TYPE, DTYPE) \
   REGISTER_FACKED_EXTERN_FUNC_HELPER(                                   \
       cinn_discrete_reduce_##REDUCE_TYPE##_internal_shm, target)        \

Original file line number	Diff line number	Diff line change
`@@ -71,7 +71,7 @@ TEST(CrossThreadReductionReplacer, basic) {`
`71`	`71`	`ScheduleBlock(B)`
`72`	`72`	`{`
`73`	`73`	`i0_0, i1 = axis.bind(i, reduce_j)`
`74`		`- B[i0_0] = cinn_block_reduce_sum_fp32_internal_shm(A[i0_0, i1], _Buffer_<cinn_buffer_t*: 32>(shm32__fp32_reduce), false)`
	`74`	`+ B[i0_0] = cinn_partial_block_reduce_sum_fp32_internal_shm(A[i0_0, i1], _Buffer_<cinn_buffer_t*: 32>(shm32__fp32_reduce), false)`
`75`	`75`	`}`
`76`	`76`	`}`
`77`	`77`	`}`