apache / apache/tvm

[Bug] ThreadStorageSync Pass must be put after MergeSharedMemory Pass

Open
#17,439 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 19h
Merged PRs (30d)
111

Description

In our current lowering pipeline, ThreadSync is placed before the MergeSharedMemoryAllocations Pass, which may lead to unknown behaviors because MergeSharedMemoryAllocations will modify the buffer access region.

https://github.com/apache/tvm/blob/main/src/driver/driver_api.cc#L585-L613

```cpp
bool detect_global_barrier =
pass_ctx->GetConfig("tir.detect_global_barrier", Bool(false)).value();
if (detect_global_barrier) {
mixed_pass_list.push_back(tir::transform::ThreadSync("global"));
}

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::ThreadSync("warp"));
mixed_pass_list.push_back(tir::transform::InferFragment());
mixed_pass_list.push_back(tir::transform::LowerThreadAllreduce());

bool use_async_copy = pass_ctx->GetConfig("tir.use_async_copy", Bool(false)).value();

if (use_async_copy) {
mixed_pass_list.push_back(tir::transform::InjectPTXAsyncCopy());
}

bool ptx_ldg32 = pass_ctx->GetConfig("tir.ptx_ldg32", Bool(false)).value();
if (ptx_ldg32) {
mixed_pass_list.push_back(tir::transform::InjectPTXLDG32());
}

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());
```

Given a simple matmul schedule pipeline:

```python
Store A_shared
Store B_shared

tvm_storage_sync

Load A_shared
Load B_shared

Store C_shared
tvm_storage_sync
Load C_shared
```

The ThreadSync pass will not inject a `tvm_storage_sync` before Store C_shared, that's make sense because C_shared is a non-interfering memory with A_shared and B_shared.

However, when we merge shared memory, C_shared will reuse the memory space with `A_shared` and `B_shared`.

```python
Store A_shared
Store B_shared

tvm_storage_sync

Load A_shared
Load B_shared

Store C_shared(reuse memory space with A_shared and B_shared)
tvm_storage_sync
Load C_shared(reuse memory space with A_shared and B_shared)
```

which is supposed to be a tvm_storage_sync statement before `Store C_shared`, otherwise may lead to a unknown behavior (random and small incorrect produce) because the Store C_shared may change the elements in Load A_shared.

And the solution is quite simple, put the ThreadStorageSync Pass after MergeSharedMemory Pass.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in src/driver/driver_api.cc at the mixed_pass_list around lines 585-613 and trace the ThreadSync and MergeSharedMemoryAllocations entries. Check the matmul schedule described in the issue and confirm that the final pipeline inserts synchronization before reused shared-memory stores without introducing incorrect behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.