csullivan commented on code in PR #15684:
URL: https://github.com/apache/tvm/pull/15684#discussion_r1318954589
##########
src/target/source/codegen_cuda.cc:
##########
@@ -968,42 +968,61 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op,
std::ostream& os) {
std::string src = this->PrintExpr(op->args[2]);
std::string src_offset = this->PrintExpr(op->args[3]);
std::string size = this->PrintExpr(op->args[4]);
- std::string barrier_ptr = this->PrintExpr(op->args[5]);
- std::string barrier_offset = this->PrintExpr(op->args[6]);
- this->stream << PrintCpAsyncBulkAsm(dst, dst_offset, src, src_offset,
size, barrier_ptr,
- barrier_offset);
+ int barrier_id = Downcast<IntImm>(op->args[5])->value;
+ ICHECK(barrier_id < barrier_count_);
Review Comment:
nit: These ICHECKs could be CHECKs since this is user facing in many cases.
##########
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:
Is this comment resolved? I see that you removed align=16, which I
understand is a requirement of cp.async.build for operands. Are we just getting
this by default due to 64 byte alignment?
--
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]
