incorrect `mbarrier.arrive` used for sm100 gather_A gemm

[Yang-YiFan]

In sonic-moe for sm100, it is mentioned that there is a relay warp arriving at the peer-CTA's mbarrier to signal tcgen05.mma.2cta that the cp.async of gather_A completes. However, this mbarrier arrive violates the ptx memory model. The call stack is as follows:

• the original call to remote arrive
• which is from PipelineTmaCpAsyncUmma
• which inherents from the cute dsl PipelineTmaUmma
• which eventually does a remote arrive here
• which calls the nvvm op nvvm.mbarrier_txn
• and i think it's the mbarrier.arrive op here?.

So the equivalent ptx being called here is mbarrier.cta.arrive.release.shared::cluster (when unspecified, the default scope is cta and default semantic is release). However, this arrive does not generate a release-acquire relation between the relay warp and 2cta mma warp in the cluster scope. There are two legal ways to do this.

Producer cta:
    cp.async addr, val
    cp.async.wait_group  0
    mbarrier.cluster.arrive.release.shared::cluster
Consumer cta:
    mbarrier.cluster.try_wait.acquire.shared::cta
    tcgen05.mma.2cta addr

however, this will emit a MEMBAR.GPU in the hot loop. An optimized way to do this is:

Producer cta:
    cp.async addr, val
    cp.async.wait_group  0
    fence.release.sync_restrict::shared::cta.cluster
    mbarrier.cluster.arrive.relaxed.shared::cluster
Consumer cta:
    mbarrier.cluster.try_wait.acquire.shared::cta
    tcgen05.mma.2cta addr

which avoids the MEMBAR.GPU.

[tridao]

So the main issue is that the mbar arrive instruction we're using has scope cta?
Your suggestions make sense to me.

[tridao]

I'm still confused how this mbar arrive instruction would work normally (for standard gemm) if the scope is cta instead of cluster.

[Yang-YiFan]

So the main issue is that the mbar arrive instruction we're using has scope cta?

Yes. And I believe the cute dsl team is working on a fix to expose all the scope/memory semantic options for the mbarrier nvvm op (rather than using the default all the time).

I'm still confused how this mbar arrive instruction would work normally (for standard gemm) if the scope is cta instead of cluster.

Could you elaborate the scenario you have in mind? In a normal gemm with cluster, we never need to do a remote barrier arrive manually. TMA multicast and 2CTA TMA have their special fence built into the hardware that doesn't require the programmer to do anything special.