[PIR+CINN]Replace FlattenLoops with Fuse (PaddlePaddle#58910)

SecretXV · Nov 28, 2023 · 0442f1a · 0442f1a
1 parent 721e0ac
commit 0442f1a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -578,11 +578,7 @@ ir::Expr OpLowererImpl::DoGroupSchedule(
       // if node is horizontal with reduce or node is reduce, loop assign
       // master.
       auto loops = ir_sch.GetLoops(op_out_name);
-      if (op_kind == framework::kElementWise) {
-        ir_sch.FlattenLoops(loops, true);
-      } else if (op_kind != framework::kReduction) {
-        ir_sch.FlattenLoops(loops, false);
-      }
+      ir_sch.Fuse(loops);
 
       if (master && op_kind != framework::kReduction) {
         auto master_loops =

diff --git a/paddle/cinn/optim/transform_gpu_forloop.cc b/paddle/cinn/optim/transform_gpu_forloop.cc
@@ -396,7 +396,7 @@ class ReplaceLoopVarToGpu : public ir::IRMutator<> {
     auto bind_info = for_ir->bind_info();
 
     std::string var_name = "";
-    if (bind_info.offset == 0)
+    if (bind_info.offset <= 0)
       var_name = "x";
     else if (bind_info.offset == 1)
       var_name = "y";