how each block sync in cluster

[mengchihe]

Hi
In sm90_mma_tma_gmma_ss.hpp, for example cluster shape is <2, 2, 1>, it seems that each block in the cluster will issue two tma.multicast for both inputA and inputB in each stage, and it's full barrier's tx_count is the sum of the data needed in this stage.
My question is that, since each tma.multicast will reduce the tx_count of all barriers in the mask, one block's full barrier will reduced twice and arrived when only half of the data is ready. It will make sense when just block with cluster.x=0 or cluster.y=0 issue tma.multicast, is there some place has this condition that I didn't find out? thanks.

[IonThruster]

Thank you for your question. In your example of cluster shape <2,2,1> the #bytes required to trigger the full barrier is indeed size(A buffer +B buffer) for 1 shared memory stage.

So with each copy( ) op issued in the collective mainloop, if we use the SM90_TMA_LOAD_MULTICAST internally to construct the TMA op - it will load only 50% of A buffer for 1 stage or 50% of B. So in order to trigger the full barrier - we need the completion of 2 TMA loads from the local thread block (for 50% of A and 50% of B) + 1 load from horizontal thread-block neighbor in the cluster for remaining 50% of A + 1 load from vertical thread-block neighbor in the cluster for remaining 50% of B matrix.

As you may have also noticed, this means we need to load the appropriate "slice" of A and B tensors from global memory into shared memory based on the rank / block-id in the cluster, and this is handled by CuTe APIs as well (get_slice and partition).

Do let me if that wasn't clear, or if I have misunderstood your question.

References :

• make_tma_copy
• get_slice and partition

[mengchihe]

Thanks for your reply.
I didn't go inside get_slice and thought it just return all the data this block needs.
thank you very much