[Hotfix] Revert driver API pass ordering that breaks MLC, mark failin…

…g test (#16770)

* Revert changes that cause failures in MLC, mark and skip the failing tests



* Restore changes unrelated to driver API reordering

src/driver/driver_api.cc

@@ -590,7 +590,6 @@ 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());
@@ -608,6 +607,9 @@ 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", {}))

tests/python/tir-transform/test_tir_transform_inject_ptx_async_copy.py

@@ -482,6 +482,12 @@ def simple_compute(
     assert generated_code == expected_cuda_script
 
 
+@pytest.mark.skip(
+    reason="This test fails due to an ordering issue with MergeSharedMemoryAllocations "
+    "in device_driver_api.cc. However, fixing this causes failures in MLC. "
+    "This bug should be addressed. See discussion in https://github.com/apache/tvm/pull/16769 "
+    "and https://github.com/apache/tvm/pull/16569#issuecomment-1992720448"
+)
 @tvm.testing.requires_cuda
 def test_vectorize_cp_async_in_if_then_else(postproc_if_missing_async_support):
     @T.prim_func