正文

请到「今天看啥」查看全文

// 启动TMA传输，将共享内存复制到全局内存
if (threadIdx.x == 0 ) {
cde::cp_async_bulk_tensor_2d_shared_to_global(&tensor_map, x, y,
&smem_buffer);
// 等待TMA传输完成读取共享内存。
// 从上一个批量复制操作创建一个"批量异步组"。
cde::cp_async_bulk_commit_group();
// 等待组完成从共享内存的读取。
cde::cp_async_bulk_wait_group_read< 0 >();
}

// 销毁屏障。这会使屏障的内存区域无效。如果kernel中还有进一步的计算，
// 这允许重用共享内存屏障的内存位置。
if (threadIdx.x == 0 ) {
(&bar)->~barrier();
}
}

并且在张量映射中我们需要调整交织设置

int *d;
  CHECK_CUDA_ERROR(cudaMalloc(&d, SIZE));
  CHECK_CUDA_ERROR(cudaMemcpy(d, h_in, SIZE, cudaMemcpyHostToDevice));
void *tensor_ptr = (void *)d;

  CUtensorMap tensor_map{};
// rank是数组的维度数。
constexpruint32_t rank = 2;
uint64_t size[rank] = {GMEM_WIDTH, GMEM_HEIGHT};
// stride是从一行的第一个元素到下一行第一个元素所需遍历的字节数。它必须是16的倍数。
uint64_t stride[rank - 1] = {GMEM_WIDTH * sizeof(int)};
// box_size是用作TMA传输目标的共享内存缓冲区的大小。
uint32_t box_size[rank] = {SMEM_WIDTH, SMEM_HEIGHT};
// 元素之间的距离，以sizeof(element)为单位。例如，stride为2可以用来只加载复值张量的实部。
uint32_t elem_stride[rank] = {1, 1};

// 创建张量描述符。
  CUresult res = cuTensorMapEncodeTiled(
      &tensor_map,  // CUtensorMap *tensorMap,
      CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_INT32,
      rank,         // cuuint32_t tensorRank,
      tensor_ptr,   // void *globalAddress,
      size,         // const cuuint64_t *globalDim,
      stride,       // const cuuint64_t *globalStrides,
      box_size,     // const cuuint32_t *boxDim,
      elem_stride,  // const cuuint32_t *elementStrides,
      // 交错模式可用于加速加载小于4字节长的值。
      CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
      // 交织可用于避免共享内存bank冲突。
      CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_32B,
      // L2提升可用于将缓存策略的效果扩展到更广泛的L2缓存行集合。
      CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
      // TMA传输会将任何超出边界的元素设置为零。
      CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

  assert(res == CUDA_SUCCESS);

结论

以上我们展示了TMA在Hopper架构中的一个非常简单的应用。此外，交织相当复杂，为了更好地理解它，进一步实验并深入研究这个主题会很有帮助。如果有人有好的学习资源建议，我很乐意听取。你可以在Github(https://github.com/simveit/tma_intro)上找到代码。

Making matrix transpose really fast on Hopper GPUs

2025年5月2日

介绍

在这篇博客中，我想展示如何实现高效的矩阵转置操作，用于Hopper GPU。我将使用原生CUDA API，不使用抽象，因为我相信这是学习硬件细节的好方法。正如你将看到的，使用交织并能够将交织的索引映射到普通索引是非常重要的。不幸的是，这在其他优秀的CUDA编程指南中没有很好地记录。我希望这篇博客能帮助更多的人使用原生CUDA实现高性能的kernel。

交织

可视化交织模式

在实现矩阵转置之前，重要的是我们要理解交织。交织是一种避免共享内存冲突的技术。在以下我们将使用在Hopper GPU上TMA的以下概念。为了更好地理解可能发生的Bank冲突，让我们可视化2d int 矩阵的Bank分配。我们使用以下不使用交织的布局：

const int GMEM_WIDTH = 32;
constint GMEM_HEIGHT = 32;
constint BLOCK_SIZE = 32;
constint SMEM_WIDTH = BLOCK_SIZE;
constint SMEM_HEIGHT = BLOCK_SIZE;

// 创建张量描述符。
CUresult res = cuTensorMapEncodeTiled(
    &tensor_map,  // CUtensorMap *tensorMap,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_INT32,
    rank,         // cuuint32_t tensorRank,
    tensor_ptr,   // void *globalAddress,
    size,         // const cuuint64_t *globalDim,
    stride,       // const cuuint64_t *globalStrides,
    box_size,     // const cuuint32_t *boxDim,
    elem_stride,  // const cuuint32_t *elementStrides,
    // 交织模式可用于加速加载小于4字节长的值。




    

    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    // 交织可用于避免共享内存Bank冲突。
    CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE,
    // L2提升可用于将缓存策略的效果扩展到更广泛的L2缓存行集合。
    CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
    // 任何超出边界的元素都将被TMA设置为零。
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

然后像下面这样填充共享内存块。

smem_buffer[row * BLOCK_SIZE + col] = (row * BLOCK_SIZE + col) % 32;

我们可以可视化这个：

我们看到每一列都被分配给一个Bank。这意味着如果同一warp中的线程访问相同的列，我们将有一个Bank冲突。我们现在可以修改布局，使我们使用128B交织模式

// 创建张量描述符。
CUresult res = cuTensorMapEncodeTiled(
    &tensor_map,  // CUtensorMap *tensorMap,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_INT32,
    rank,         // cuuint32_t tensorRank,
    tensor_ptr,   // void *globalAddress,
    size,         // const cuuint64_t *globalDim,
    stride,       // const cuuint64_t *globalStrides,
    box_size,     // const cuuint32_t *boxDim,
    elem_stride,  // const cuuint32_t *elementStrides,
    // 交织模式可用于加速加载小于4字节长的值。
    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    // 交织可用于避免共享内存Bank冲突。
    CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
    // L2提升可用于将缓存策略的效果扩展到更广泛的L2缓存行集合。
    CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
    // 任何超出边界的元素都将被TMA设置为零。
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

assert(res == CUDA_SUCCESS);

对SMEM中的值进行相同的赋值将产生以下图片：

我们可以看到我们现在有显著更少的潜在Bank冲突。交织模式是周期性的，在矩阵中每 8 * 32 * sizeof(int) = 128 个元素重复一次。

使用正确的索引修改共享内存。

TMA在从全局内存传输数据到共享内存时会自动为我们进行数据交织。那么我们如何恢复这些被交织的索引呢？虽然NVIDIA的官方文档(https://docs.nvidia.com/cuda/cuda-c-programming-guide/#tma-swizzle)中并没有详细说明TMA交织的具体实现，但在Igor Terentyev的GTC Talk(https://www.nvidia.com/en-us/on-demand/session/gtc24-s62192/)中提供了相关的计算公式。具体公式如下：

图片解析开始

这张图详细描述了在 128 字节段内对 16 字节块进行 swizzle（交织）操作的索引计算方法。

1. 约束条件（Constraints）

• NX * sizeof(T) == SWIZZLE_SIZE 其中 T array[][NX] ，即二维数组的每一行有 NX 个元素， T

请到「今天看啥」查看全文