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

Aurelius84 · web-flow · commit bd2e3d399100 · 2023-11-11T11:30:46.000+08:00
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";
