JosephTheOctonaut commented on code in PR #15684:
URL: https://github.com/apache/tvm/pull/15684#discussion_r1318980467
##########
tests/python/unittest/test_tir_ptx_cp_async.py:
##########
@@ -126,32 +116,20 @@ def ptx_cp_async_bulk(A: T.Buffer((32, 128), "float16"),
B: T.Buffer((32, 128),
T.launch_thread(bx, 1)
T.launch_thread(tx, 32)
with T.block():
- # Shared memory targets for cp.async.bulk must be 16 byte aligned
- # Problem: CUDA codegen does not support allocation alignment
- # Workaround: Ensure that `A_shared` occurs before `barrier` in
program order
- # by allocating and initializing `A_shared` before
`barrier`
- # which should result in `A_shared` being 16+ byte aligned
- # given it will be the first shared memory allocation
- # TODO(Straw) Add CUDA codegen support for allocation alignment
- A_shared = T.alloc_buffer([32, 128], "float16", scope="shared",
align=16)
- A_shared[0, 0] = 0
-
- barrier = T.alloc_buffer([1], "uint64", scope="shared")
- barrier[0] = 0
+ A_shared = T.alloc_buffer([32, 128], "float16", scope="shared")
Review Comment:
The barrier array initialization now rounds up the array size so it's always
16-byte aligned
--
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]
