cute（4）第一个Kernel

第 4 课：第一个 GPU Kernel

对应代码：04_first_kernel.cu 需要 GPU。

核心概念

在 GPU kernel 中使用 CuTe Tensor：

• make_gmem_ptr 包装全局内存指针
• local_tile 按 block 切分工作
• 每个线程用一维 index 访问 tile 中的元素

---

代码解析

Kernel 1：向量加法 C = A + B

// ============================================================
// Kernel 1：向量加法 C = A + B
// 每个 block 处理 256 个元素，每个线程处理若干个
// ============================================================
template <int kBlockSize>
__global__ void vector_add_kernel(float* C, const float* A, const float* B, int N) {
  using namespace cute;

  // 把原始指针包装成 CuTe Tensor
  auto gA = make_tensor(make_gmem_ptr(A), make_shape(N));  // (N,)
  auto gB = make_tensor(make_gmem_ptr(B), make_shape(N));
  auto gC = make_tensor(make_gmem_ptr(C), make_shape(N));

  // 按 kBlockSize 切 tile，取当前 block 的那一块
  auto tA = local_tile(gA, make_tile(Int<kBlockSize>{}), make_coord(blockIdx.x));  // (kBlockSize,)
  auto tB = local_tile(gB, make_tile(Int<kBlockSize>{}), make_coord(blockIdx.x));
  auto tC = local_tile(gC, make_tile(Int<kBlockSize>{}), make_coord(blockIdx.x));

  auto thread_layout = make_layout(Int<kBlockSize>{});

  // 每个线程处理的元素
  int elements_per_thread = kBlockSize / blockDim.x;

  for (int i = 0; i < elements_per_thread; ++i) {
    int idx = threadIdx.x + i * blockDim.x;
    if (idx < size(tA)) {
      tC(idx) = tA(idx) + tB(idx);
    }
  }
}

template <int kBlockSize>
__global__ void vector_add_kernel(float* C, const float* A, const float* B, int N) {

第一步：包装全局内存指针

auto gA = make_tensor(make_gmem_ptr(A), make_shape(N));

• make_gmem_ptr(A) 把原始 CUDA 指针标记为"全局内存指针"
• CuTe 用这个标记来选择正确的内存访问指令（global load vs shared load）
• 创建了一个一维 tensor，shape = (N,)

第二步：local_tile 切分

auto tA = local_tile(gA, make_tile(Int<kBlockSize>{}), make_coord(blockIdx.x));

• 把长度 N 的数组按 kBlockSize=256 切成 N/256 个 tile
• 取第 blockIdx.x 个 tile
• 结果 tA 的 shape = (256,)，就是当前 block 负责的 256 个元素

第三步：线程分配

int elements_per_thread = kBlockSize / blockDim.x;  // 256/256 = 1

for (int i = 0; i < elements_per_thread; ++i) {
    int idx = threadIdx.x + i * blockDim.x;
    if (idx < size(tA)) {
        tC(idx) = tA(idx) + tB(idx);
    }
}

这里手动用 threadIdx.x 计算每个线程负责的 index。每个线程处理 kBlockSize/blockDim.x 个元素。

当 kBlockSize == blockDim.x 时（本例就是），每个线程只处理 1 个元素。

Launch 配置：

constexpr int kBlock = 256;
int num_blocks = N / kBlock;                           // 1024/256 = 4 个 block
vector_add_kernel<kBlock><<<num_blocks, kBlock>>>(dC, dA, dB, N);

Kernel 2：矩阵转置

// ============================================================
// Kernel 2：2D 矩阵转置
// 展示如何用 CuTe 处理二维 Tensor
// ============================================================
constexpr int kTileM = 32;
constexpr int kTileN = 32;

__global__ void transpose_kernel(float* dst, const float* src, int M, int N) {
  using namespace cute;

  // src: (M, N) row-major，dst: (N, M) row-major
  auto gSrc = make_tensor(make_gmem_ptr(src), make_shape(M, N), make_stride(N, 1));
  auto gDst = make_tensor(make_gmem_ptr(dst), make_shape(N, M), make_stride(M, 1));

  // 取当前 block 负责的 tile
  auto tileSrc = local_tile(gSrc, make_tile(Int<kTileM>{}, Int<kTileN>{}),
                            make_coord(blockIdx.y, blockIdx.x));  // (kTileM, kTileN)
  auto tileDst = local_tile(gDst, make_tile(Int<kTileN>{}, Int<kTileM>{}),
                            make_coord(blockIdx.x, blockIdx.y));  // (kTileN, kTileM)

  // 每个线程处理 tile 中的若干元素
  int tid = threadIdx.y * blockDim.x + threadIdx.x;
  int total_threads = blockDim.x * blockDim.y;
  int total_elements = kTileM * kTileN;

  for (int i = tid; i < total_elements; i += total_threads) {
    int m = i % kTileM;
    int n = i / kTileM;
    tileDst(n, m) = tileSrc(m, n);  // 转置：src(m,n) -> dst(n,m)
  }
}

__global__ void transpose_kernel(float* dst, const float* src, int M, int N) {

全局 Tensor 创建：

auto gSrc = make_tensor(make_gmem_ptr(src), make_shape(M, N), make_stride(N, 1));
auto gDst = make_tensor(make_gmem_ptr(dst), make_shape(N, M), make_stride(M, 1));

• src 是 (M, N) row-major：stride = (N, 1)
• dst 是 (N, M) row-major：stride = (M, 1)
• 转置后 dst 的行数=N，列数=M

local_tile 切分：

auto tileSrc = local_tile(gSrc, make_tile(Int<kTileM>{}, Int<kTileN>{}),
                          make_coord(blockIdx.y, blockIdx.x));  // (32, 32)
auto tileDst = local_tile(gDst, make_tile(Int<kTileN>{}, Int<kTileM>{}),
                          make_coord(blockIdx.x, blockIdx.y));  // (32, 32)

注意坐标和 tile 大小的对应关系：

• src 用 (blockIdx.y, blockIdx.x) 取 M 方向第 y 个、N 方向第 x 个 tile
• dst 用 (blockIdx.x, blockIdx.y) ← 行列互换，因为转置后 N 和 M 交换了

线程分配 + 转置：

int tid = threadIdx.y * blockDim.x + threadIdx.x;      // 线程编号
int total_threads = blockDim.x * blockDim.y;            // 32×8 = 256
int total_elements = kTileM * kTileN;                   // 32×32 = 1024

for (int i = tid; i < total_elements; i += total_threads) {
    int m = i % kTileM;
    int n = i / kTileM;
    tileDst(n, m) = tileSrc(m, n);  // 转置核心：src(m,n) → dst(n,m)
}

1024 个元素 / 256 个线程 = 每个线程处理 4 个元素。

转置的核心就一行：tileDst(n, m) = tileSrc(m, n)，交换 m 和 n 的位置。

Launch 配置：

dim3 block(32, 8);                                     // 256 线程
dim3 grid(N / kTileN, M / kTileM);                     // (128/32, 64/32) = (4, 2) 个 block
transpose_kernel<<<grid, block>>>(dDst, dSrc, M, N);

---

关键 API

API	作用
make_gmem_ptr(ptr)	标记指针为全局内存，CuTe 据此选择访问指令
make_tensor(gmem_ptr, shape, stride)	在全局内存上创建 tensor
local_tile(tensor, tile, coord)	按 tile 切分并取指定 tile
tensor(i) / tensor(i, j)	一维/二维访问元素

04 vs 05 的线程分配方式对比

	04（手动分配）	05（TiledCopy）
线程映射	手动算 threadIdx.x + i * blockDim.x	get_slice(threadIdx.x) 自动映射
partition	无，自己算 index	partition_S/D 自动切分
copy	逐元素赋值	cute::copy(tiled_copy, src, dst)
灵活性	简单场景够用	复杂场景（128bit 向量化等）必须用

04 是手动挡，05 开始进入自动挡。