Reorder MergeSharedMemoryAllocations in GPU codegen

apache · Mar 19, 2024 · e77ef6e · e77ef6e
1 parent 3684fd3
commit e77ef6e
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc
@@ -590,6 +590,7 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target)
 
   mixed_pass_list.push_back(tir::transform::ThreadSync("shared"));
   mixed_pass_list.push_back(tir::transform::ThreadSync("shared.dyn"));
+  mixed_pass_list.push_back(tir::transform::MergeSharedMemoryAllocations());
   mixed_pass_list.push_back(tir::transform::ThreadSync("warp"));
   mixed_pass_list.push_back(tir::transform::InferFragment());
   mixed_pass_list.push_back(tir::transform::LowerThreadAllreduce());
@@ -607,9 +608,6 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target)
 
   mixed_pass_list.push_back(tir::transform::AnnotateDeviceRegions());
   mixed_pass_list.push_back(tir::transform::SplitHostDevice());
-  // MergeSharedMemoryAllocations must be applied after SplitHostDevice
-  // because the merged allocation site is at the beginning of each device function
-  mixed_pass_list.push_back(tir::transform::MergeSharedMemoryAllocations());
 
   bool unpacked_api = mixed_mod->GetAttr<relay::Executor>(tvm::attr::kExecutor)
                           .value_or(relay::Executor::Create("graph", {}))