diff --git a/paddle/cinn/optim/optimize.cc b/paddle/cinn/optim/optimize.cc
index b93d06678c9068..74c3fdb2928298 100644
--- a/paddle/cinn/optim/optimize.cc
+++ b/paddle/cinn/optim/optimize.cc
@@ -122,9 +122,15 @@ ir::LoweredFunc Optimize(ir::LoweredFunc fn,
         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());
+        func_pass_manager.Run(copied);
+        VLOG(10) << "After Optimize TransBufferWithDynamicShape:" << copied;
 #endif
       },
-      [](auto) {});
+      [&](std::variant<common::UnknownArch, common::X86Arch, common::ARMArch>) {
+      });
 
   SimplifyUnitBlock(&copied->body);
   VLOG(4) << "After SimplifyUnitBlock:" << copied;
@@ -157,7 +163,14 @@ ir::LoweredFunc Optimize(ir::LoweredFunc fn,
         func_pass_manager.Run(copied);
         VLOG(4) << "After Optimize RearrangeLoadInstruction:" << copied;
       },
-      [](auto) {});
+      [&](std::variant<common::HygonDCUArchHIP, common::HygonDCUArchSYCL>) {
+        FuncPassManager func_pass_manager;
+        func_pass_manager.AddPass(CreateRearrangeLoadInstructionPass());
+        func_pass_manager.Run(copied);
+        VLOG(4) << "After Optimize RearrangeLoadInstruction:" << copied;
+      },
+      [&](std::variant<common::UnknownArch, common::X86Arch, common::ARMArch>) {
+      });
 
   VectorizeForTrans(&copied->body);
   VLOG(10) << "After Optimize vectorize" << copied;