在这篇博文中，我将逐步向您展示如何使用 Mojo 为Hopper架构实现高效的转置内核。最佳内核的带宽为2775.49 GB/s ，即84.1056% 。这些优化与我在本文使用的同一台H100上使用纯CUDA实现2771.35 GB/s带宽时所应用的优化相同。这表明 Mojo 可以在完全相同的任务上实现与CUDA相当的性能。您可以将这些内核与我之前编写的内核进行比较，也可以阅读我的其他博文，我在其中详细解释了这些概念。在这里，我将仅简要回顾它们，而是重点介绍实现细节。对于不了解如何在 Mojo 中使用TMA读者，我建议您参考我之前关于此主题的博文。
幼稚的方法
在调用内核之前，我们需要初始化两个TMA descriptors ，这个概念类似于我们可以在CUDA中使用的cuTensorMapEncodeTiled 。
var descriptor = create_tma_descriptor[DType.float32, 2]( gmem_dev, (GMEM_HEIGHT, GMEM_WIDTH), (GMEM_WIDTH, 1), (SMEM_HEIGHT, SMEM_WIDTH), ) var descriptor_tr = create_tma_descriptor[DType.float32, 2]( gmem_tr_dev, (GMEM_WIDTH, GMEM_HEIGHT), (GMEM_HEIGHT, 1), (SMEM_WIDTH, SMEM_HEIGHT), )
我们有两个描述符。它们都是行主格式，一个是另一个的转置。相应的smems也与转置有关。简单提醒一下，这就是我们要实现的算法。我们取一个tile，在tile内部执行转置，然后将其放在矩阵中的相反位置，即转置后的位置。
下面是存档的代码。
加载到共享内存
@__llvm_arg_metadata(descriptor, `nvvm.grid_constant`) @__llvm_arg_metadata(descriptor_tr, `nvvm.grid_constant`) fn transpose_kernel_naive[ block_size: Int ](descriptor: TMADescriptor, descriptor_tr: TMADescriptor): var shmem = stack_allocation[ block_size * block_size, DType.float32, alignment=1024, address_space = _GPUAddressSpace.SHARED, ]() var shmem_tr = stack_allocation[ block_size * block_size, DType.float32, alignment=1024, address_space = _GPUAddressSpace.SHARED, ]() var mbar = stack_allocation[ 1, Int64, address_space = _GPUAddressSpace.SHARED ]() var descriptor_ptr = UnsafePointer(to=descriptor).bitcast[NoneType]() var descriptor_tr_ptr = UnsafePointer(to=descriptor_tr).bitcast[NoneType]() x = block_idx.x * block_size y = block_idx.y * block_size col = thread_idx.x % block_size row = thread_idx.x // block_size # LOAD if thread_idx.x == 0: mbarrier_init(mbar, 1) mbarrier_arrive_expect_tx_shared(mbar, block_size * block_size * 4) cp_async_bulk_tensor_shared_cluster_global( shmem, descriptor_ptr, mbar, Index(x, y) ) barrier() mbarrier_try_wait_parity_shared(mbar, 0, 10000000)我们使用nvvm.grid_constant注释描述符，类似于我们在CUDA中所做的操作。分配共享内存后，我们使用x和y定义图块的左上坐标，并获取当前线程负责的row和column 。然后，我们将图块复制到共享内存阵列。此内核的带宽为1056.08 GB/s ，比我们使用CUDA实现的875.46 GB/s要快。我认为这是因为我们使用了PTX api 在 Mojo 中进行TMA传输。您可以在这篇精彩的博文中了解CUDA api 中它们之间的区别。
在共享内存中计算转置
# COMPUTE shmem_tr[col * block_size + row] = shmem[row * block_size + col] # FENCE barrier() tma_store_fence()我们使用这两个数组计算转置。然后创建一个fence ，让TMA知道计算已完成。
存储到 gmem
# STORE if thread_idx.x == 0: cp_async_bulk_tensor_global_shared_cta( shmem_tr, descriptor_tr_ptr, Index(y, x) ) cp_async_bulk_commit_group() cp_async_bulk_wait_group[0]()
我们使用转置的 TMA 描述符将转置结果存储到 GMEM。
调酒
有关 swizzling 的更详细说明及其工作原理，请参阅我之前关于矩阵转置的博客文章，其概念与 Mojo 相同。我在文末链接的 repo 中也提供了一个程序，您可以用它来理解 swizzling。只需调整两处即可使 swizzling 正常工作：
- 描述符需要提供适当的 swizzling 模式
- 在内核中我们需要使用混合索引
可以按如下方式实现
var descriptor = create_tma_descriptor[ DType.float32, 2, TensorMapSwizzle.SWIZZLE_128B ]( gmem_dev, (GMEM_HEIGHT, GMEM_WIDTH), (GMEM_WIDTH, 1), (SMEM_HEIGHT, SMEM_WIDTH), ) var descriptor_tr = create_tma_descriptor[ DType.float32, 2, TensorMapSwizzle.SWIZZLE_128B ]( gmem_tr_dev, (GMEM_WIDTH, GMEM_HEIGHT), (GMEM_HEIGHT, 1), (SMEM_WIDTH, SMEM_HEIGHT), )我们可以像这样计算混合索引：
fn calculate_row_swizzle[block_size: Int](col: Int, row: Int) -> Int: i16_tr = (col * BLOCK_SIZE + row) * 4 >> 4 y16_tr = i16_tr >> 3 x16_tr = i16_tr & 7 x16_swz_tr = y16_tr ^ x16_tr return ((x16_swz_tr * 4) & (BLOCK_SIZE - 1)) + (row & 3) fn calculate_col_swizzle[block_size: Int](col: Int, row: Int) -> Int: i16 = (row * BLOCK_SIZE + col) * 4 >> 4 y16 = i16 >> 3 x16 = i16 & 7 x16_swz = y16 ^ x16 return ((x16_swz * 4) & (block_size - 1)) + (col & 3)
然后像这样使用内核中的混合索引：
col_swizzle = calculate_col_swizzle[block_size](col, row) row_swizzle = calculate_row_swizzle[block_size](col, row) ... # COMPUTE shmem_tr[col * block_size + row_swizzle] = shmem[ row * block_size + col_swizzle ]
其余一切都完全相同。
该内核的速度为1437.55 GB/s而CUDA的速度为1251.76 GB/s 。
每个线程处理一批列
在内存受限的内核中，一个重要且常见的优化方法是线程粗化，本质上是在每个线程上增加更多工作量。我们可以按如下方式修改之前的内核来实现这一点：
# COMPUTE @parameter for i in range(batch_size): col_ = col + i row_ = row col_swizzle = calculate_col_swizzle[block_size](col_, row_) row_swizzle = calculate_row_swizzle[block_size](col_, row_) shmem_tr[col * block_size + row_swizzle] = shmem[ row * block_size + col_swizzle ]请注意，我们使用此方法启动较少的线程（我们除以batch_size的一个因子），以解释我们现在每个线程处理多个列的事实。
该内核的带宽为2775.49 GB/s而等效CUDA内核的带宽为2771.35 GB/s 。
结论
希望这篇博文能帮助您了解如何使用 Mojo 在 GPU 计算的常见任务上实现高性能。欢迎通过Linkedin与我联系，讨论 GPU 编程或其他与 MLSys 相关的话题。
博客文章的完整代码可以在我的Github上找到。
原文： https://veitner.bearblog.dev/highly-efficient-matrix-transpose-in-mojo/