4.11 异步数据复制（Asynchronous Data Copies）

__global__ void stencil_kernel(const float *left, const float *center, const float *right)
{
    // 左 halo（8 个元素）- 中心（32 个元素）- 右 halo（8 个元素）
    __shared__ float buffer[8 + 32 + 8];
    const int tid = threadIdx.x;

    if (tid < 8) {
        buffer[tid] = left[tid]; // 左 halo
    } else if (tid >= 32 - 8) {
        buffer[tid + 16] = right[tid]; // 右 halo
    }
    if (tid < 32) {
        buffer[tid + 8] = center[tid]; // 中心
    }
    __syncthreads();

    // 计算 stencil
}

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

__global__ void stencil_kernel(const float *left, const float *center, const float *right)
{
    auto block = cooperative_groups::this_thread_block();
    auto thread = cooperative_groups::this_thread();
    using barrier_t = cuda::barrier<cuda::thread_scope_block>;
    __shared__ barrier_t barrier;
    __shared__ float buffer[8 + 32 + 8];

    if (block.thread_rank() == 0) {
        init(&barrier, block.size());
    }
    __syncthreads();

    // 版本 1：由各个线程单独发起复制。
    if (tid < 8) {
        cuda::memcpy_async(buffer + tid, left + tid, cuda::aligned_size_t<4>(sizeof(float)), barrier);
        // 或：cuda::memcpy_async(thread, buffer + tid, left + tid,
        //       cuda::aligned_size_t<4>(sizeof(float)), barrier);
    } else if (tid >= 32 - 8) {
        cuda::memcpy_async(buffer + tid + 16, right + tid,
                           cuda::aligned_size_t<4>(sizeof(float)), barrier);
        // 或：cuda::memcpy_async(thread, buffer + tid + 16, right + tid,
        //       cuda::aligned_size_t<4>(sizeof(float)), barrier);
    }
    if (tid < 32) {
        cuda::memcpy_async(buffer + tid + 8, center + tid,
                           cuda::aligned_size_t<4>(sizeof(float)), barrier);
        // 或：cuda::memcpy_async(thread, buffer + tid + 8, center + tid,
        //       cuda::aligned_size_t<4>(sizeof(float)), barrier);
    }

    // 版本 2：在所有线程之间协作发起复制。
    cuda::memcpy_async(block, buffer, left,
                       cuda::aligned_size_t<4>(8 * sizeof(float)), barrier);
    cuda::memcpy_async(block, buffer + 8, center,
                       cuda::aligned_size_t<4>(32 * sizeof(float)), barrier);
    cuda::memcpy_async(block, buffer + 40, right,
                       cuda::aligned_size_t<4>(8 * sizeof(float)), barrier);

    barrier.arrive_and_wait();
    __syncthreads();

    // 计算 stencil
}

#include <cooperative_groups.h>
#include <cooperative_groups/memcpy_async.h>

namespace cg = cooperative_groups;

__global__ void stencil_kernel(const float *left, const float *center, const float *right)
{
    cg::thread_block block = cg::this_thread_block();
    __shared__ float buffer[8 + 32 + 8];

    cg::memcpy_async(block, buffer, left, 8 * sizeof(float));
    cg::memcpy_async(block, buffer + 8, center, 32 * sizeof(float));
    cg::memcpy_async(block, buffer + 40, right, 8 * sizeof(float));
    cg::wait(block);
    __syncthreads();

    // 计算 stencil。
}

#include <cuda_pipeline.h>

__global__ void stencil_kernel(const float *left, const float *center, const float *right)
{
    __shared__ float buffer[8 + 32 + 8];
    const int tid = threadIdx.x;

    if (tid < 8) {
        __pipeline_memcpy_async(buffer + tid, left + tid, sizeof(float));
    } else if (tid >= 32 - 8) {
        __pipeline_memcpy_async(buffer + tid + 16, right + tid, sizeof(float));
    }
    if (tid < 32) {
        __pipeline_memcpy_async(buffer + tid + 8, center + tid, sizeof(float));
    }
    __pipeline_commit();
    __pipeline_wait_prior(0);
    __syncthreads();

    // 计算 stencil。
}

#include <cooperative_groups.h>
#include <cuda/pipeline>

template <size_t num_stages = 2>
__global__ void prefetch_kernel(int* global_out, int const* global_in,
                                size_t size, size_t batch_size) {
    auto grid = cooperative_groups::this_grid();
    auto block = cooperative_groups::this_thread_block();
    assert(size == batch_size * grid.size());

    extern __shared__ int shared[];
    size_t shared_offset[num_stages];
    for (int s = 0; s < num_stages; ++s) shared_offset[s] = s * block.size();

    cuda::pipeline<cuda::thread_scope_thread> pipeline = cuda::make_pipeline();
    auto block_batch = [&](size_t batch) -> int {
        return block.group_index().x * block.size() + grid.size() * batch;
    };

    for (int s = 0; s < num_stages; ++s) {
        pipeline.producer_acquire();
        cuda::memcpy_async(shared + shared_offset[s] + tid,
                           global_in + block_batch(s) + tid,
                           cuda::aligned_size_t<4>(sizeof(int)), pipeline);
        pipeline.producer_commit();
    }

    int stage = 0;
    for (size_t compute_batch = 0, fetch_batch = num_stages;
         compute_batch < batch_size; ++compute_batch, ++fetch_batch) {
        constexpr size_t pending_batches = num_stages - 1;
        cuda::pipeline_consumer_wait_prior<pending_batches>(pipeline);
        __syncthreads();
        compute(global_out + block_batch(compute_batch) + tid,
                shared + shared_offset[stage] + tid);
        pipeline.consumer_release();
        __syncthreads();
        pipeline.producer_acquire();
        if (fetch_batch < batch_size) {
            cuda::memcpy_async(shared + shared_offset[stage] + tid,
                               global_in + block_batch(fetch_batch) + tid,
                               cuda::aligned_size_t<4>(sizeof(int)), pipeline);
        }
        pipeline.producer_commit();
        stage = (stage + 1) % num_stages;
    }
}

#include <cooperative_groups.h>
#include <cooperative_groups/memcpy_async.h>

namespace cg = cooperative_groups;

template <size_t num_stages = 2>
__global__ void prefetch_kernel(int* global_out, int const* global_in,
                                size_t size, size_t batch_size) {
    auto grid = cg::this_grid();
    auto block = cg::this_thread_block();
    assert(size == batch_size * grid.size());
    extern __shared__ int shared[];
    size_t shared_offset[num_stages];
    for (int s = 0; s < num_stages; ++s) shared_offset[s] = s * block.size();

    for (int s = 0; s < num_stages; ++s) {
        cg::memcpy_async(block, shared + shared_offset[s],
                         global_in + block_batch(s),
                         cuda::aligned_size_t<4>(sizeof(int)) * block.size());
    }
    int stage = 0;
    for (size_t compute_batch = 0, fetch_batch = num_stages;
         compute_batch < batch_size; ++compute_batch, ++fetch_batch) {
        size_t pending_batches =
            (fetch_batch < batch_size - num_stages) ? num_stages - 1
                                                   : batch_size - fetch_batch - 1;
        cg::wait_prior(pending_batches);
        __syncthreads();
        compute(global_out + block_batch(compute_batch) + tid,
                shared + shared_offset[stage] + tid);
        __syncthreads();
        if (fetch_batch < batch_size) {
            cg::memcpy_async(block, shared + shared_offset[stage],
                             global_in + block_batch(fetch_batch),
                             cuda::aligned_size_t<4>(sizeof(int)) * block.size());
        }
        stage = (stage + 1) % num_stages;
    }
}

#include <cooperative_groups.h>
#include <cuda_awbarrier_primitives.h>

template <size_t num_stages = 2>
__global__ void prefetch_kernel(int* global_out, int const* global_in,
                                size_t size, size_t batch_size) {
    auto grid = cooperative_groups::this_grid();
    auto block = cooperative_groups::this_thread_block();
    extern __shared__ int shared[];
    size_t shared_offset[num_stages];
    for (int s = 0; s < num_stages; ++s) shared_offset[s] = s * block.size();

    auto block_batch = [&](size_t batch) -> int {
        return block.group_index().x * block.size() + grid.size() * batch;
    };
    for (int s = 0; s < num_stages; ++s) {
        __pipeline_memcpy_async(shared + shared_offset[s] + tid,
                                global_in + block_batch(s) + tid,
                                cuda::aligned_size_t<4>(sizeof(int)));
        __pipeline_commit();
    }
    int stage = 0;
    for (size_t compute_batch = 0, fetch_batch = num_stages;
         compute_batch < batch_size; ++compute_batch, ++fetch_batch) {
        constexpr size_t pending_batches = num_stages - 1;
        __pipeline_wait_prior<pending_batches>();
        __syncthreads();
        compute(global_out + block_batch(compute_batch) + tid,
                shared + shared_offset[stage] + tid);
        __syncthreads();
        if (fetch_batch < batch_size) {
            __pipeline_memcpy_async(shared + shared_offset[stage] + tid,
                                    global_in + block_batch(fetch_batch) + tid,
                                    cuda::aligned_size_t<4>(sizeof(int)));
        }
        __pipeline_commit();
        stage = (stage + 1) % num_stages;
    }
}

#include <cooperative_groups.h>
#include <cuda/pipeline>

using pipeline = cuda::pipeline<cuda::thread_scope_block>;

__device__ void produce(pipeline &pipe, int num_stages, int stage,
                        int num_batches, int batch, float *buffer,
                        int buffer_len, float *in, int N) {
  if (batch < num_batches) {
    pipe.producer_acquire();
    cuda::memcpy_async(buffer + stage * buffer_len + threadIdx.x,
                       in + batch * buffer_len + threadIdx.x,
                       cuda::aligned_size_t<4>(sizeof(float)), pipe);
    pipe.producer_commit();
  }
}

__device__ void consume(pipeline &pipe, int num_stages, int stage,
                        int num_batches, int batch, float *buffer,
                        int buffer_len, float *out, int N) {
  pipe.consumer_wait();
  // 消费 buffer(stage) 并更新 out(batch)
  pipe.consumer_release();
}

__global__ void producer_consumer_pattern(float *in, float *out,
                                          int N, int buffer_len) {
  auto block = cooperative_groups::this_thread_block();
  constexpr int warpSize = 32;
  __shared__ extern float buffer[]; // 大小为 2 * buffer_len
  const int num_batches = N / buffer_len;

  constexpr auto scope = cuda::thread_scope_block;
  constexpr int num_stages = 2;
  cuda::std::size_t producer_count = warpSize;
  __shared__ cuda::pipeline_shared_state<scope, num_stages> shared_state;
  pipeline pipe = cuda::make_pipeline(block, &shared_state, producer_count);

  if (block.thread_rank() < producer_count)
    for (int s = 0; s < num_stages; ++s)
      produce(pipe, num_stages, s, num_batches, s, buffer, buffer_len, in, N);

  int stage = 0;
  for (size_t b = 0; b < num_batches; ++b) {
    if (block.thread_rank() < producer_count)
      produce(pipe, num_stages, stage, num_batches, b + num_stages,
              buffer, buffer_len, in, N);
    else
      consume(pipe, num_stages, stage, num_batches, b,
              buffer, buffer_len, out, N);
    stage = (stage + 1) % num_stages;
  }
}

#include <cooperative_groups.h>
#include <cuda_awbarrier_primitives.h>

__device__ void produce(__mbarrier_t ready[], __mbarrier_t filled[],
                        float *buffer, int buffer_len, float *in, int N) {
  for (int i = 0; i < N / buffer_len; ++i) {
    __mbarrier_token_t token = __mbarrier_arrive(&ready[i % 2]);
    while (!__mbarrier_try_wait(&ready[i % 2], token, 1000)) {}
    __pipeline_memcpy_async(buffer + i * buffer_len + threadIdx.x,
                            in + i * buffer_len + threadIdx.x,
                            cuda::aligned_size_t<4>(sizeof(float)));
    __pipeline_arrive_on(filled[i % 2]);
    __mbarrier_arrive(filled[i % 2]);
  }
}

__device__ void consume(__mbarrier_t ready[], __mbarrier_t filled[],
                        float *buffer, int buffer_len, float *out, int N) {
  __mbarrier_arrive(&ready[0]);
  __mbarrier_arrive(&ready[1]);
  for (int i = 0; i < N / buffer_len; ++i) {
    __mbarrier_token_t token = __mbarrier_arrive(&filled[i % 2]);
    while (!__mbarrier_try_wait(&filled[i % 2], token, 1000)) {}
    __mbarrier_arrive(&ready[i % 2]);
  }
}

__global__ void producer_consumer_pattern(int N, float *in, float *out,
                                          int buffer_len) {
  __shared__ extern float buffer[];
  __shared__ __mbarrier_t bar[4];
  auto block = cooperative_groups::this_thread_block();
  if (block.thread_rank() < 4)
    __mbarrier_init(bar + block.thread_rank(), block.size());
  __syncthreads();
  if (block.thread_rank() < warpSize)
    produce(bar, bar + 2, buffer, buffer_len, in, N);
  else
    consume(bar, bar + 2, buffer, buffer_len, out, N);
}

#if defined(__CUDA_MINIMUM_ARCH__) && __CUDA_MINIMUM_ARCH__ < 900
static_assert(false, "Device code is being compiled with older architectures that are incompatible with TMA.");
#endif // __CUDA_MINIMUM_ARCH__

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

using barrier = cuda::barrier<cuda::thread_scope_block>;
namespace ptx = cuda::ptx;
static constexpr size_t buf_len = 1024;

__device__ inline bool is_elected() {
    unsigned int tid = threadIdx.x;
    unsigned int warp_id = tid / 32;
    unsigned int uniform_warp_id = __shfl_sync(0xFFFFFFFF, warp_id, 0);
    return (uniform_warp_id == 0 && ptx::elect_sync(0xFFFFFFFF));
}

__global__ void add_one_kernel(int* data, size_t offset) {
  __shared__ alignas(16) int smem_data[buf_len];
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ barrier bar;
  if (threadIdx.x == 0) init(&bar, blockDim.x);
  __syncthreads();

  if (is_elected()) {
    cuda::memcpy_async(smem_data, data + offset,
                       cuda::aligned_size_t<16>(sizeof(smem_data)), bar);
    // 也可以使用 cuda::device::memcpy_async_tx 或 cuda::ptx::cp_async_bulk。
  }

  barrier::arrival_token token = bar.arrive();
  bar.wait(std::move(token));

  for (int i = threadIdx.x; i < buf_len; i += blockDim.x)
    smem_data[i] += 1;

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

  if (is_elected()) {
    ptx::cp_async_bulk(ptx::space_global, ptx::space_shared,
                       data + offset, smem_data, sizeof(smem_data));
    ptx::cp_async_bulk_commit_group();
    ptx::cp_async_bulk_wait_group_read(ptx::n32_t<0>());
  }
}

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

namespace ptx = cuda::ptx;
namespace cg = cooperative_groups;

__device__ inline bool is_elected() {
    unsigned int tid = threadIdx.x;
    unsigned int warp_id = tid / 32;
    unsigned int uniform_warp_id = __shfl_sync(0xFFFFFFFF, warp_id, 0);
    return (uniform_warp_id == 0 && ptx::elect_sync(0xFFFFFFFF));
}

template <int block_size, int num_stages>
__global__ void prefetch_kernel(int* global_out, int const* global_in,
                                size_t size, size_t batch_size) {
    auto grid = cg::this_grid();
    auto block = cg::this_thread_block();
    const int tid = threadIdx.x;
    extern __shared__ int shared[];
    size_t shared_offset[num_stages];
    for (int s = 0; s < num_stages; ++s) shared_offset[s] = s * block.size();
    auto block_batch = [&](size_t batch) -> int {
        return block.group_index().x * block.size() + grid.size() * batch;
    };

    #pragma nv_diag_suppress static_var_with_dynamic_init
    __shared__ cuda::barrier<cuda::thread_scope_block> bar[num_stages];
    if (tid == 0) {
        #pragma unroll num_stages
        for (int i = 0; i < num_stages; ++i) init(&bar[i], 1);
    }
    __syncthreads();

    if (is_elected()) {
        size_t num_bytes = block_size * sizeof(int);
        #pragma unroll num_stages
        for (int s = 0; s < num_stages; ++s) {
            cuda::device::memcpy_async_tx(&shared[shared_offset[s]],
                &global_in[block_batch(s)],
                cuda::aligned_size_t<16>(num_bytes), bar[s]);
            (void)cuda::device::barrier_arrive_tx(bar[s], 1, num_bytes);
        }
    }

    int stage = 0;
    uint32_t parity = 0;
    for (size_t compute_batch = 0, fetch_batch = num_stages;
         compute_batch < batch_size; ++compute_batch, ++fetch_batch) {
        while (!ptx::mbarrier_try_wait_parity(
                   ptx::sem_acquire, ptx::scope_cta,
                   cuda::device::barrier_native_handle(bar[stage]), parity)) {}
        compute(global_out + block_batch(compute_batch) + tid,
                shared + shared_offset[stage] + tid);
        __syncthreads();
        if (is_elected() && fetch_batch < batch_size) {
            size_t num_bytes = block_size * sizeof(int);
            cuda::device::memcpy_async_tx(&shared[shared_offset[stage]],
                &global_in[block_batch(fetch_batch)],
                cuda::aligned_size_t<16>(num_bytes), bar[stage]);
            (void)cuda::device::barrier_arrive_tx(bar[stage], 1, num_bytes);
        }
        stage++;
        if (stage == num_stages) { stage = 0; parity ^= 1; }
    }
}

#include <cudaTypedefs.h>

PFN_cuTensorMapEncodeTiled_v12000 get_cuTensorMapEncodeTiled() {
  cudaDriverEntryPointQueryResult driver_status;
  void* cuTensorMapEncodeTiled_ptr = nullptr;
  CUDA_CHECK(cudaGetDriverEntryPointByVersion(
      "cuTensorMapEncodeTiled", &cuTensorMapEncodeTiled_ptr, 12000,
      cudaEnableDefault, &driver_status));
  assert(driver_status == cudaDriverEntryPointSuccess);
  return reinterpret_cast<PFN_cuTensorMapEncodeTiled_v12000>(cuTensorMapEncodeTiled_ptr);
}

CUtensorMap tensor_map{};
constexpr uint32_t rank = 2;
uint64_t size[rank] = {GMEM_WIDTH, GMEM_HEIGHT};
uint64_t stride[rank - 1] = {GMEM_WIDTH * sizeof(int)};
uint32_t box_size[rank] = {SMEM_WIDTH, SMEM_HEIGHT};
uint32_t elem_stride[rank] = {1, 1};

auto cuTensorMapEncodeTiled = get_cuTensorMapEncodeTiled();
CUresult res = cuTensorMapEncodeTiled(
  &tensor_map,
  CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_INT32,
  rank, tensor_ptr, size, stride, box_size, elem_stride,
  CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
  CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE,
  CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
  CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

#include <cuda.h>

__global__ void kernel(const __grid_constant__ CUtensorMap tensor_map) {
   // 在这里使用 tensor_map。
}
int main() {
  CUtensorMap map;
  // 初始化 map
  kernel<<<1, 1>>>(map);
}

#include <cuda.h>

__constant__ CUtensorMap global_tensor_map;
__global__ void kernel() {
  // 在这里使用 global_tensor_map。
}
int main() {
  CUtensorMap local_tensor_map;
  // 初始化 map
  cudaMemcpyToSymbol(global_tensor_map, &local_tensor_map, sizeof(CUtensorMap));
  kernel<<<1, 1>>>();
}

#include <cuda.h>
#include <cuda/ptx>
namespace ptx = cuda::ptx;

__device__ CUtensorMap global_tensor_map;
__global__ void kernel(CUtensorMap *tensor_map) {
  ptx::n32_t<128> size_bytes;
  ptx::fence_proxy_tensormap_generic(
      ptx::sem_acquire, ptx::scope_sys, tensor_map, size_bytes);
  // fence 后可以安全使用 tensor_map。
}

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

using barrier = cuda::barrier<cuda::thread_scope_block>;
namespace ptx = cuda::ptx;

__device__ inline bool is_elected() {
    unsigned int tid = threadIdx.x;
    unsigned int warp_id = tid / 32;
    unsigned int uniform_warp_id = __shfl_sync(0xFFFFFFFF, warp_id, 0);
    return (uniform_warp_id == 0 && ptx::elect_sync(0xFFFFFFFF));
}

__global__ void kernel(const __grid_constant__ CUtensorMap tensor_map,
                       int x, int y) {
  __shared__ alignas(128) int smem_buffer[SMEM_HEIGHT][SMEM_WIDTH];
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ barrier bar;
  if (threadIdx.x == 0) init(&bar, blockDim.x);
  __syncthreads();

  barrier::arrival_token token;
  if (is_elected()) {
    int32_t tensor_coords[2] = {x, y};
    ptx::cp_async_bulk_tensor(
        ptx::space_shared, ptx::space_global, &smem_buffer,
        &tensor_map, tensor_coords, cuda::device::barrier_native_handle(bar));
    token = cuda::device::barrier_arrive_tx(bar, 1, sizeof(smem_buffer));
  } else {
    token = bar.arrive();
  }
  bar.wait(std::move(token));

  smem_buffer[0][threadIdx.x] += threadIdx.x;
  ptx::fence_proxy_async(ptx::space_shared);
  __syncthreads();

  if (is_elected()) {
    int32_t tensor_coords[2] = {x, y};
    ptx::cp_async_bulk_tensor(
        ptx::space_global, ptx::space_shared, &tensor_map,
        tensor_coords, &smem_buffer);
    ptx::cp_async_bulk_commit_group();
    ptx::cp_async_bulk_wait_group_read(ptx::n32_t<0>());
  }
  if (threadIdx.x == 0) (&bar)->~barrier();
}

// 初始化设备上下文
CUDA_CHECK(cudaDeviceSynchronize());

// 使用 cuTensorMapEncodeTiled Driver API 创建 tensor map 模板
CUtensorMap template_tensor_map = make_tensormap_template();

// 在全局内存中分配 tensor map 和 tensor
CUtensorMap* global_tensor_map;
CUDA_CHECK(cudaMalloc(&global_tensor_map, sizeof(CUtensorMap)));
char* global_buf;
CUDA_CHECK(cudaMalloc(&global_buf, 8 * 256));

// 使用数据填充全局 buffer
fill_global_buf<<<1, 1>>>(global_buf);

// 定义要在设备端创建的 tensor map 参数
tensormap_params p{};
p.global_address    = global_buf;
p.rank              = 2;
p.box_dim[0]        = 128;
p.box_dim[1]        = 4;
p.global_dim[0]     = 256;
p.global_dim[1]     = 8;
p.global_stride[0]  = 256;
p.element_stride[0] = 1;
p.element_stride[1] = 1;

// 在设备端编码 global_tensor_map
encode_tensor_map<<<1, 32>>>(template_tensor_map, p, global_tensor_map);

// 在另一个 kernel 中使用
consume_tensor_map<<<1, 1>>>(global_tensor_map);

// 检查错误
CUDA_CHECK(cudaDeviceSynchronize());

struct tensormap_params {
  void* global_address;
  int rank;
  uint32_t box_dim[5];
  uint64_t global_dim[5];
  size_t global_stride[4];
  uint32_t element_stride[5];
};

// 初始化设备上下文、创建模板、分配 tensor map 和数据，然后：
encode_tensor_map<<<1, 32>>>(template_tensor_map, p, global_tensor_map);
consume_tensor_map<<<1, 1>>>(global_tensor_map);
CUDA_CHECK(cudaDeviceSynchronize());

#include <cuda/ptx>
namespace ptx = cuda::ptx;

__launch_bounds__(32)
__global__ void encode_tensor_map(
    const __grid_constant__ CUtensorMap template_tensor_map,
    tensormap_params p, CUtensorMap* out) {
  __shared__ alignas(128) CUtensorMap smem_tmap;
  if (threadIdx.x == 0) {
    smem_tmap = template_tensor_map;
    const auto space_shared = ptx::space_shared;
    ptx::tensormap_replace_global_address(space_shared, &smem_tmap, p.global_address);
    ptx::tensormap_replace_rank(space_shared, &smem_tmap, p.rank - 1);
    if (0 < p.rank) ptx::tensormap_replace_box_dim(space_shared, &smem_tmap, ptx::n32_t<0>{}, p.box_dim[0]);
    if (1 < p.rank) ptx::tensormap_replace_box_dim(space_shared, &smem_tmap, ptx::n32_t<1>{}, p.box_dim[1]);
    if (2 < p.rank) ptx::tensormap_replace_box_dim(space_shared, &smem_tmap, ptx::n32_t<2>{}, p.box_dim[2]);
    if (3 < p.rank) ptx::tensormap_replace_box_dim(space_shared, &smem_tmap, ptx::n32_t<3>{}, p.box_dim[3]);
    if (4 < p.rank) ptx::tensormap_replace_box_dim(space_shared, &smem_tmap, ptx::n32_t<4>{}, p.box_dim[4]);
    if (0 < p.rank) ptx::tensormap_replace_global_dim(space_shared, &smem_tmap, ptx::n32_t<0>{}, (uint32_t)p.global_dim[0]);
    if (1 < p.rank) ptx::tensormap_replace_global_dim(space_shared, &smem_tmap, ptx::n32_t<1>{}, (uint32_t)p.global_dim[1]);
    if (2 < p.rank) ptx::tensormap_replace_global_dim(space_shared, &smem_tmap, ptx::n32_t<2>{}, (uint32_t)p.global_dim[2]);
    if (3 < p.rank) ptx::tensormap_replace_global_dim(space_shared, &smem_tmap, ptx::n32_t<3>{}, (uint32_t)p.global_dim[3]);
    if (4 < p.rank) ptx::tensormap_replace_global_dim(space_shared, &smem_tmap, ptx::n32_t<4>{}, (uint32_t)p.global_dim[4]);
    if (1 < p.rank) ptx::tensormap_replace_global_stride(space_shared, &smem_tmap, ptx::n32_t<0>{}, p.global_stride[0]);
    if (2 < p.rank) ptx::tensormap_replace_global_stride(space_shared, &smem_tmap, ptx::n32_t<1>{}, p.global_stride[1]);
    if (3 < p.rank) ptx::tensormap_replace_global_stride(space_shared, &smem_tmap, ptx::n32_t<2>{}, p.global_stride[2]);
    if (4 < p.rank) ptx::tensormap_replace_global_stride(space_shared, &smem_tmap, ptx::n32_t<3>{}, p.global_stride[3]);
    if (0 < p.rank) ptx::tensormap_replace_element_size(space_shared, &smem_tmap, ptx::n32_t<0>{}, p.element_stride[0]);
    if (1 < p.rank) ptx::tensormap_replace_element_size(space_shared, &smem_tmap, ptx::n32_t<1>{}, p.element_stride[1]);
    if (2 < p.rank) ptx::tensormap_replace_element_size(space_shared, &smem_tmap, ptx::n32_t<2>{}, p.element_stride[2]);
    if (3 < p.rank) ptx::tensormap_replace_element_size(space_shared, &smem_tmap, ptx::n32_t<3>{}, p.element_stride[3]);
    if (4 < p.rank) ptx::tensormap_replace_element_size(space_shared, &smem_tmap, ptx::n32_t<4>{}, p.element_stride[4]);
    auto zero = ptx::n32_t<0>{};
    ptx::tensormap_replace_elemtype(space_shared, &smem_tmap, zero);
    ptx::tensormap_replace_interleave_layout(space_shared, &smem_tmap, zero);
    ptx::tensormap_replace_swizzle_mode(space_shared, &smem_tmap, zero);
    ptx::tensormap_replace_fill_mode(space_shared, &smem_tmap, zero);
  }
  __syncwarp();
  ptx::n32_t<128> bytes_128;
  ptx::tensormap_cp_fenceproxy(ptx::sem_release, ptx::scope_gpu,
                               out, &smem_tmap, bytes_128);
}

__global__ void consume_tensor_map(CUtensorMap* tensor_map) {
  ptx::n32_t<128> size_bytes;
  ptx::fence_proxy_tensormap_generic(
      ptx::sem_acquire, ptx::scope_sys, tensor_map, size_bytes);
  // fence 后可以安全使用 tensor_map。

  __shared__ uint64_t bar;
  __shared__ alignas(128) char smem_buf[4][128];
  if (threadIdx.x == 0) {
    ptx::mbarrier_init(&bar, 1);
    ptx::cp_async_bulk_tensor(ptx::space_cluster, ptx::space_global,
                              smem_buf, tensor_map, {0, 0}, &bar);
    ptx::mbarrier_arrive_expect_tx(ptx::sem_release, ptx::scope_cta,
                                   ptx::space_shared, &bar, sizeof(smem_buf));
  }
  const int parity = 0;
  while (!ptx::mbarrier_try_wait_parity(&bar, parity)) {}
  for (int j = 0; j < 4; ++j) {
    for (int i = 0; i < 128; ++i) {
      printf("%3d ", smem_buf[j][i]);
      if (i % 32 == 31) printf("\n");
    }
    printf("\n");
  }
}

CUtensorMap make_tensormap_template() {
  CUtensorMap template_tensor_map{};
  auto cuTensorMapEncodeTiled = get_cuTensorMapEncodeTiled();
  uint32_t dims_32 = 16;
  uint64_t dims_strides_64 = 16;
  uint32_t elem_strides = 1;
  CUresult res = cuTensorMapEncodeTiled(
    &template_tensor_map,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8,
    1, nullptr, &dims_strides_64, &dims_strides_64,
    &dims_32, &elem_strides,
    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE,
    CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
  CU_CHECK(res);
  return template_tensor_map;
}

__global__ void kernel_tma(const __grid_constant__ CUtensorMap tensor_map) {
   __shared__ alignas(1024) int4 smem_buffer[8][8];
   __shared__ alignas(1024) int4 smem_buffer_tr[8][8];
   #pragma nv_diag_suppress static_var_with_dynamic_init
   __shared__ barrier bar;
   if (threadIdx.x == 0) init(&bar, blockDim.x);
   __syncthreads();

   barrier::arrival_token token;
   if (is_elected()) {
     int32_t tensor_coords[2] = {0, 0};
     ptx::cp_async_bulk_tensor(ptx::space_shared, ptx::space_global,
                               &smem_buffer, &tensor_map, tensor_coords,
                               cuda::device::barrier_native_handle(bar));
     token = cuda::device::barrier_arrive_tx(bar, 1, sizeof(smem_buffer));
   } else token = bar.arrive();
   bar.wait(std::move(token));

   for (int sidx_j = threadIdx.x; sidx_j < 8; sidx_j += blockDim.x) {
      for (int sidx_i = 0; sidx_i < 8; ++sidx_i) {
         const int swiz_j_idx = (sidx_i % 8) ^ sidx_j;
         const int swiz_i_idx_tr = (sidx_j % 8) ^ sidx_i;
         smem_buffer_tr[sidx_j][swiz_i_idx_tr] = smem_buffer[sidx_i][swiz_j_idx];
      }
   }
   ptx::fence_proxy_async(ptx::space_shared);
   __syncthreads();
   if (is_elected()) {
       int32_t tensor_coords[2] = {x, y};
       ptx::cp_async_bulk_tensor(ptx::space_global, ptx::space_shared,
                                 &tensor_map, tensor_coords, &smem_buffer_tr);
       ptx::cp_async_bulk_commit_group();
       ptx::cp_async_bulk_wait_group_read(ptx::n32_t<0>());
   }
   if (threadIdx.x == 0) (&bar)->~barrier();
}

void* tensor_ptr = d_data;
CUtensorMap tensor_map{};
constexpr uint32_t rank = 2;
uint64_t size[rank] = {4 * 8, 8};
uint64_t stride[rank - 1] = {8 * sizeof(int4)};
// 内层 shared-memory box 维度（字节数）等于 swizzle span。
uint32_t box_size[rank] = {4 * 8, 8};
uint32_t elem_stride[rank] = {1, 1};

CUresult res = cuTensorMapEncodeTiled(
    &tensor_map,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_INT32,
    rank, tensor_ptr, size, stride, box_size, elem_stride,
    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
    CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

kernel_tma<<<1, 8>>>(tensor_map);

data_t* smem_ptr = &smem[0][0];
int offset = (reinterpret_cast<uintptr_t>(smem_ptr)/128)%8;
smem[y][((y+offset)%8)^x] = ...;

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

__global__ __cluster_dims__(8, 1, 1) void producer_consumer_kernel() {
    using namespace cooperative_groups;
    using namespace cuda::device;
    using namespace cuda::ptx;
    using barrier_t = cuda::barrier<cuda::thread_scope_block>;
    auto cluster = this_cluster();
    __shared__ int buffer[BLOCK_SIZE];
    __shared__ barrier_t filled;
    __shared__ barrier_t ready;

    if (threadIdx.x == 0) {
        init(&filled, 1);
        init(&ready, BLOCK_SIZE);
    }
    cluster.sync();

    int rk = cluster.block_rank();
    int rk_next = (rk + 1) % 8;
    int rk_prev = (rk + 7) % 8;
    auto buffer_next = cluster.map_shared_rank(buffer, rk_next);
    auto bar_next = cluster.map_shared_rank(barrier_native_handle(filled), rk_next);
    auto bar_prev = cluster.map_shared_rank(barrier_native_handle(ready), rk_prev);

    int phase = 0;
    for (int it = 0; it < 1000; ++it) {
        st_async(&buffer_next[threadIdx.x], rk, bar_next);
        if (threadIdx.x == 0) {
            mbarrier_arrive_expect_tx(sem_release, scope_cluster, space_shared,
                                      barrier_native_handle(filled), sizeof(buffer));
        }
        while (!mbarrier_try_wait_parity(barrier_native_handle(filled), phase, 1000)) {}
        int r = buffer[threadIdx.x];
        // 使用数据执行某些操作。
        mbarrier_arrive(sem_release, scope_cluster, space_cluster, bar_prev);
        while (!mbarrier_try_wait_parity(barrier_native_handle(ready), phase, 1000)) {}
        phase ^= 1;
    }
}