You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "adstraw (via GitHub)" <gi...@apache.org> on 2023/09/06 14:10:13 UTC
[GitHub] [tvm] adstraw opened a new pull request, #15684: [Hopper TMA] Add intrinsic to create barriers for synchronization
adstraw opened a new pull request, #15684:
URL: https://github.com/apache/tvm/pull/15684
This PR adds an intrinsic to create barriers that can be used with existing barrier intrinsics for synchronization. The prior method of barrier allocation was to use `alloc_buffer` e.g. as follows `barrier = T.alloc_buffer([1], "uint64", scope="shared")` and then pass the pointer and offset to that barrier allocation for use in the barrier intrinsics. This was a functional interface, but also caused problems with alignment of other non-barrier shared memory allocations. See removed workarounds marked with `TODO` in the tests in this PR. At the expense of the additional `create_barriers` intrinsic we get a simplified interface using barrier ID rather than pointer / offset passed to the barrier intrinsics and more low-level codegen control which is used in this PR to solve the alignment issue.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] adstraw commented on a diff in pull request #15684: [Hopper TMA] Add intrinsic to create barriers for synchronization
Posted by "adstraw (via GitHub)" <gi...@apache.org>.
adstraw commented on code in PR #15684:
URL: https://github.com/apache/tvm/pull/15684#discussion_r1319144691
##########
src/target/source/codegen_cuda.h:
##########
@@ -109,6 +109,13 @@ class CodeGenCUDA final : public CodeGenC {
// Op attribute map
OpAttrMap<bool> op_need_warp_shuffle_ = Op::GetAttrMap<bool>("cuda.need_warp_shuffle");
+ // The name of the barrier array in shared memory
+ const std::string barrier_name_ = "barrier";
+ // The size of the barrier array in shared memory
+ int barrier_count_ = -1;
+ // The alignment of the barrier array in shared memory
+ const int barrier_alignment_bytes_ = 64;
Review Comment:
I chose 64 to match `kAllocAlignment` but the actual alignment requirement for bulk copy is 16. It would work either way. I will use 16 and add a comment noting that it's the minimum.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] JosephTheOctonaut commented on a diff in pull request #15684: [Hopper TMA] Add intrinsic to create barriers for synchronization
Posted by "JosephTheOctonaut (via GitHub)" <gi...@apache.org>.
JosephTheOctonaut commented on code in PR #15684:
URL: https://github.com/apache/tvm/pull/15684#discussion_r1318944190
##########
include/tvm/tir/builtin.h:
##########
@@ -681,43 +680,51 @@ TVM_DLL const Op& ptx_wait_group();
/*!
* \brief tvm intrinsics for ptx async copy barrier using cp.async.mbarrier.arrive
*
- * ptx_cp_async_barrier(Var barrier_ptr, Expr barrier_offset)
+ * ptx_cp_async_barrier(int barrier_id)
*
*/
TVM_DLL const Op& ptx_cp_async_barrier();
/*!
* \brief tvm intrinsics for ptx barrier initialization of thread count using mbarrier.init
*
- * ptx_init_barrier_thread_count(Var barrier_ptr, Expr barrier_offset, int thread_count)
+ * ptx_init_barrier_thread_count(int barrier_id, int thread_count)
*
*/
TVM_DLL const Op& ptx_init_barrier_thread_count();
/*!
* \brief tvm intrinsics for ptx barrier arrival using mbarrier.arrive
*
- * ptx_arrive_barrier(Var barrier_ptr, Expr barrier_offset)
+ * ptx_arrive_barrier(int barrier_id)
*
*/
TVM_DLL const Op& ptx_arrive_barrier();
/*!
* \brief tvm intrinsic for ptx barrier arrival with expect tx using mbarrier.arrive.expect_tx
*
- * ptx_arrive_barrier_expect_tx(Var barrier_ptr, Expr barrier_offset, int byte_count)
+ * ptx_arrive_barrier_expect_tx(int barrier_id, int byte_count)
Review Comment:
Right now this is the only way to add expected transaction counts to a barrier, correct? Do we need a mechanism aside from `arrive` for adding to the count? This way makes sense when a thread issues a bulk async copy and arrives right away, but if the arrival isn't right away, the transaction count isn't added to the barrier until later.
##########
src/target/source/codegen_cuda.h:
##########
@@ -109,6 +109,13 @@ class CodeGenCUDA final : public CodeGenC {
// Op attribute map
OpAttrMap<bool> op_need_warp_shuffle_ = Op::GetAttrMap<bool>("cuda.need_warp_shuffle");
+ // The name of the barrier array in shared memory
+ const std::string barrier_name_ = "barrier";
Review Comment:
This is definitely fine for now, but thinking ahead a little bit, will this cause issues if we try to have multiple Blocks in the same Thread Block Cluster, each with their own "barrier"? Depending on scoping it might be fine.
##########
src/target/source/codegen_cuda.h:
##########
@@ -109,6 +109,13 @@ class CodeGenCUDA final : public CodeGenC {
// Op attribute map
OpAttrMap<bool> op_need_warp_shuffle_ = Op::GetAttrMap<bool>("cuda.need_warp_shuffle");
+ // The name of the barrier array in shared memory
+ const std::string barrier_name_ = "barrier";
+ // The size of the barrier array in shared memory
+ int barrier_count_ = -1;
+ // The alignment of the barrier array in shared memory
+ const int barrier_alignment_bytes_ = 64;
Review Comment:
Why is this 64 instead of 16?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] csullivan merged pull request #15684: [Hopper TMA] Add intrinsic to create barriers for synchronization
Posted by "csullivan (via GitHub)" <gi...@apache.org>.
csullivan merged PR #15684:
URL: https://github.com/apache/tvm/pull/15684
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] adstraw commented on a diff in pull request #15684: [Hopper TMA] Add intrinsic to create barriers for synchronization
Posted by "adstraw (via GitHub)" <gi...@apache.org>.
adstraw commented on code in PR #15684:
URL: https://github.com/apache/tvm/pull/15684#discussion_r1319140645
##########
src/target/source/codegen_cuda.h:
##########
@@ -109,6 +109,13 @@ class CodeGenCUDA final : public CodeGenC {
// Op attribute map
OpAttrMap<bool> op_need_warp_shuffle_ = Op::GetAttrMap<bool>("cuda.need_warp_shuffle");
+ // The name of the barrier array in shared memory
+ const std::string barrier_name_ = "barrier";
Review Comment:
Good point. Scoping is kernel scope. Let's leave this for now and address problems as they come up.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] csullivan commented on a diff in pull request #15684: [Hopper TMA] Add intrinsic to create barriers for synchronization
Posted by "csullivan (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] JosephTheOctonaut commented on a diff in pull request #15684: [Hopper TMA] Add intrinsic to create barriers for synchronization
Posted by "JosephTheOctonaut (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] adstraw commented on a diff in pull request #15684: [Hopper TMA] Add intrinsic to create barriers for synchronization
Posted by "adstraw (via GitHub)" <gi...@apache.org>.
adstraw commented on code in PR #15684:
URL: https://github.com/apache/tvm/pull/15684#discussion_r1319144179
##########
include/tvm/tir/builtin.h:
##########
@@ -681,43 +680,51 @@ TVM_DLL const Op& ptx_wait_group();
/*!
* \brief tvm intrinsics for ptx async copy barrier using cp.async.mbarrier.arrive
*
- * ptx_cp_async_barrier(Var barrier_ptr, Expr barrier_offset)
+ * ptx_cp_async_barrier(int barrier_id)
*
*/
TVM_DLL const Op& ptx_cp_async_barrier();
/*!
* \brief tvm intrinsics for ptx barrier initialization of thread count using mbarrier.init
*
- * ptx_init_barrier_thread_count(Var barrier_ptr, Expr barrier_offset, int thread_count)
+ * ptx_init_barrier_thread_count(int barrier_id, int thread_count)
*
*/
TVM_DLL const Op& ptx_init_barrier_thread_count();
/*!
* \brief tvm intrinsics for ptx barrier arrival using mbarrier.arrive
*
- * ptx_arrive_barrier(Var barrier_ptr, Expr barrier_offset)
+ * ptx_arrive_barrier(int barrier_id)
*
*/
TVM_DLL const Op& ptx_arrive_barrier();
/*!
* \brief tvm intrinsic for ptx barrier arrival with expect tx using mbarrier.arrive.expect_tx
*
- * ptx_arrive_barrier_expect_tx(Var barrier_ptr, Expr barrier_offset, int byte_count)
+ * ptx_arrive_barrier_expect_tx(int barrier_id, int byte_count)
Review Comment:
Good point. I am tempted to leave this as-is for now and solve it when the need arises.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org