`cuda::memcpy_async` with `cuda::barrier` implementation is inefficient on sm90+

[bernhardmgruber]

I received the following piece of code (shortened) using cuda::memcpy_async and its author claims the codegen is inefficient:

#include <cuda/barrier>
#include <cuda/ptx>
#include <cooperative_groups.h>

__device__ static inline bool elect_sync(const ::cuda::std::uint32_t& __membermask) { ... }

constexpr int elem_per_block = 128;
constexpr int num_bytes = elem_per_block * sizeof(float);

using barrier = cuda::barrier<cuda::thread_scope_block>;
using aligned_16 = cuda::aligned_size_t<16>;
namespace cg = cooperative_groups;

__global__ void tma_add(
  const float* a, const float* b, float* c, size_t len) {
    __shared__ alignas(16) float smem_a[elem_per_block];
    __shared__ alignas(16) float smem_b[elem_per_block];
    __shared__ barrier bar;
    auto block = cg::this_thread_block();

    if (threadIdx.x == 0) {
        init(&bar, blockDim.x);
    }
    __syncthreads();

    cuda::memcpy_async(
        block,
        smem_a,
        a,
        aligned_16(num_bytes),
        bar
    );
    bar.arrive_and_wait();
}

Full version here: https://godbolt.org/z/jGz884cE5

If the call to cuda::memcpy_async is replaced with

    int warp_id = threadIdx.x / 32;
    bool elected = warp_id == 0 && elect_sync(~0);
    if (elected) {
      cuda::ptx::cp_async_bulk(cuda::ptx::space_cluster, cuda::ptx::space_global, 
        smem_a, a, num_bytes, cuda::device::barrier_native_handle(bar));
      cuda::ptx::mbarrier_expect_tx(cuda::ptx::sem_relaxed, cuda::ptx::scope_cta, cuda::ptx::space_shared,
        cuda::device::barrier_native_handle(bar), num_bytes);
    }

the codegen is much better. So what's wrong with cuda::memcpy_async?

After some investigation I noticed the following problems:

• our __cp_async_bulk_shared_global elects with __g.thread_rank() == 0 and not with elect_sync, leading to bad codegen because the compiler does not recognize the uniform data path

• inside the bulk copy code path of cuda::memcpy_async, we elect twice (using __g.thread_rank() == 0), once in __cp_async_bulk_shared_global, once for the __memcpy_completion_impl::__defer, both need to elect the same thread, so we cannot use elect_sync. Optimally, this should be in a single if(elected) { ... } containing both the scheduling of the bulk copy and the barrier update

• bar.arrive_and_wait() waits for phase completion with backoff (using cuda::std::chrono, timers and thread yielding). Since sm90, we have hardware support (SYNCS) for waiting for a parity change, so we should use that.

[bernhardmgruber]

Comparing a TMA copy example entirely written with PTX vs. one using cuda::barrier (https://godbolt.org/z/qa1cPbb4f), it looks like only wait_parity is generating more code than necessary.