[Bug] ThreadStorageSync Pass must be put after MergeSharedMemory Pass
- 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