PaddlePaddle
diff --git a/‎paddle/cinn/hlir/pe/reduction.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions b/‎paddle/cinn/hlir/pe/reduction.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/pe/reduction.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/cinn/hlir/pe/reduction.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/cinn/ir/lowered_func.h
Copy file name to clipboard
+26Lines changed: 26 additions & 0 deletions b/‎paddle/cinn/ir/lowered_func.h
Copy file name to clipboard
+26Lines changed: 26 additions & 0 deletions
diff --git a/‎paddle/cinn/optim/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/cinn/optim/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/cinn/optim/optimize.cc
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/cinn/optim/optimize.cc
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
@@ -1205,6 +1205,31 @@ std::string DiscreteReduceExternalFuncName(const ir::Expr& op,
   return "";
 }
 
+std::string GridReduceExternalFuncName(const ir::Expr& op,
+                                       const cinn::common::Type type) {
+  if (op.As<ir::Add>()) {
+    if (type.is_bool()) {
+      return "cinn_grid_reduce_any";
+    }
+    return "cinn_grid_reduce_sum" + Type2StrForReduce(type);
+  } else if (op.As<ir::Mul>()) {
+    if (type.is_bool()) {
+      return "cinn_grid_reduce_all";
+    }
+    return "cinn_grid_reduce_prod" + Type2StrForReduce(type);
+  } else if (op.As<ir::Max>()) {
+    return "cinn_grid_reduce_max" + Type2StrForReduce(type);
+  } else if (op.As<ir::Min>()) {
+    return "cinn_grid_reduce_min" + Type2StrForReduce(type);
+  } else if (op.As<ir::And>()) {
+    return "cinn_grid_reduce_all";
+  } else if (op.As<ir::Or>()) {
+    return "cinn_grid_reduce_any";
+  }
+  PADDLE_THROW(::common::errors::InvalidArgument(
+      "No matching grid reduce template for op: %s, type: %s", op, type));
+}
+
 }  // namespace pe
 }  // namespace hlir
 }  // namespace cinn
@@ -474,6 +474,9 @@ std::string CrossThreadReduceExternalFuncName(const ir::Expr& op,
 std::string DiscreteReduceExternalFuncName(const ir::Expr& op,
                                            const ir::Expr& tensor);
 
+std::string GridReduceExternalFuncName(const ir::Expr& op,
+                                       const cinn::common::Type type);
+
 std::string Type2StrForReduce(cinn::common::Type type);
 }  // namespace pe
 }  // namespace hlir
 
@@ -113,6 +113,28 @@ struct CudaAxisInfo {
 
 std::ostream& operator<<(std::ostream& os, const CudaAxisInfo& x);
 
+/**
+ * A struct representing a temporary global buffer (allocated on the heap) that
+ * is used as staging space during kernel execution.
+ */
+struct TempSpaceInfo {
+  TempSpaceInfo() = default;
+  TempSpaceInfo(const Expr& size, int arg_idx, bool need_zero_init = false)
+      : size_(size), arg_idx_(arg_idx), need_zero_init_(need_zero_init) {}
+
+  Expr size() const { return size_; }
+  int arg_idx() const { return arg_idx_; }
+  bool need_zero_init() const { return need_zero_init_; }
+
+ private:
+  // size of the space in bytes
+  Expr size_;
+  // index in the function's argument list
+  int arg_idx_;
+  // whether this space need to be zero-initialized
+  bool need_zero_init_;
+};
+
 /**
  * Definition of a lowered function. Note that, it should be functional.
  *
@@ -131,6 +153,10 @@ struct _LoweredFunc_ : ExprNode<_LoweredFunc_> {
   //! function's argument list, but will be used in the body.
   std::vector<Buffer> temp_bufs;
 
+  //! Temporary global buffers. These buffers will appear in the function's
+  //! argument list.
+  std::vector<TempSpaceInfo> temp_spaces;
+
   //! Body of this function.
   Expr body;
 
 
@@ -24,6 +24,7 @@ gather_srcs(
   cast_bool_to_int8.cc
   var_mod_simplify.cc
   remove_schedule_block.cc
+  replace_cross_block_reduction.cc
   replace_cross_thread_reduction.cc
   replace_mod_to_max.cc
   resize_buffer.cc
 
@@ -31,6 +31,7 @@
 #include "paddle/cinn/optim/rearrange_load_instruction.h"
 #include "paddle/cinn/optim/remove_schedule_block.h"
 #include "paddle/cinn/optim/replace_const_param_to_integer.h"
+#include "paddle/cinn/optim/replace_cross_block_reduction.h"
 #include "paddle/cinn/optim/replace_cross_thread_reduction.h"
 #include "paddle/cinn/optim/trans_buffer_with_dynamic_shape.h"
 #include "paddle/cinn/optim/transform_gpu_forloop.h"
@@ -59,6 +60,9 @@ Expr Optimize(Expr e,
   // Simplify already contains CastSimplify
   Simplify(&copied);
   ReplaceCrossThreadReduction(&copied);
+  VLOG(4) << "After Optimize ReplaceCrossThreadReduction:" << copied;
+  ReplaceCrossBlockReduction(&copied);
+  VLOG(4) << "After Optimize ReplaceCrossBlockReduction:" << copied;
   UnrollLoop(&copied);
   VLOG(4) << "After Optimize UnrollLoop:" << copied;