Skip to content

[Hotfix] Revert driver API pass ordering that breaks MLC, mark failing test #16770

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tqchen merged 3 commits into from
Mar 23, 2024
Merged

[Hotfix] Revert driver API pass ordering that breaks MLC, mark failing test #16770

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ea9b181
Revert changes that cause failures in MLC, mark and skip the failing …
slyubomirsky Mar 22, 2024
428544d
Lint
slyubomirsky Mar 22, 2024
de1e9a6
Restore changes unrelated to driver API reordering
slyubomirsky Mar 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/driver/driver_api.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand All @@ -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", {}))
Expand Down
6 changes: 6 additions & 0 deletions tests/python/tir-transform/test_tir_transform_inject_ptx_async_copy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -482,6 +482,12 @@ def simple_compute(
assert generated_code == expected_cuda_script


@pytest.mark.skip(
Copy link
Contributor

@Lunderberg Lunderberg Mar 22, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like the documentation, both to indicate what needs to be resolved in order to re-enable functionality, and where the existing discussions are

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
Expand Down