【CINN】Optimize block reduce (#58344)

Courtesy-Xs · web-flow · commit 8c1bcfe22efe · 2023-11-08T00:32:23.000-06:00
* optimize block reduce

* fix var name

* add a buffer as a arg for extern call

* fix cinn_block_reduce because of coupling in old reduce schedule logic

* fix replace_cross_thread_reduction test case

* fix lost buffer of shared

* fix duplicated buffer error

* fix merge conflict
diff --git a/paddle/cinn/hlir/pe/reduction.cc b/paddle/cinn/hlir/pe/reduction.cc
@@ -1082,20 +1082,20 @@ std::string CrossThreadReduceExternalFuncName(const ir::Expr& op,
   CHECK_NOTNULL(tensor.as_tensor());
   if (op.As<ir::Add>()) {
     return "cinn_block_reduce_sum" +
-           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal";
+           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
   } else if (op.As<ir::Mul>()) {
     return "cinn_block_reduce_prod" +
-           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal";
+           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
   } else if (op.As<ir::Max>()) {
     return "cinn_block_reduce_max" +
-           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal";
+           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
   } else if (op.As<ir::Min>()) {
     return "cinn_block_reduce_min" +
-           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal";
+           Type2StrForReduce(tensor.as_tensor()->type()) + "_internal_shm";
   } else if (op.As<ir::And>()) {
-    return "cinn_block_reduce_all_internal";
+    return "cinn_block_reduce_all_internal_shm";
   } else if (op.As<ir::Or>()) {
-    return "cinn_block_reduce_any_internal";
+    return "cinn_block_reduce_any_internal_shm";
   } else {
     LOG(FATAL) << "Reduce type: " << op << " Not supported yet!";
   }
diff --git a/paddle/cinn/hlir/pe/reduction.h b/paddle/cinn/hlir/pe/reduction.h
@@ -470,6 +470,8 @@ std::vector<ir::Tensor> TwoStepBlockReduceAny(
 
 std::string CrossThreadReduceExternalFuncName(const ir::Expr& op,
                                               const ir::Expr& tensor);
+
+std::string Type2StrForReduce(common::Type type);
 }  // namespace pe
 }  // namespace hlir
 }  // namespace cinn
diff --git a/paddle/cinn/optim/replace_cross_thread_reduction.cc b/paddle/cinn/optim/replace_cross_thread_reduction.cc
@@ -31,7 +31,15 @@ namespace cinn {
 namespace optim {
 namespace {
 
-struct CrossThreadReductionReplacer : public ir::IRMutator<Expr*> {
+struct BufferCmp {
+  bool operator()(const ir::Buffer& a, const ir::Buffer& b) const {
+    if (a->name == b->name) return false;
+    return true;
+  }
+};
+
+thread_local std::set<ir::Buffer, BufferCmp> shm_buffer_;
+struct CrossThreadReductionReplacer : public ir::IRMutator<> {
   void operator()(ir::Expr* expr) { Visit(expr); }
 
  private:
@@ -91,6 +99,23 @@ struct CrossThreadReductionReplacer : public ir::IRMutator<Expr*> {
 
   void Visit(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
 
+  void Visit(const ir::_LoweredFunc_* expr, ir::Expr* op) override {
+    ir::IRMutator<>::Visit(expr, op);
+    if (std::find_if(op->as_lowered_func()->temp_bufs.begin(),
+                     op->as_lowered_func()->temp_bufs.end(),
+                     [&](const ir::Buffer& buf) -> bool {
+                       for (auto& tmp_buf : shm_buffer_) {
+                         if (buf->name == tmp_buf->name) return true;
+                       }
+                       return false;
+                     }) == op->as_lowered_func()->temp_bufs.end())
+      op->as_lowered_func()->temp_bufs.insert(
+          op->as_lowered_func()->temp_bufs.end(),
+          shm_buffer_.begin(),
+          shm_buffer_.end());
+    shm_buffer_.clear();
+  }
+
   void Visit(const ir::ScheduleBlockRealize* expr, ir::Expr* op) override {
     if (!CanReplace(expr)) {
       VLOG(6) << "Can't replace cross thread reduction: " << *op;
@@ -113,17 +138,24 @@ struct CrossThreadReductionReplacer : public ir::IRMutator<Expr*> {
       original_update_stmt = original_update_body;
     }
 
-#define REPLACE_TO_EXTERNAL_CALL(Op)                                   \
-  if (original_update_stmt.As<ir::Store>()->value.As<Op>()) {          \
-    auto* node = original_update_stmt.As<ir::Store>()->value.As<Op>(); \
-    CHECK(node);                                                       \
-    auto& operand = node->b();                                         \
-    std::string reduce_func_name =                                     \
-        hlir::pe::CrossThreadReduceExternalFuncName(                   \
-            original_update_stmt.As<ir::Store>()->value,               \
-            operand.As<ir::Load>()->tensor);                           \
-    original_update_stmt.As<ir::Store>()->value =                      \
-        lang::CallExtern(reduce_func_name, {node->b()});               \
+#define REPLACE_TO_EXTERNAL_CALL(Op)                                     \
+  if (original_update_stmt.As<ir::Store>()->value.As<Op>()) {            \
+    auto* node = original_update_stmt.As<ir::Store>()->value.As<Op>();   \
+    CHECK(node);                                                         \
+    auto& operand = node->b();                                           \
+    std::string reduce_func_name =                                       \
+        hlir::pe::CrossThreadReduceExternalFuncName(                     \
+            original_update_stmt.As<ir::Store>()->value,                 \
+            operand.As<ir::Load>()->tensor);                             \
+    auto tmp_dtype = operand.As<ir::Load>()->tensor.as_tensor()->type(); \
+    auto tmp_buffer = ir::_Buffer_::Make(                                \
+        "shm32_" + hlir::pe::Type2StrForReduce(tmp_dtype) + "_reduce",   \
+        {ir::Expr(32)});                                                 \
+    tmp_buffer->dtype = tmp_dtype;                                       \
+    tmp_buffer->memory_type = ir::MemoryType::GPUShared;                 \
+    shm_buffer_.insert(tmp_buffer);                                      \
+    original_update_stmt.As<ir::Store>()->value =                        \
+        lang::CallExtern(reduce_func_name, {node->b(), tmp_buffer});     \
   }
 
     REPLACE_TO_EXTERNAL_CALL(ir::Add)
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(A[i0_0, i1])
+          B[i0_0] = cinn_block_reduce_sum_fp32_internal_shm(A[i0_0, i1], _Buffer_<cinn_buffer_t*: 32>(shm32__fp32_reduce))
         }
       }
     }
diff --git a/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh b/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh
@@ -572,13 +572,57 @@ EXPAND_REDUCE_FP16_MACRO(CINN_BLOCK_REDUCE_INTERNAL_MACRO)
 #undef CINN_BLOCK_REDUCE_INTERNAL_IMPL
 #undef CINN_BLOCK_REDUCE_INTERNAL_MACRO
 
+#define CINN_BLOCK_REDUCE_INTERNAL_SHM_IMPL(TYPE, value, init_value, cinn_warp_shuffle_internal) \
+  int warp_id = threadIdx.x / 32;                                                            \
+  TYPE tmp_val = cinn_warp_shuffle_internal(value);                                          \
+  if (blockDim.x <= 32) {                                                                    \
+    return tmp_val;                                                                          \
+  }                                                                                          \
+  if (warp_id == 0) {                                                                        \
+    shm[threadIdx.x] = init_value;                                                           \
+  }                                                                                          \
+  __syncthreads();                                                                           \
+  if ((threadIdx.x & 31) == 0) {                                                             \
+    shm[warp_id] = tmp_val;                                                                  \
+  }                                                                                          \
+  __syncthreads();                                                                           \
+  if (warp_id == 0) {                                                                        \
+    tmp_val = shm[threadIdx.x];                                                              \
+    shm[threadIdx.x] = cinn_warp_shuffle_internal(tmp_val);                                  \
+  }                                                                                          \
+  __syncthreads();                                                                           \
+  return shm[0];
+
+#define CINN_BLOCK_REDUCE_INTERNAL_SHM_MACRO(REDUCE_TYPE, INITIAL_VALUE, DTYPE)                                                        \
+  __device__ inline DTYPE cinn_block_reduce_##REDUCE_TYPE##_internal_shm(const DTYPE value, DTYPE* shm) {                              \
+    CINN_BLOCK_REDUCE_INTERNAL_SHM_IMPL(DTYPE, value, (DTYPE)(INITIAL_VALUE), cinn_warp_shuffle_##REDUCE_TYPE##_internal);             \
+  }
+
+EXPAND_REDUCE_INT32_MARCO(CINN_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+EXPAND_REDUCE_INT64_MARCO(CINN_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+EXPAND_REDUCE_FP32_MACRO(CINN_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+EXPAND_REDUCE_FP64_MACRO(CINN_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+EXPAND_REDUCE_BOOL_MACRO(CINN_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+
+#ifdef CINN_CUDA_BF16
+EXPAND_REDUCE_BF16_MACRO(CINN_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+#endif
+
+#ifdef CINN_CUDA_FP16
+EXPAND_REDUCE_FP16_MACRO(CINN_BLOCK_REDUCE_INTERNAL_SHM_MACRO)
+#endif
+
+#undef CINN_BLOCK_REDUCE_INTERNAL_SHM_IMPL
+#undef CINN_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) { \
+    __shared__ DTYPE shm[32];                                                                         \
     DTYPE tmp_val = (DTYPE)(INITIAL_VALUE);                                                           \
     for (int i = threadIdx.x; i < extend; i += blockDim.x) {                                          \
       tmp_val = cinn_##REDUCE_TYPE(tmp_val, buf[offset + i]);                                         \
     }                                                                                                 \
-    return cinn_block_reduce_##REDUCE_TYPE##_internal(tmp_val);                                       \
+    return cinn_block_reduce_##REDUCE_TYPE##_internal_shm(tmp_val,shm);                                       \
   }
 
 EXPAND_REDUCE_INT32_MARCO(CINN_BLOCK_REDUCE_IMPL)
diff --git a/paddle/cinn/runtime/cuda/cuda_intrinsics_reduce.cc b/paddle/cinn/runtime/cuda/cuda_intrinsics_reduce.cc
@@ -110,6 +110,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_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_BLOCK_REDUCE_FUNC_IMPL(REDUCE_TYPE, DTYPE)                   \
   REGISTER_FACKED_EXTERN_FUNC_HELPER(cinn_block_reduce_##REDUCE_TYPE, target) \
       .SetRetType<DTYPE>()                                                    \

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(A[i0_0, i1])`
	`74`	`+ B[i0_0] = cinn_block_reduce_sum_fp32_internal_shm(A[i0_0, i1], _Buffer_<cinn_buffer_t*: 32>(shm32__fp32_reduce))`
`75`	`75`	`}`
`76`	`76`	`}`
`77`	`77`	`}`