nox-410 opened a new pull request, #16215:
URL: https://github.com/apache/tvm/pull/16215
The MergeDynamicSharedMemoryAllocations currently will first free and then
allocate for each scope. However, when a buffer is allocated and freed within a
leaf scope, it will run into a free before alloc bug.
This commit solves this bug by delaying the leaf free after the alloc is
done.
A simple test case is also added
```python
import tvm
import tvm.testing
from tvm.script import tir as T
class TestLeafAllocFree(tvm.testing.CompareBeforeAfter):
transform = tvm.tir.transform.MergeDynamicSharedMemoryAllocations()
def before(self):
@T.prim_func
def func():
threadIdx_x = T.launch_thread("threadIdx.x", 128)
A_sh_data = T.allocate([128], "float32", "shared.dyn")
B_sh_data = T.allocate([128], "float32", "shared.dyn")
A_sh = T.decl_buffer([128], "float32", data=A_sh_data,
scope="shared.dyn")
B_sh = T.decl_buffer([128], "float32", data=B_sh_data,
scope="shared.dyn")
B_sh[threadIdx_x] = A_sh[threadIdx_x]
return func
def expected(self):
@T.prim_func
def func():
threadIdx_x = T.launch_thread("threadIdx.x", 128)
buf_dyn_shmem = T.allocate([1024], "uint8", "shared.dyn")
A_sh = T.decl_buffer((128,), data=buf_dyn_shmem,
scope="shared.dyn")
B_sh = T.decl_buffer((128,), data=buf_dyn_shmem,
scope="shared.dyn")
B_sh[threadIdx_x + 128] = A_sh[threadIdx_x]
return func
if __name__ == "__main__":
tvm.testing.main()
```
--
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]
