cute（5）tiled_copy

鱿鱼圈 Lv4
  2026-06-10 22:13:32 2026-06-10 22:13:32 创建   2026-06-14 23:33:01 2026-06-14 23:33:01 更新      

前置阅读

5 Tiled Copy

对应代码：05_tiled_copy.cu 需要 GPU。

核心概念

CuTe 的数据搬运体系：

  • Copy_Atom：一条拷贝指令的最小单元
  • make_tiled_copy：把 atom 扩展到多线程协作
  • partition_S / partition_D：自动计算每个线程的读/写位置
  • get_slice：从 tiled_copy 中取出"当前线程"的那一份

整体数据流

1
global memory (src)  ──copy──>  shared memory  ──copy──>  global memory (dst)

目的：验证 tiled_copy 机制能正确搬运数据。src 和 dst 内容应完全一致。

代码解析

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
// ============================================================
// Kernel：用 tiled_copy 把 global 数据搬到 shared memory，再写回
// ============================================================
template <int kTileM, int kTileN>
__global__ void copy_kernel(float* dst, const float* src, int M, int N) {
  using namespace cute;

  // 定义 shared memory
  __shared__ float smem_buf[kTileM * kTileN];

  // Global Tensor
  auto gSrc = make_tensor(make_gmem_ptr(src), make_shape(M, N), make_stride(N, Int<1>{}));
  auto gDst = make_tensor(make_gmem_ptr(dst), make_shape(M, N), make_stride(N, Int<1>{}));

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

  // Shared Memory Tensor
  auto sA = make_tensor(make_smem_ptr(smem_buf),
                        make_layout(make_shape(Int<kTileM>{}, Int<kTileN>{}),
                                    make_stride(Int<kTileN>{}, Int<1>{})));

  // ============================================================
  // 构造 TiledCopy
  // ============================================================
  // Copy_Atom: 每条指令拷贝一个 float（4 字节）
  using CopyAtom = Copy_Atom<UniversalCopy<float>, float>;

  // TiledCopy: 用 (32, 4) = 128 个线程来覆盖 (kTileM, kTileN) 的 tile
  // 每个线程每次拷贝 (1, 1) 个元素
  auto tiled_copy = make_tiled_copy(
      CopyAtom{},
      make_layout(make_shape(Int<32>{}, Int<4>{}),      // 线程排布：32行4列
                  make_stride(Int<4>{}, Int<1>{})),      // 行优先
      make_layout(make_shape(Int<1>{}, Int<1>{})));      // 每线程每次拷贝 1x1

  auto thr_copy = tiled_copy.get_slice(threadIdx.x);

  // 切分：当前线程要从 src 读什么、往 smem 写什么
  auto tSrc = thr_copy.partition_S(src_tile);   // 当前线程在 global 中要读的部分
  auto tDst = thr_copy.partition_D(sA);          // 当前线程在 smem 中要写的部分

  // 打印第一个线程的 partition 信息（调试用）
  if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0) {
    printf("=== Thread 0 partition info ===\n");
    printf("  tSrc shape: "); print(shape(tSrc)); printf("\n");
    printf("  tDst shape: "); print(shape(tDst)); printf("\n");
    printf("  tSrc layout: "); print(tSrc.layout()); printf("\n");
    printf("  tDst layout: "); print(tDst.layout()); printf("\n");
  }

  // ============================================================
  // 执行拷贝：global -> shared
  // ============================================================
  cute::copy(tiled_copy, tSrc, tDst);
  __syncthreads();

  // ============================================================
  // 写回：shared -> global (反方向)
  // ============================================================
  auto tSrc2 = thr_copy.partition_S(sA);
  auto tDst2 = thr_copy.partition_D(dst_tile);
  cute::copy(tiled_copy, tSrc2, tDst2);
}

1. 全局 Tensor + local_tile

1
2
3
auto gSrc = make_tensor(make_gmem_ptr(src), make_shape(M, N), make_stride(N, Int<1>{}));
auto src_tile = local_tile(gSrc, make_tile(Int<kTileM>{}, Int<kTileN>{}),
                           make_coord(blockIdx.y, blockIdx.x));
  • 全局 tensor (M, N) row-major
  • local_tile 取当前 block 的 (32, 32) tile
  • blockIdx.y 对应 M 方向，blockIdx.x 对应 N 方向

2. Shared Memory Tensor

1
2
3
4
__shared__ float smem_buf[kTileM * kTileN];
auto sA = make_tensor(make_smem_ptr(smem_buf),
                      make_layout(make_shape(Int<kTileM>{}, Int<kTileN>{}),
                                  make_stride(Int<kTileN>{}, Int<1>{})));
  • make_smem_ptr 标记为 shared memory 指针（类似 make_gmem_ptr）
  • layout 是 (32, 32):(32, 1) row-major

3. 构造 TiledCopy（核心）

1
2
3
4
5
6
7
using CopyAtom = Copy_Atom<UniversalCopy<float>, float>;

auto tiled_copy = make_tiled_copy(
    CopyAtom{},                                                    // 一条指令拷 1 个 float
    make_layout(make_shape(Int<32>{}, Int<4>{}),                   // 线程排布：32 行 4 列
                make_stride(Int<4>{}, Int<1>{})),                   // row-major
    make_layout(make_shape(Int<1>{}, Int<1>{})));                   // 每线程每次拷 1×1

三个参数的含义：

参数 含义 本例
CopyAtom 一条指令能搬多少 1 个 float（4 字节）
线程 layout 128 个线程在 tile 上怎么排列 32 行 × 4 列，row-major
值 layout 每个线程每次负责多少元素 1×1

一次能覆盖多少：32×4 × 1×1 = 128 个元素。

tile 总共多少：32×32 = 1024 个元素。

每个线程要跑多少轮：1024 / 128 = 8 轮。

4. get_slice —— 取出当前线程的工作描述

1
auto thr_copy = tiled_copy.get_slice(threadIdx.x);

get_slice(threadIdx.x) 做了什么：

从 128 线程的 tiled_copy 中，根据 threadIdx.x 找到"我是第几号线程"，返回一个 thread-level 的 copy 对象。

这个对象知道当前线程在 tile 中的位置，之后可以用它来切分 Source 和 Destination。

类比：tiled_copy 是"全队的工作分配表"，get_slice 是"从表里找到我的那一行"。

5. partition_S / partition_D —— 自动计算读写位置

1
2
auto tSrc = thr_copy.partition_S(src_tile);   // 从 global 切出"我要读的"
auto tDst = thr_copy.partition_D(sA);          // 从 smem 切出"我要写的"
  • partition_S (Source)：把 src_tile 按线程分配，返回当前线程要读的部分
  • partition_D (Destination)：把 sA 按线程分配，返回当前线程要写的部分

结果 shape：(CPY, CPY_M, CPY_N) = (1, 1, 8)

1
2
3
CPY   = 1  ← 每次拷 1 个元素
CPY_M = 1  ← M 方向：32 行 / 32 行线程 = 1 轮
CPY_N = 8  ← N 方向：32 列 / 4 列线程 = 8 轮

6. 每个线程的具体工作

线程 layout (32, 4):(4, 1) row-major，线程号到 tile 位置的映射：

1
2
3
4
5
6
7
8
9
threadIdx.x = row * 4 + col

thread 0   → (row=0,  col=0)
thread 1   → (row=0,  col=1)
thread 2   → (row=0,  col=2)
thread 3   → (row=0,  col=3)
thread 4   → (row=1,  col=0)
...
thread 127 → (row=31, col=3)

tile 是 32×32，线程一次覆盖 32×4，N 方向要跑 8 轮：

1
2
3
4
5
6
7
8
9
10
32×32 tile 的分配：

           col 0-3      col 4-7     col 8-11    ...   col 28-31
          +-----------+-----------+-----------+-----+-----------+
 row 0    | t0 t1 t2 t3 | t0 t1 t2 t3 | t0 t1 t2 t3 | ... | t0 t1 t2 t3 |
 row 1    | t4 t5 t6 t7 | t4 t5 t6 t7 | t4 t5 t6 t7 | ... | t4 t5 t6 t7 |
 ...      |             |             |             |     |             |
 row 31   | t124-t127   | t124-t127   | t124-t127   | ... | t124-t127   |
          +-----------+-----------+-----------+-----+-----------+
            第 0 轮       第 1 轮       第 2 轮            第 7 轮

Thread 0（row=0, col=0）负责的 8 个元素：

1
2
3
4
5
6
7
8
第 0 轮：tile(0,  0)
第 1 轮：tile(0,  4)
第 2 轮：tile(0,  8)
第 3 轮：tile(0, 12)
第 4 轮：tile(0, 16)
第 5 轮：tile(0, 20)
第 6 轮：tile(0, 24)
第 7 轮：tile(0, 28)

Thread 5（row=1, col=1）负责的 8 个元素：

1
2
3
4
第 0 轮：tile(1,  1)
第 1 轮：tile(1,  5)
...
第 7 轮：tile(1, 29)

7. 执行拷贝

1
2
3
4
5
6
7
8
// global → shared
cute::copy(tiled_copy, tSrc, tDst);
__syncthreads();

// shared → global
auto tSrc2 = thr_copy.partition_S(sA);
auto tDst2 = thr_copy.partition_D(dst_tile);
cute::copy(tiled_copy, tSrc2, tDst2);

cute::copy 对当前线程遍历所有 8 个元素逐个拷贝，等价于：

1
2
3
for (int i = 0; i < 8; ++i) {
    tDst(i) = tSrc(i);
}

__syncthreads() 确保所有线程都写完 smem 后再读。

反方向拷贝时，Source 变成 smem，Destination 变成 global。

make_tiled_copy 参数详解

1
make_tiled_copy(CopyAtom, ThreadLayout, ValLayout)

一次覆盖的 tile 大小 = ThreadLayout的shape × ValLayout的shape

1
2
3
ThreadLayout = (32, 4)
ValLayout    = (1, 1)
一次覆盖 = (32×1, 4×1) = (32, 4) = 128 个元素

如果把 ValLayout 改成 (1, 4)：

  • 每个线程每次拷 4 个 float（128bit 向量化）
  • 一次覆盖 (32×1, 4×4) = (32, 16) = 512 个元素
  • 只需 2 轮就能覆盖 32×32 tile

高性能 kernel 中常见配置：

1
2
3
4
// 128bit 异步拷贝
Copy_Atom<CP_ASYNC_CACHEGLOBAL<uint128_t>, half_t>  // 一次拷 128bit = 8 个 half
make_layout(make_shape(Int<32>{}, Int<4>{}))         // 128 线程
make_layout(make_shape(Int<1>{}, Int<8>{}))          // 每线程 8 个元素 = 128bit

API 总结

API 作用
Copy_Atom<Op, T> 定义一条拷贝指令的最小单元
make_tiled_copy(atom, thr_layout, val_layout) 创建多线程协作的 tiled copy
tiled_copy.get_slice(threadIdx.x) 取出当前线程的工作描述
thr_copy.partition_S(tensor) 切分 Source：当前线程要读哪些
thr_copy.partition_D(tensor) 切分 Destination：当前线程要写哪些
cute::copy(tiled_copy, src, dst) 执行拷贝
make_smem_ptr(ptr) 标记为 shared memory 指针

与 04 的对比

04（手动） 05（TiledCopy）
线程→元素映射 手动算 index get_slice + partition 自动
拷贝操作 逐元素赋值 cute::copy
向量化 自己实现 换 CopyAtom 即可（如 uint128_t）
适用场景 简单 kernel 复杂数据搬运（G2S, S2R, R2S, S2G）

从 05 开始，CuTe 自动管理线程分配和数据搬运，你只需描述"怎么排线程"和"每次搬多少"，剩下的全由 partition 自动计算。

  • 标题: cute（5）tiled_copy
  • 作者: 鱿鱼圈
  • 创建于 : 2026-06-10 22:13:32
  • 更新于 : 2026-06-14 23:33:01
  • 链接: https://yuyanqi.com/2026/06/10/cute(5)tiled_copy/
  • 版权声明: 本文章采用 CC BY-NC-SA 4.0 进行许可。
评论
目录
cute（5）tiled_copy
  1. 前置阅读
  2. 5 Tiled Copy
    1. 核心概念
    2. 整体数据流
    3. 代码解析
      1. 1. 全局 Tensor + local_tile
      2. 2. Shared Memory Tensor
      3. 3. 构造 TiledCopy（核心）
      4. 4. get_slice —— 取出当前线程的工作描述
      5. 5. partition_S / partition_D —— 自动计算读写位置
      6. 6. 每个线程的具体工作
      7. 7. 执行拷贝
    4. make_tiled_copy 参数详解
    5. API 总结
    6. 与 04 的对比