Fix

Author: co63oc

paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc

@@ -56,7 +56,6 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.h"
 #include "paddle/fluid/pir/transforms/build_cinn_pass.h"
 #include "paddle/fluid/pir/transforms/general/dead_code_elimination_pass.h"
-#include "paddle/fluid/pir/transforms/general/identity_op_clean_pass.h"
 #include "paddle/fluid/pir/transforms/gpu/fused_gemm_epilogue_pass.h"
 
 COMMON_DECLARE_bool(cinn_specify_input_dynamic_dim);
@@ -95,19 +94,6 @@ bool HasDynamicShape(const pir::Program& program) {
 }
 }  // namespace
 
-void ApplyIdentityOpCleanPass(
-    ::pir::Program* program,
-    const std::function<std::shared_ptr<::pir::PassManager>()>&
-        CreatePassManager) {
-  // NOTE(gongshaotian):Before Paddle 3.0, useless full op and scale op are
-  // inserted at the end of the Program when export models using Paddle. This
-  // Pass is designed to address the above-mentioned issues encountered when
-  // open CINN execution in some models that cannot be reexported.
-  std::shared_ptr<pir::PassManager> pass_manager = CreatePassManager();
-  pass_manager->AddPass(pir::CreateIdentityOpCleanPass());
-  pass_manager->Run(program);
-}
-
 void ApplyShapeOptimizationPass(
     ::pir::Program* program,
     const std::function<std::shared_ptr<::pir::PassManager>()>&
@@ -278,7 +264,6 @@ void ApplyCinnPass(::pir::Program* program,
       .file_name("original_programs.py")
       .dump_symbolic_shape(FLAGS_logging_pir_py_code_dump_symbolic_dims)
       .SaveIfFlagEnabled();
-  ApplyIdentityOpCleanPass(program, CreatePassManager);
   ApplyShapeOptimizationPass(program, CreatePassManager);
   ApplyPdToCinnPass(program, CreatePassManager);
   ApplyCinnPreprocessPass(program, CreatePassManager);

paddle/cinn/hlir/framework/pir/trivial_op_impl.cc

@@ -36,7 +36,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"
 
-PD_DECLARE_bool(group_schedule_tiling_first);
+PD_DECLARE_bool(cinn_enable_grid_reduce);
 
 namespace cinn {
 namespace hlir {
@@ -808,8 +808,10 @@ std::shared_ptr<FusionGroupInfo> GetFusionGroupInfo(
                    });
   }
 
-  group_info->can_apply_grid_reduce =
-      GetCanApplyGridReduce(op_compute_bodies, group_info->reduce_axis);
+  if (FLAGS_cinn_enable_grid_reduce) {
+    group_info->can_apply_grid_reduce =
+        GetCanApplyGridReduce(op_compute_bodies, group_info->reduce_axis);
+  }
 
   VLOG(4) << group_info->DebugPrint();
   return group_info;

paddle/cinn/hlir/framework/pir/trivial_op_util.cc

@@ -471,10 +471,9 @@ ExprTransformer UnsqueezeForTransformer(
       schedule_block->body =
           WrapForTransformer(to_append_var)(schedule_block->body);
     } else {
-      PADDLE_THROW(
+      PADDLE_THROW(::common::errors::PreconditionNotMet(
           "UnsqueezeForTransformer: only support insert after a (For / "
-          "ScheduleBlockRealizer): %s",
-          followed_expr);
+          "ScheduleBlockRealizer)"));
     }
     VLOG(6) << "UnsqueezeForTransformer: After changed: " << copied_e;
     return copied_e;
@@ -547,19 +546,24 @@ ExprTransformer RemoveVarInScheduleBlockRealize(const ir::Var& target_vars,
     VLOG(4) << "Start RemoveVarInScheduleBlockRealize(" << target_vars << ", "
             << replaced_expr << ")";
     VLOG(4) << "      Input is " << e;
-    PADDLE_ENFORCE(e.As<ir::ScheduleBlockRealize>() != nullptr,
-                   "RemoveVarInScheduleBlockRealize: input expr is not a "
-                   "ScheduleBlockRealize.");
+    PADDLE_ENFORCE_NE(
+        e.As<ir::ScheduleBlockRealize>(),
+        nullptr,
+        ::common::errors::InvalidArgument(
+            "RemoveVarInScheduleBlockRealize: input expr is not a "
+            "ScheduleBlockRealize."));
     auto copied_ir = ir::ir_utils::IRCopy(e);
     auto schedule_block_iter_vars =
         copied_ir.As<ir::ScheduleBlockRealize>()->iter_values;
     auto block_bound_vars = copied_ir.As<ir::ScheduleBlockRealize>()
                                 ->schedule_block.As<ir::ScheduleBlock>()
                                 ->iter_vars;
     for (const auto& i_var : schedule_block_iter_vars) {
-      PADDLE_ENFORCE(
+      PADDLE_ENFORCE_EQ(
           i_var.is_var(),
-          "RemoveVarInScheduleBlockRealize: axes.bind rhs is is not a Var.");
+          true,
+          ::common::errors::InvalidArgument("RemoveVarInScheduleBlockRealize: "
+                                            "axes.bind rhs is is not a Var."));
     }
     // find replace idx
     int target_idx = -1;
@@ -686,10 +690,11 @@ ExprTransformer RemoveOneTransformer(int one) {
       VLOG(4) << "RemoveOneTransformer: father block is root realize";
       ir::Expr shedule_block =
           target_block.As<ir::ScheduleBlockRealize>()->schedule_block;
-      PADDLE_ENFORCE_EQ(shedule_block.As<ir::ScheduleBlock>()->body,
-                        target_for,
-                        ::common::errors::PreconditionNotMet(
-                            "Root realize body should be equal to target for"));
+      PADDLE_ENFORCE_EQ(
+          shedule_block.As<ir::ScheduleBlock>()->body,
+          target_for,
+          ::common::errors::InvalidArgument(
+              "Root realize body should be equal to target for."));
       const auto for_body = target_for.As<ir::For>()->body;
       const auto for_body_stmts = for_body.As<ir::Block>()->stmts;
       if (for_body_stmts.size() == 1 &&
@@ -747,12 +752,17 @@ ExprTransformer RemoveOnesTransformer(const std::vector<int32_t>& ones) {
 ExprTransformer TransposeForsTransformer(const std::vector<int32_t>& perm) {
   const auto& f = [=](const ir::Expr& root) -> ir::Expr {
     const auto& iters = GetNonReduceLoopVars(root);
-    PADDLE_ENFORCE_EQ(iters.size(),
-                      perm.size(),
-                      "Transposed iters size and perm size should be equal.");
+    PADDLE_ENFORCE_EQ(
+        iters.size(),
+        perm.size(),
+        ::common::errors::InvalidArgument(
+            "Transposed iters size and perm size should be equal."));
     for (size_t i = 0; i < perm.size(); ++i) {
       if (iters[i]->is_reduce_axis) {
-        PADDLE_ENFORCE_EQ(i, perm[i], "Can only transpose non reduce iters.");
+        PADDLE_ENFORCE_EQ(i,
+                          perm[i],
+                          ::common::errors::InvalidArgument(
+                              "Can only transpose non reduce iters."));
       }
     }
     const auto transposed_iters = cinn::fusion::TransposeVector(iters, perm);
@@ -773,7 +783,7 @@ ExprTransformer InsertForsTransformer(const std::vector<int32_t>& axis,
         axis.size(),
         vars.size(),
         ::common::errors::InvalidArgument(
-            "The number of axis to insert and vars should be equal"));
+            "The number of axis to insert and vars should be equal."));
     const size_t reduce_size =
         std::count_if(iters.begin(), iters.end(), [](const ir::Var& v) {
           return v->is_reduce_axis;
@@ -782,7 +792,7 @@ ExprTransformer InsertForsTransformer(const std::vector<int32_t>& axis,
       PADDLE_ENFORCE_LE(axis[i],
                         iters.size() - reduce_size,
                         ::common::errors::OutOfRange(
-                            "Insert axis should not be behind reduce axis"));
+                            "Insert axis should not be behind reduce axis."));
       iters.insert(iters.begin() + axis[i], vars[i]);
     }
     const auto non_reduce_iters =
@@ -837,7 +847,7 @@ int InplaceMutateSingleExpr(ir::Expr* root,
   PADDLE_ENFORCE_EQ(
       source.size(),
       1,
-      ::common::errors::InvalidArgument("Only one expr should be found"));
+      ::common::errors::InvalidArgument("Only one expr should be found."));
   const auto& target = transformer(source[0]);
   ComposeUtils::MappingTargetExprToDestExprMutator(source[0], target)(root);
   return 1;  // operation number.
@@ -880,7 +890,10 @@ void CheckFusionInputValid(const std::vector<ir::Expr>& op_compute_bodies,
   VLOG(4) << "      op_patterns.size() = " << op_compute_bodies.size();
   VLOG(4) << "op_compute_bodies.size() = " << op_patterns.size();
   PADDLE_ENFORCE_EQ(
-      op_patterns.size(), op_compute_bodies.size(), "ops and  size not equal");
+      op_patterns.size(),
+      op_compute_bodies.size(),
+      ::common::errors::InvalidArgument(
+          "The number of op_compute_bodies and op_patterns should be equal."));
 }
 
 std::vector<ir::Var> AppendBound(const std::vector<ir::Var> vars,

paddle/cinn/ir/group_schedule/config/group_tile_util.cc

@@ -34,7 +34,6 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,
   // A tensor is downstream of reduce either if it is produced by a reduce, or
   // if it has data dependency on another tensor that is downstream of reduce.
   std::unordered_set<std::string> reduce_downstream_tensor_names;
-  int reduce_count = 0;
 
   const auto IsReduceDownstream = [&](const ir::Expr& expr_block) {
     for (auto& expr_load : ChildTensorLoads(expr_block)) {
@@ -90,9 +89,6 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,
     bool is_reduce_downstream = IsReduceDownstream(expr_block);
     bool output_has_reduce_axis = CheckOutputHasReduceAxis(body, expr_block);
 
-    if (is_reduce) {
-      ++reduce_count;
-    }
     if (is_reduce_downstream || is_reduce) {
       AddReduceDownstream(expr_block);
     }
@@ -105,8 +101,7 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,
       return false;
     }
   }
-
-  return reduce_count == 1;
+  return true;
 }
 
 }  // namespace ir

paddle/cinn/ir/group_schedule/config/group_tile_util.h

@@ -19,8 +19,8 @@ namespace cinn {
 namespace ir {
 
 // Check whether we can apply grid reduce in this group.
-// We can apply grid reduce if there is exactly one reduce, and whose result is
-// not broadcasted before output.
+// We can apply grid reduce if there is no reduce-then-broadcast dependency
+// in this group.
 bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,
                            const std::vector<int64_t>& reduce_axis);
 

paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.cc

@@ -79,6 +79,7 @@ class TileFirstGeneralTactic final : public ScheduleTactic {
   std::vector<int32_t> vec_flatten_axis_;
   std::vector<int32_t> vec_reduce_axis_;
   std::unordered_map<std::string, std::string> map_rf_block_;
+  std::unordered_map<std::string, std::string> map_global_rf_block_;
 };
 
 void TileFirstGeneralTactic::Init(ScheduleContext* context) {
@@ -381,19 +382,46 @@ void TileFirstGeneralTactic::SplitSptialInner(ir::IRSchedule* sch,
 
 void TileFirstGeneralTactic::SplitReduceInner(ir::IRSchedule* sch,
                                               const std::string& block_id) {
+  const int64_t rd_block = context_->config.tile_config.grid_reduce_num;
+  const int64_t rd_thread = 16;
+  const int cur_reduce_axis = 2;
+
+  // [ R ] => [ rd_block*rd_thread, rd_inner ]
   auto loops = sch->GetLoops(block_id);
-  // [S(-1), S(32), R] => [S(-1), S(32), R(16), R(-1)]
-  sch->Split(loops[2], std::vector<int>{16, -1});
+  sch->Split(loops[cur_reduce_axis],
+             std::vector<int>{-1, rd_block * rd_thread});
+  loops = sch->GetLoops(block_id);
+  sch->Reorder({loops[cur_reduce_axis + 1], loops[cur_reduce_axis]});
 
   loops = sch->GetLoops(block_id);
   if (IsReductionSBlock(sch->GetBlock(block_id)) &&
       ir::GetLoopExtent(loops[2]) != 1) {
     ir::Expr rf_tensor =
-        sch->FactorizeReduction(loops[2],
-                                0,
+        sch->FactorizeReduction(loops[cur_reduce_axis],
+                                /* rf_axis = */ 0,
                                 /* with_write_back_block_init = */ false);
     map_rf_block_[block_id] = rf_tensor.as_tensor_ref()->name;
   }
+
+  // [ rd_block*rd_thread ] => [ rd_block, rd_thread ]
+  if (rd_block > 1) {
+    loops = sch->GetLoops(block_id);
+    sch->Split(loops[cur_reduce_axis], {rd_block, rd_thread});
+
+    if (IsReductionSBlock(sch->GetBlock(block_id))) {
+      loops = sch->GetLoops(map_rf_block_[block_id]);
+      sch->Split(loops[cur_reduce_axis], {rd_block, rd_thread});
+
+      loops = sch->GetLoops(block_id);
+      ir::Expr rf_tensor =
+          sch->FactorizeReduction(loops[cur_reduce_axis],
+                                  /* rf_axis = */ 0,
+                                  /* with_write_back_block_init = */ false);
+      std::string tensor_name = rf_tensor.as_tensor_ref()->name;
+      map_global_rf_block_[block_id] = tensor_name;
+      rf_tensor.as_tensor_ref()->WithBuffer("global", "_" + tensor_name);
+    }
+  }
 }
 
 void TileFirstGeneralTactic::VariableTypeAssignment(
@@ -435,6 +463,12 @@ void TileFirstGeneralTactic::SetDiscreteReduceType(
                      ->schedule_block.As<ir::ScheduleBlock>();
     block->reduce_method = cinn::ir::DiscreteReduceMethod();
   }
+  if (map_global_rf_block_.count(block_id) > 0) {
+    auto block = sch->GetBlock(map_global_rf_block_[block_id])
+                     .As<ir::ScheduleBlockRealize>()
+                     ->schedule_block.As<ir::ScheduleBlock>();
+    block->reduce_method = cinn::ir::DiscreteReduceMethod();
+  }
 }
 
 void TileFirstGeneralTactic::BindCudaInfo(ir::IRSchedule* sch,
@@ -446,14 +480,24 @@ void TileFirstGeneralTactic::BindCudaInfo(ir::IRSchedule* sch,
   const auto DoBind = [&](const std::vector<ir::Expr>& loops) {
     sch->Bind(loops[0], "blockIdx.x");
     sch->Bind(loops[1], "threadIdx.x");
-    sch->Bind(loops[2], "threadIdx.y");
+    if (context_->config.tile_config.grid_reduce_num > 1) {
+      sch->Bind(loops[2], "blockIdx.y");
+      if (loops.size() > 3) {
+        sch->Bind(loops[3], "threadIdx.y");
+      }
+    } else {
+      sch->Bind(loops[2], "threadIdx.y");
+    }
   };
 
   DoBind(sch->GetLoops(block_id));
 
   if (map_rf_block_.count(block_id) > 0) {
     DoBind(sch->GetLoops(map_rf_block_[block_id]));
   }
+  if (map_global_rf_block_.count(block_id) > 0) {
+    DoBind(sch->GetLoops(map_global_rf_block_[block_id]));
+  }
 }
 
 std::unique_ptr<ScheduleTactic> CreateTileFirstGeneralTactic() {

paddle/cinn/operator_fusion/pir_graph_analyzing/fusion_iters.h

@@ -36,8 +36,11 @@ struct FusionItersManager {
                      ShardableAxesInfoManager* axes_info)
       : shape_analysis_(shape_analysis), axes_info_(axes_info) {
     PADDLE_ENFORCE_NOT_NULL(shape_analysis,
-                            "shape_analysis should not be nullptr.");
-    PADDLE_ENFORCE_NOT_NULL(axes_info, "axes_info should not be nullptr.");
+                            ::common::errors::InvalidArgument(
+                                "shape_analysis should not be nullptr."));
+    PADDLE_ENFORCE_NOT_NULL(
+        axes_info,
+        ::common::errors::InvalidArgument("axes_info should not be nullptr."));
   }
   FusionItersSignature GetItersSignature(pir::Operation* op);
 

paddle/cinn/operator_fusion/policy/iters_fusion_policy.h

@@ -25,7 +25,8 @@ struct ItersFusionPolicy final : public PolicyBase {
   ItersFusionPolicy(std::shared_ptr<FusionItersManager> iters_manager)
       : iters_manager_(iters_manager) {
     PADDLE_ENFORCE_NOT_NULL(iters_manager,
-                            "iters_manager should not be nullptr.");
+                            ::common::errors::InvalidArgument(
+                                "iters_manager should not be nullptr."));
   }
   static constexpr PolicyKind Kind = PolicyKind::ItersFusion;
   std::string Name() { return "ItersFusionPolicy"; }