TMA with Host memory #2

  int *d_data;
  cudaHostGetDevicePointer(&d_data, h_data, 0);

  CUtensorMap tensor_map =
      make_2d_tma_desc(d_data, global_dim, GLOBAL_K * sizeof(int), tile_dim,
                       CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE);

#include <cuda.h>
#include <cudaTypedefs.h>

#include <cuda/barrier>
#include <iostream>
#pragma nv_diag_suppress static_var_with_dynamic_init

using barrier_t = cuda::barrier<cuda::thread_scope_block>;
namespace cde = cuda::device::experimental;

constexpr size_t GLOBAL_M = 64;
constexpr size_t GLOBAL_K = 32;
constexpr size_t TILE_M = 8;
constexpr size_t TILE_K = 16;

inline PFN_cuTensorMapEncodeTiled get_cuTensorMapEncodeTiled() {
  cudaDriverEntryPointQueryResult driver_status;
  void *cuTensorMapEncodeTiled_ptr = nullptr;

  cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled",
                                   &cuTensorMapEncodeTiled_ptr, 12000,
                                   cudaEnableDefault, &driver_status);
  if (driver_status != cudaDriverEntryPointSuccess)
    throw std::runtime_error("driver_status != cudaDriverEntryPointSuccess");
  return reinterpret_cast<PFN_cuTensorMapEncodeTiled>(
      cuTensorMapEncodeTiled_ptr);
}

CUtensorMap make_2d_tma_desc(int32_t *global_address, uint64_t global_dim[2],
                             uint64_t stride, uint32_t smem_dim[2],
                             CUtensorMapSwizzle swizzle) {
  CUtensorMap tensor_map = {};
  uint64_t global_stride[1] = {stride};
  uint32_t elem_stride[2] = {1, 1};

  auto encode = get_cuTensorMapEncodeTiled();

  auto res =
      encode(&tensor_map, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_INT32,
             2,  // rank =2
             global_address, global_dim, global_stride, smem_dim, elem_stride,
             CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle,
             CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
             CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

  assert(res == CUDA_SUCCESS && "make tma descriptor failed.");
  return tensor_map;
}

__global__ void tma_kernel(const __grid_constant__ CUtensorMap tensor_map,
                           uint32_t x, uint32_t y) {
  __shared__ alignas(128) int tile_smem[TILE_M * TILE_K];
  __shared__ barrier_t bar;

  // 初始化Barrier
  if (threadIdx.x == 0) {
    init(&bar, blockDim.x);
    // 由于TMA调用路径为async proxy, 需要fence保持可见
    // cuda::ptx::fence_proxy_async(cuda::ptx::space_shared);
    cde::fence_proxy_async_shared_cta();  // b)
  }
  __syncthreads();

  barrier_t::arrival_token token;
  if (threadIdx.x == 0) {
    cde::cp_async_bulk_tensor_2d_global_to_shared(tile_smem, &tensor_map, x, y,
                                                  bar);
    token = cuda::device::barrier_arrive_tx(bar, 1, sizeof(tile_smem));
  } else {
    token = bar.arrive();
  }

  // Wait for the data to have arrived.
  bar.wait(std::move(token));
  printf("[tma_kernel] threadIdx.x %d arrived\n", threadIdx.x);
  for (int i = 0; i < TILE_M * TILE_K; i += blockDim.x) {
    tile_smem[i + threadIdx.x] += threadIdx.x * 100;
  }

  cuda::ptx::fence_proxy_async(cuda::ptx::space_shared);
  __syncthreads();

  if (threadIdx.x == 0) {
    cde::cp_async_bulk_tensor_2d_shared_to_global(&tensor_map, x, y, tile_smem);

    cuda::ptx::cp_async_bulk_commit_group();
    // wait all group finished
    cuda::ptx::cp_async_bulk_wait_group_read(cuda::ptx::n32_t<0>());
  }
  printf("thread %d done\n", threadIdx.x);
}

int main(int argc, char **argv) {
  uint64_t global_dim[2] = {GLOBAL_M, GLOBAL_K};
  size_t GLOBAL_SIZE = GLOBAL_K * GLOBAL_M;
  uint32_t tile_dim[2] = {TILE_M, TILE_K};

  int *h_data = nullptr;
  cudaHostAlloc(&h_data, GLOBAL_SIZE * sizeof(int), cudaHostAllocMapped);
  for (size_t i = 0; i < GLOBAL_SIZE; ++i) {
    h_data[i] = 1;
  }

  int *d_data;
  cudaHostGetDevicePointer(&d_data, h_data, 0);

  CUtensorMap tensor_map =
      make_2d_tma_desc(d_data, global_dim, GLOBAL_K * sizeof(int), tile_dim,
                       CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE);

  uint32_t coord_x = 16;
  uint32_t coord_y = 16;

  tma_kernel<<<1, TILE_M * TILE_K>>>(tensor_map, coord_x, coord_y);
  cudaDeviceSynchronize();

  cudaError_t err = cudaGetLastError();
  std::cout << cudaGetErrorString(err) << std::endl;

  for (size_t i = 0; i < GLOBAL_M; ++i) {
    for (size_t j = 0; j < GLOBAL_K; ++j) {
      printf("%5d ", h_data[i * GLOBAL_K + j]);
    }
    printf("\n");
  }

  cudaFree(d_data);

  return 0;
}