Replace inline PTX by cuda::ptx in cuda::barrier<thread_scope_block>

[bernhardmgruber]

Due to lack of a good example for a SASS test, I used this simple example:

#include <cuda/barrier>
#include <cuda/ptx>

// selects a single leader thread from the block
__device__ bool elect_one() {
  // elect_sync is important to help the optimizer generate a uniform datapath
  return cuda::ptx::elect_sync(~0) && threadIdx.x < 32;
}

__global__ void example_kernel(int* gmem1, double* gmem2) {
  constexpr int tile_size = 1024;
  __shared__ alignas(16)    int smem1[tile_size];
  __shared__ alignas(16) double smem2[tile_size];
  #pragma nv_diag_suppress static_var_with_dynamic_init
  using barrier_t = cuda::barrier<cuda::thread_scope_block>;
  __shared__  barrier_t bar;
  // setup the barrier where only the leader thread arrives
  if (elect_one()) {
    init(&bar, 1);
    // issue two TMA bulk copy operations
    cuda::device::memcpy_async_tx(smem1, gmem1, cuda::aligned_size_t<16>(tile_size * sizeof(int)   ), bar);
    cuda::device::memcpy_async_tx(smem2, gmem2, cuda::aligned_size_t<16>(tile_size * sizeof(double)), bar);
    // arrive and update the barrier's expect_tx with the **total** number of loaded bytes
    (void)cuda::device::barrier_arrive_tx(bar, 1, tile_size * (sizeof(int) + sizeof(double)));
  }
  __syncthreads(); // need to sync so the barrier is set up when the other threads arrive and wait
  // wait for the current barrier phase to complete
  bar.wait_parity(0);
  // process data in smem ...
}

Compiled for sm100, it does differ in SASS a little bit, but the compiler just flipped a branch:

[miscco]

looks technically correct, but the formatting is atrocious. Could we add an else to the conditions, as all early branches are returning?

[miscco]

This is strange, because the first condition takes anything but cluster_shared, so the second one seems wrong

[bernhardmgruber]

Anything that is shared is also cluster_shared. Because the shared memory address space is part of the cluster shared memory space.

[bernhardmgruber]

Here, we trap for any barrier that is in cluster shared memory, but not in the shared memory of the current CTA.

[miscco]

could we turn that into an else if

[fbusato]

could not be a precondition instead?

[github-actions]

😬 CI Workflow Results

🟥 Finished in 4h 41m: Pass: 42%/84 | Total: 7h 57m | Max: 39m 16s | Hits: 99%/25847

See results here.

[fbusato]

cuda::std::terminate() ?

[fbusato]

could be a _CCCL_ASSERT instead?

[bernhardmgruber]

I agree, we should probably assert here. And check the documentation whether we make it clear that barriers ought not live in cluster shared memory.

[fbusato]

Probably not worth to move to their C++ versions

[fbusato]

could not be a precondition instead?