LeiWang1999 opened a new pull request, #16342:
URL: https://github.com/apache/tvm/pull/16342
#759 proposed a pass `storage_rewrite` and provided a trivial storage reuse
plan based on liveness analysis, just as #9341 mentioned, the solution has some
limitations:
1. storage_rewrite can't handle buffer with different dtypes.
```c++
int8 A_shared[32];
int8 B_shared[32];
int32 C_shared[4]; // will not be reused even we have enough workspace
as different types.
```
2. storage_rewrite can't allocate a buffer in the place of another 2 buffers.
```c++
int8 A_shared[32];
int8 B_shared[32];
int8 C_shared[64];
// will be reused as A_shared[32], B_shared[64], results in 32 half
elements space waste.
```
#8571 and #9341 introduced a pass `MergeDynamicSharedMemoryAllocations`,
which can support efficient memory reuse solely for dynamic shared memory.
However, sometimes we do not want to use dynamic shared memory for codegen, so
this pull request made a simple extend to `MergeDynamicSharedMemoryAllocations`
to support both dynamic and static shared memory optimal reuse.
By default, the static shared memory merge is disabled to maintain
consistency, to enable the static part:
```python
with tvm.transform.PassContext(config={"tir.merge_static_smem": True}):
cuda_mod = tvm.build(sch.mod, target="cuda")
```
Take int8xint8=int32 tensorcore gemm as an example, we have a big tile and
used static shared memory, before the pass:
```cpp
__global__ void __launch_bounds__(128) Fused(int8_t* __restrict__ input0,
int8_t* __restrict__ input1, int* __restrict__ output0) {
int mediate0_shared_warp[128];
__shared__ signed char input0_shared[16384];
__shared__ signed char input1_shared[16384];
signed char input0_shared_warp[64];
signed char input1_shared_warp[64];
signed char input0_shared_warp_1[64];
signed char input1_shared_warp_1[64];
__shared__ int mediate0_shared[6400];
```
it will exceed the maximum available static shared memory, and compilation
will fail. After this pass
```cpp
__global__ void __launch_bounds__(128) Fused(int8_t* __restrict__ input0,
int8_t* __restrict__ input1, int* __restrict__ output0) {
__shared__ uchar buf_shmem[32768];
int mediate0_shared_warp[128];
signed char input0_shared_warp[64];
signed char input1_shared_warp[64];
signed char input0_shared_warp_1[64];
signed char input1_shared_warp_1[64];
```
we can save around 50% shared memory and the code generation perf with
fastdlight can achieve 510+Tflops (without the pass, the best tile is around
420TFlops on A100), this pass will enable us to explore more tile configs under
static shared memory.
Moreover, the pass can optimize the dynamic shared memory plan as well, as
the storage_rewrite pass will merge C_shared to B_shared in this example, which
is not friendly for further memory plan analysis, the flag `merge_static_smem`
will disable the trivial reuse behavior by (don't know if the flag can be
improved):
```python
if (!enable_reuse || is_small_array || !is_flat_memory_space) {
return NewAlloc(op, attach_scope, scope, const_nbits);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
