Compiler Explorer

Source code

// how to get SASS:
//
// nvcc -O2 --keep-device-functions --device-c --gpu-architecture=sm_90a -I
// cutlass/include cp_async_bulk.cu
//
// cuobjdump -sass cp_async_bulk.o | c++filt >cp_async_bulk.sass
#include "cute/arch/cluster_sm90.hpp"
#include "cute/arch/copy_sm90.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/uint128.h"

#include <cstdint>

namespace test_sass {
  __device__ uint32_t lane_idx() {                                              
    uint32_t reg;                                                              
    asm("mov.u32 %0, %%laneid;" : "=r"(reg));                            
    return reg;                                                                
  }

__device__ void cp_async_bulk_commit_group() { cute::tma_store_arrive(); }

__device__ void cp_async_wait_group_0() { cute::tma_store_wait<0>(); }

__device__ void cp_async_wait_group_1() { cute::tma_store_wait<1>(); }

__device__ void cp_async_wait_group_2() { cute::tma_store_wait<2>(); }

__device__ void cp_async_wait_group_3() { cute::tma_store_wait<3>(); }

__device__ void cp_async_bulk_g2s(void const *gmem_ptr, uint64_t *mbar_ptr,
                                  void *smem_ptr, int32_t load_bytes) {
  cute::SM90_BULK_COPY_G2S::copy(gmem_ptr, mbar_ptr, smem_ptr, load_bytes);
}

__device__ void cp_async_bulk_g2s_sync_aligned(void const *gmem_ptr,
                                               uint64_t *mbar_ptr,
                                               void *smem_ptr,
                                               int32_t load_bytes) {
  if (cute::elect_one_sync()) {
    cute::SM90_BULK_COPY_G2S::copy(gmem_ptr, mbar_ptr, smem_ptr, load_bytes);
  }
}

__device__ void cp_async_bulk_g2s_lane0(void const *gmem_ptr,
                                        uint64_t *mbar_ptr, void *smem_ptr,
                                        int32_t load_bytes) {
  if (lane_idx() == 0) {
    cute::SM90_BULK_COPY_G2S::copy(gmem_ptr, mbar_ptr, smem_ptr, load_bytes);
  }
}

__device__ void cp_async_bulk_s2g(void const *smem_ptr, void *gmem_ptr,
                                  int32_t store_bytes) {
  cute::SM90_BULK_COPY_S2G::copy(smem_ptr, gmem_ptr, store_bytes);
}

__device__ void cp_async_bulk_s2g_sync_aligned(void const *smem_ptr,
                                               void *gmem_ptr,
                                               int32_t store_bytes) {
  if (cute::elect_one_sync()) {
    cute::SM90_BULK_COPY_S2G::copy(smem_ptr, gmem_ptr, store_bytes);
  }
}

__device__ void cp_async_bulk_s2g_lane0(void const *smem_ptr, void *gmem_ptr,
                                        int32_t store_bytes) {
  if (lane_idx() == 0) {
    cute::SM90_BULK_COPY_S2G::copy(smem_ptr, gmem_ptr, store_bytes);
  }
}

__device__ void cp_async_bulk_prefetch(void const *gmem_ptr,
                                       int32_t load_bytes) {
  cute::SM90_BULK_COPY_G2S::PREFETCH::copy(gmem_ptr, load_bytes);
}

__device__ void cp_async_bulk_prefetch_sync_aligned(void const *gmem_ptr,
                                                    int32_t load_bytes) {
  if (cute::elect_one_sync()) {
    cute::SM90_BULK_COPY_G2S::PREFETCH::copy(gmem_ptr, load_bytes);
  }
}

__device__ void cp_async_bulk_prefetch_lane0(void const *gmem_ptr,
                                             int32_t load_bytes) {
  if (lane_idx() == 0) {
    cute::SM90_BULK_COPY_G2S::PREFETCH::copy(gmem_ptr, load_bytes);
  }
}

} // namespace test_sass