[CINN][Backend Pass Update No.12] Update transform_gpu_forloop pass (#70883)

* Update transform_gpu_forloop pass

* Update op_lowering_impl.cc

* Update CudaSyncThreadsDropIfThenElse pass

* Disable EliminateCommonGlobalMemoryRead pass

Author: Albresky

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

@@ -40,6 +40,7 @@
 #include "paddle/cinn/optim/eliminate_common_global_memory_read.h"
 #include "paddle/cinn/optim/schedule_block_dce.h"
 #include "paddle/cinn/optim/transform_gpu_forloop.h"
+#include "paddle/cinn/pass/pass_manager.h"
 #include "paddle/common/ddim.h"
 #include "paddle/common/enforce.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
@@ -393,12 +394,26 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(
           [&](common::NVGPUArch) {
 #ifdef CINN_WITH_CUDA
             // optim::EliminateCommonGlobalMemoryRead(&(func_body));
-            optim::OptimizeExprGPU(&(func_body));
+            ir::stmt::BlockRef func_body_block =
+                ir::ConvertExprBlockToStmtBlock(func_body);
+            LOG(INFO) << "Before OptimizeExprGPU in op_lowering_impl: \n"
+                      << func_body_block;
+            optim::OptimizeExprGPU(func_body_block);
+            LOG(INFO) << "After OptimizeExprGPU in op_lowering_impl: \n"
+                      << func_body_block;
+            func_body = ir::ConvertStmtBlockToExprBlock(func_body_block);
 #endif
           },
           [&](std::variant<common::HygonDCUArchHIP, common::HygonDCUArchSYCL>) {
             // optim::EliminateCommonGlobalMemoryRead(&(func_body));
-            optim::OptimizeExprGPU(&(func_body));
+            ir::stmt::BlockRef func_body_block =
+                ir::ConvertExprBlockToStmtBlock(func_body);
+            LOG(INFO) << "Before OptimizeExprGPU in op_lowering_impl: \n"
+                      << func_body_block;
+            optim::OptimizeExprGPU(func_body_block);
+            LOG(INFO) << "After OptimizeExprGPU in op_lowering_impl: \n"
+                      << func_body_block;
+            func_body = ir::ConvertStmtBlockToExprBlock(func_body_block);
           });
     }
 

paddle/cinn/optim/optimize.cc

@@ -85,9 +85,17 @@ ir::LoweredFunc Optimize(ir::LoweredFunc fn,
 #ifdef CINN_WITH_CUDA
         ir::SetCudaAxisInfo(copied);
         if (remove_gpu_for_loops) {
-          RemoveGpuForLoops(copied);
+          LOG(INFO) << "Before removing GPU for loops:\n" << copied;
+          FuncPassManager func_pass_manager;
+          func_pass_manager.AddPass(CreateRemoveGpuForLoopsPass());
+          func_pass_manager.Run(copied);
+          LOG(INFO) << "After removing GPU for loops:\n" << copied;
         }
-        CudaSyncThreadsDropIfThenElse(copied);
+        VLOG(10) << "Before Optimize CudaSyncThreadsDropIfThenElse:" << copied;
+        BlockPassManager blk_pass_manager;
+        blk_pass_manager.AddPass(CreateCudaSyncThreadsDropIfThenElsePass());
+        blk_pass_manager.Run(copied->body_block);
+        VLOG(10) << "After Optimize CudaSyncThreadsDropIfThenElse:" << copied;
         FuncPassManager func_pass_manager;
         VLOG(10) << "Before Optimize TransBufferWithDynamicShape:" << copied;
         func_pass_manager.AddPass(CreateTransBufferWithDynamicShapePass());
@@ -99,10 +107,17 @@ ir::LoweredFunc Optimize(ir::LoweredFunc fn,
 #ifdef CINN_WITH_HIP
         ir::SetCudaAxisInfo(copied);
         if (remove_gpu_for_loops) {
-          RemoveGpuForLoops(copied);
+          LOG(INFO) << "Before removing GPU for loops:\n" << copied;
+          FuncPassManager func_pass_manager;
+          func_pass_manager.AddPass(CreateRemoveGpuForLoopsPass());
+          func_pass_manager.Run(copied);
+          LOG(INFO) << "After removing GPU for loops:\n" << copied;
         }
-        CudaSyncThreadsDropIfThenElse(copied);
-    // CudaTransBufferWithDynamicShape(&copied);
+        VLOG(10) << "Before Optimize CudaSyncThreadsDropIfThenElse:" << copied;
+        BlockPassManager blk_pass_manager;
+        blk_pass_manager.AddPass(CreateCudaSyncThreadsDropIfThenElsePass());
+        blk_pass_manager.Run(copied->body_block);
+        VLOG(10) << "After Optimize CudaSyncThreadsDropIfThenElse:" << copied;
 #endif
       },
       [&](common::HygonDCUArchSYCL) { CINN_NOT_IMPLEMENTED },

paddle/cinn/optim/replace_var_with_expr.cc

@@ -118,6 +118,13 @@ struct ReplaceVarWithExprMutator : public ir::IRMutator<>,
         ir::IRMutator<>::Visit(&var->upper_bound, &var->upper_bound);
       }
     }
+
+    std::vector<Expr> iter_values = stmt->iter_values();
+    for (ir::Expr& iter_value : iter_values) {
+      ir::IRMutator<>::Visit(&iter_value, &iter_value);
+    }
+    stmt->set_iter_values(iter_values);
+
     std::vector<Expr> new_read_buffers = stmt->read_buffers();
     for (Expr& read_buffer : new_read_buffers) {
       ir::IRMutator<>::Visit(&read_buffer, &read_buffer);