📅 创建时间：2026-06-03 🏷️ 标签：#TVM #TensorExpression #Schedule #TE #Axis #Tiling #Vectorization #ComputeSchedule分离 📚 前置知识：[[/04-ai/01-llm-engineering/07-llm-evolution]]（LLM 发展脉络）[[17-kernel-primer]]（Kernel 开发入门） 📚 相关知识：[[21-auto-scheduling]]（自动调度）[[13-codegen-architecture]]（代码生成架构）[[20-triton]]（Triton）

TVM Tensor Expression 与计算调度分离 / TVM Tensor Expressions and Compute-Schedule Separation

┌──────────────────────────────────────────────────────────────────────────────┐ │ 📖 场景：TVM TE 写的 GEMM 比手写 CUDA 慢 2 倍 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ 你用 TVM TE 写了一个 GEMM，但运行出来比手写 CUDA 慢 2 倍。 │ │ 你知道 schedule 可以优化，但 split/reorder/vectorize/fuse 各种参数组合爆炸， │ │ 不知道哪个是对的。 │ └──────────────────────────────────────────────────────────────────────────────┘

第1节 Compute-Schedule 分离——描述"算什么"和"怎么算"分开

1.1 传统编程模型 vs TVM TE 编程模型

传统 CUDA 编程： "算什么"和"怎么算"混在一起

cuda

// 问题：计算逻辑和调度逻辑无法分离
// 如果你想尝试不同的 tiling，需要重写整个 kernel

__global__ void matmul_kernel(float* C, float* A, float* B) {
    // 硬编码的 tiling: block(16, 16), thread(8, 8)
    __shared__ float As[16][16];
    __shared__ float Bs[16][16];
    
    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;
    
    // 固定的三层循环
    for (int kk = 0; kk < K; kk += 16) {
        // 固定 Cooperative loading
        As[ty][tx] = A[(by*16+ty)*K + kk + tx];
        Bs[ty][tx] = B[(kk+ty)*N + bx*16 + tx];
        __syncthreads();
        
        for (int k = 0; k < 16; k++)
            C[(by*16+ty)*N + bx*16+tx] += As[ty][k] * Bs[k][tx];
        __syncthreads();
    }
}
// 如果想试试 tile(32,32)，需要重写
// 如果想试试不同的 loop order，需要重写

TVM TE 编程模型： 分离"算什么"和"怎么算"

┌─────────────────────────────────────────────────────────────────────┐
│  Tensor Expression (TE)                                            │
│                                                                     │
│   ┌─────────────────────┐        ┌─────────────────────┐          │
│   │  Compute (算什么呢)  │        │   Schedule (怎么算呢) │          │
│   │                     │        │                     │          │
│   │  C[i,j] = sum_k     │   →   │  split(i, 4)       │          │
│   │    A[i,k] * B[k,j]  │        │  reorder(j, i_outer,│          │
│   │                     │        │         i_inner)     │          │
│   │  (声明式，只描述计算) │        │  bind(threadIdx.x,  │          │
│   │                     │        │         i_outer)    │          │
│   └─────────────────────┘        └─────────────────────┘          │
│                                                                     │
│   优势：同一个 Compute 可以用不同的 Schedule 尝试                    │
│         不需要重写计算逻辑                                           │
└─────────────────────────────────────────────────────────────────────┘

1.2 为什么分离很重要

分离的价值:

1. 可组合优化
   - 先调 tile，再调 unroll，再调 vectorize
   - 每个优化可以独立验证

2. 自动搜索成为可能
   - 同一个 compute 定义，可以尝试 1000 种 schedule
   - AutoTVM / Ansor 就是建立在这个基础上

3. 硬件抽象
   - Compute 是通用的（任何硬件）
   - Schedule 是硬件相关的（CUDA/OpenCL/Vulkan）

1.3 TVM TE 的两个核心 API

python

# TVM TE 两个核心 API

# 1. te.compute() — 描述"算什么"
#    输入: output shape + compute function
#    输出: Tensor expression

# 2. te.create_schedule() + schedule primitives
#    输入: Tensor expression
#    输出: Schedule (描述"怎么算")

第2节 Tensor Expression 基础

2.1 placeholder——声明输入张量

python

import tvm
from tvm import te

# placeholder: 声明一个输入张量
# 类似于在 numpy 中: A = numpy.zeros((n, k))
# tvm 会根据 shape 和 dtype 分配内存

n = te.var("n")    # 符号变量，表示任意维度
m = te.var("m")
k = te.var("k")

# 声明矩阵 A: shape (n, k), dtype "float32"
A = te.placeholder((n, k), name="A", dtype="float32")
#  A: 一个 n×k 的输入张量，TVM 不知道具体值，只知道 shape

# 声明矩阵 B: shape (k, m)
B = te.placeholder((k, m), name="B", dtype="float32")

print(A)  # Tensor(shape=[n, k], name=A)
print(B)  # Tensor(shape=[k, m], name=B)

2.2 compute——声明计算

python

# compute: 声明一个计算操作
# te.compute(output_shape, lambda_function, name="...")

# 声明 reduce_axis: 遍历 k 维度进行求和
# reduce_axis 是 TVM 中的规约轴，类似于 for 循环遍历
k_axis = te.reduce_axis((0, k), name="k")
#  k_axis: 遍历 0..k-1

# 声明 GEMM 计算: C[i,j] = sum_k(A[i,k] * B[k,j])
# lambda i, j: te.sum(...) 中 i,j 是输出坐标
# te.sum 的参数是求和表达式
C = te.compute(
    (n, m),                    # 输出 shape: n×m
    lambda i, j: te.sum(       # compute function
        A[i, k_axis] * B[k_axis, j],  # A[i,k] * B[k,j]
        axis=k_axis            # 对 k_axis 求和
    ),
    name="C",
    dtype="float32"
)

# C 的语义:
# for i in range(n):
#     for j in range(m):
#         C[i,j] = sum_k(A[i,k] * B[k,j])

# 验证输出张量的结构
print(C.op)        # AttrStmt (attr_key=gpu, ...)
print(C.op.body)   # [sum(k_axis, [(...)] * (...))]

2.3 te.compute 的高级用法

python

# 场景 1: 带 bias 的 GEMM
C = te.compute(
    (n, m),
    lambda i, j: te.sum(A[i, k_axis] * B[k_axis, j], axis=k_axis) + bias[j],
    name="C_with_bias"
)

# 场景 2: 带激活函数的 GEMM
C = te.compute(
    (n, m),
    lambda i, j: te.relu(te.sum(A[i, k_axis] * B[k_axis, j], axis=k_axis)),
    name="C_with_relu"
)

# 场景 3: 多输出 (如 SqueezeExcitation)
# 定义一个产生多个输出的 compute
def conv2d_relu_conv2d(inputs, weight1, weight2):
    # 中间结果
    conv1 = te.compute(
        (batch, channels_out1, H, W),
        lambda b, c, h, w: te.sum(
            inputs[b, c_in, h + kh, w + kw] * weight1[c_in, c, kh, kw],
            axis=[c_in, kh, kw]
        ),
        name="conv1"
    )
    
    relu1 = te.compute(
        conv1.shape,
        lambda *args: te.max(conv1[args], 0),  # ReLU
        name="relu1"
    )
    
    # 输出
    out = te.compute(
        (batch, channels_out2, H, W),
        lambda b, c, h, w: te.sum(
            relu1[b, c_in, h + kh, w + kw] * weight2[c_in, c, kh, kw],
            axis=[c_in, kh, kw]
        ),
        name="out"
    )
    
    return relu1, out  # 多输出

第3节 Schedule 原语

3.1 Schedule 基础——create_schedule

python

# create_schedule: 从 compute graph 创建 schedule
# schedule 描述了"怎么算"

# 获取 C 的 op (operation)
s = te.create_schedule(C.op)
# s: Schedule 对象，包含 compute graph 的执行顺序

print(s[C.op])  # 打印当前的 schedule
# Stage(C, scope=root):
#   for i: i.pos = 0, i.extent = n
#     for j: j.pos = 0, j.extent = m
#       for k: k.pos = 0, k.extent = k
#         produce C
#           C[i, j] = sum k

3.2 split——拆分 axis

python

# split: 将一个 axis 拆分为 inner 和 outer 两层
# 语法: split(axis, factor) 或 split(axis, [outer_factor, inner_factor])

s = te.create_schedule(C.op)

# 获取当前 schedule 中的 axis
print(s[C.op].op.axis)      # [i, j] — 输出 axis
print(s[C.op].op.reduce_axis)  # [k_axis] — reduce axis

i, j = s[C.op].op.axis
k = s[C.op].op.reduce_axis[0]

# 拆分 i axis: i → i_outer, i_inner
# 假设 n = 1024, factor = 64
# i_outer = 1024 / 64 = 16
# i_inner = 64
i_outer, i_inner = s[C.op].split(i, factor=64)

# 或者用 nparts (按数量拆分)
i_outer2, i_inner2 = s[C.op].split(i, nparts=16)
# i_outer2 有 16 个迭代
# i_inner2 有 1024/16 = 64 个迭代

# 拆分 reduce axis k
ko, ki = s[C.op].split(k, factor=16)

# 验证拆分
print(i_outer.extent)  # 16 (n/64)
print(i_inner.extent)  # 64

3.3 reorder——改变 loop 嵌套顺序

python

# reorder: 改变 loop 嵌套顺序
# 不同的嵌套顺序会导致完全不同的内存访问模式

# 当前默认顺序: i → j → k
# 重排为: i_outer → j → i_inner → k
s = te.create_schedule(C.op)

i, j = s[C.op].op.axis
k = s[C.op].op.reduce_axis[0]

# 先 split i
i_outer, i_inner = s[C.op].split(i, factor=64)

# 重排 loop 顺序
# 原来: for i: for j: for k:
# 重排后: for i_outer: for j: for i_inner: for k:
s[C.op].reorder(
    i_outer,  # 第一层: i_outer
    j,         # 第二层: j
    i_inner,   # 第三层: i_inner
    k          # 第四层: k
)

# 为什么这样重排?
# - i_outer 在最外层: blockIdx.x
# - j 在第二层: blockIdx.y
# - i_inner 在第三层: threadIdx.x
# - k 在最内层: 连续内存访问

# 对应的 CUDA grid/block 映射:
# s[C.op].bind(i_outer, te.thread_axis("blockIdx.x"))
# s[C.op].bind(j, te.thread_axis("blockIdx.y"))
# s[C.op].bind(i_inner, te.thread_axis("threadIdx.x"))

3.4 tile——二维分块

python

# tile: 简化版的 split + reorder，用于二维分块
# tile(bx, by, tx, ty, [bm, bn], [tm, tn])
# 等价于: split + split + reorder

s = te.create_schedule(C.op)

i, j = s[C.op].op.axis
k = s[C.op].op.reduce_axis[0]

# tile i 和 j: 2D block × 2D thread
# 把 (i, j) 拆分为 (i_outer, j_outer) × (i_inner, j_inner)
# 效果: 把矩阵分成 block_size × block_size 的块
#       每个 block 内部再分成 thread_size × thread_size 的小块

bx, ty = s[C.op].split(i, factor=64)   # i → bx(i_outer), ty(i_inner)
by, tx = s[C.op].split(j, factor=64)   # j → by(j_outer), tx(j_inner)

# 再次拆分 i_inner 和 j_inner (如果需要)
# 或者直接 reorder
s[C.op].reorder(bx, by, ty, tx, k)

# tile 等价的简化写法:
# s[C.op].tile(i, j, x_factor=64, y_factor=64)
# 但更推荐用 split + reorder 以获得更多控制

3.5 fuse——合并相邻 axis

python

# fuse: 合并两个相邻的 axis
# 合并后变成一个一维 index

s = te.create_schedule(C.op)

i, j = s[C.op].op.axis

# 假设 i=1024, j=512
# fuse(i, j) → fused 范围: 1024 * 512 = 524288

fused = s[C.op].fuse(i, j)

# fused 的 extent = i.extent * j.extent = 1024 * 512
print(fused.extent)  # 524288

# fused index 可以用除法和取模恢复原来的 i, j
# i = fused // j.extent
# j = fused % j.extent

# 应用: 把 2D thread grid 映射到 1D
# blockIdx.x * blockDim.x + threadIdx.x → 单一的 thread index

3.6 vectorize——向量化

python

# vectorize: 标记可以被 SIMD 指令向量化处理的 axis
# TVM 会生成 vectorized load/store 指令

s = te.create_schedule(C.op)

i, j = s[C.op].op.axis
k = s[C.op].op.reduce_axis[0]

# 假设 j 方向有连续的内存访问
# 使用 vectorize(8) 让 TVM 生成 8-wide 的 vectorized load
# 原来: load A[i, j]
# 向量化后: vectorized_load A[i, j:j+8]

# vectorize 只能用于连续访问的 axis
# 所以一般用于最内层 axis

# 示例: 如果 j 是连续的，可以:
s[C.op].vectorize(j)  # TVM 会尝试 vectorize j

# 注意: vectorize 需要满足 alignment 要求
# TVM 会自动检查，或者用 pad_einsum 填充到对齐

3.7 unroll——循环展开

python

# unroll: 展开循环，减少分支开销，提高 ILP (指令级并行)

s = te.create_schedule(C.op)

i, j = s[C.op].op.axis
k = s[C.op].op.reduce_axis[0]

# unroll k axis (reduce axis)
# 如果 k < 某个阈值 (如 8)，TVM 会完全展开
s[C.op].unroll(k)

# unroll 的三种模式:
# - auto: TVM 自动决定
# - unroll_extent: 展开到这个值
# - full_unroll: 完全展开

# 示例: 展开小于 8 的循环
# 相当于:
# if k.extent < 8:
#     for k in range(k.extent):
#         C[i,j] += A[i,k] * B[k,j]  # 展开成多条语句
# else:
#     for k in range(k.extent):  # 保持循环
#         C[i,j] += A[i,k] * B[k,j]

# 典型用法: unroll 小 tile 减少开销
# tile_size < 8 → unroll
# tile_size > 32 → 保持循环（展开太大反而增加寄存器压力）

3.8 bind——绑定到 GPU thread/block

python

# bind: 将 axis 绑定到 GPU thread 或 block

s = te.create_schedule(C.op)

i, j = s[C.op].op.axis
k = s[C.op].op.reduce_axis[0]

# split i 和 j 以便并行
i_outer, i_inner = s[C.op].split(i, factor=64)
j_outer, j_inner = s[C.op].split(j, factor=64)

# bind 到 GPU 维度
s[C.op].bind(i_outer, te.thread_axis("blockIdx.x"))  # i_outer → block.x
s[C.op].bind(j_outer, te.thread_axis("blockIdx.y"))  # j_outer → block.y
s[C.op].bind(i_inner, te.thread_axis("threadIdx.x"))  # i_inner → thread.x
s[C.op].bind(j_inner, te.thread_axis("threadIdx.y"))  # j_inner → thread.y

# thread_axis 常用的选项:
# - "blockIdx.x/y/z": CUDA block index
# - "threadIdx.x/y/z": CUDA thread index
# - "vthread.x/y/z": virtual thread (用于更细粒度调度)
# - "ctr": 标量迭代器（用于常量循环）

第4节 Compute at 策略——inline vs. compute_at

4.1 compute_at 的概念

python

# compute_at: 把中间结果放到"哪里"计算
# 控制中间结果的生存周期和存储位置

s = te.create_schedule(C.op)

# 默认: 所有 tensor 在 root level 计算
# C 的全部结果计算完后再写回 global memory

# compute_at(target, axis):
# 把当前 tensor 的计算"移到" target tensor 的 loop 内
# 可以放在 shared memory 或 register 中，而不是 global memory

# 场景: 如果有一个中间 tensor T
# T = te.compute((n, m), lambda i, j: A[i, j] * 2)
# C = te.compute((n, m), lambda i, j: T[i, j] + B[i, j])

# 如果不 compute_at: T 会被计算并存储到 global memory，然后 C 再读取
# 如果 compute_at: T 在 C 的 loop 内计算（inline），不占用 global memory

# 示例: 把 GEMM 的 K 方向循环展开后，把 A/B tile compute_at 到合适位置
# ...

# 关键: compute_at 只能用于有共同 parent loop 的情况

第5节 Cache Read/Write——手动管理缓存

5.1 cache_read / cache_write

python

# cache_read: 把 global memory 的读取缓存到 shared memory
# cache_write: 把结果先写到 shared memory，再写回 global memory

s = te.create_schedule(C.op)

# 获取 stage
s_A = s.cache_read(A, "shared")
s_B = s.cache_read(B, "shared")
s_C = s.cache_write(C, "shared")

# 然后对 s_A, s_B 应用 schedule
# 把 A 的读取放到 C 的 loop 内，并绑定到 shared memory

# 示例 GEMM with shared memory:
# 1. 把 A, B 读取到 shared memory
# 2. 使用 cache_read 后，A, B 会在 C 的 loop 内被读取

# 更详细的 shared memory 使用:
# see Section 6 中的完整 GEMM 示例

第6节完整的 TE GEMM 示例

6.1 带 Shared Memory 的 GEMM Schedule

python

import tvm
from tvm import te
import numpy as np

# ============================================================
# 第一步：定义 Compute
# ============================================================

N = 1024  # 矩阵维度
K = 1024

# placeholder
A = te.placeholder((N, K), name="A", dtype="float32")
B = te.placeholder((K, N), name="B", dtype="float32")

# reduce axis
k_axis = te.reduce_axis((0, K), name="k")

# GEMM: C = A @ B
C = te.compute(
    (N, N),
    lambda i, j: te.sum(A[i, k_axis] * B[k_axis, j], axis=k_axis),
    name="C"
)

# ============================================================
# 第二步：定义 Schedule
# ============================================================

s = te.create_schedule(C.op)

# Tile 的大小配置
BLOCK_SIZE = 32   # 每个 block 处理 32×32 的 C
THREAD_SIZE = 8   # 每个 thread 处理 8×8 的 C

# 获取所有 axis
i, j = s[C.op].op.axis
k = s[C.op].op.reduce_axis[0]

# 第一层 tiling: 把 i,j 分成 block 级别
i0, i1 = s[C.op].split(i, factor=BLOCK_SIZE)
j0, j1 = s[C.op].split(j, factor=BLOCK_SIZE)

# 第二层 tiling: block 内部分成 thread 级别
i0_outer, i0_inner = s[C.op].split(i0, factor=THREAD_SIZE)
j0_outer, j0_inner = s[C.op].split(j0, factor=THREAD_SIZE)

# 把 k 分成 tile，减少 reduce 范围
ko, ki = s[C.op].split(k, factor=BLOCK_SIZE)

# 重排 loop 顺序: 按 (i1, j1, i0_outer, j0_outer, i0_inner, j0_inner, ki) 执行
# i1, j1: block index
# i0_outer, j0_outer: thread block 内的 tile index
# i0_inner, j0_inner: thread index
# ki: reduce axis
s[C.op].reorder(
    i1,          # blockIdx.x
    j1,          # blockIdx.y
    i0_outer,    # 0..BLOCK_SIZE/THREAD_SIZE
    j0_outer,    # 0..BLOCK_SIZE/THREAD_SIZE
    i0_inner,    # threadIdx.x
    j0_inner,    # threadIdx.y
    ko,          # k outer (reduce)
    ki           # k inner (reduce)
)

# 绑定到 GPU thread/block
s[C.op].bind(i1, te.thread_axis("blockIdx.x"))
s[C.op].bind(j1, te.thread_axis("blockIdx.y"))
s[C.op].bind(i0_inner, te.thread_axis("threadIdx.x"))
s[C.op].bind(j0_inner, te.thread_axis("threadIdx.y"))

# ============================================================
# 第三步：缓存读取到 Shared Memory
# ============================================================

# 把 A 和 B 的读取缓存到 shared memory
# 这样同一 block 内的线程可以共享数据
s_A = s.cache_read(A, "shared", [C.op])
s_B = s.cache_read(B, "shared", [C.op])

# 对缓存的读取应用 schedule
# 把 A 的读取放到 C 的 loop 中
# s_A 的 axis: A_i0_outer, A_i0_inner, k
# 需要 compute_at 到 C 的合适位置

# s_A 的 axis
A_i, A_k = s[s_A].op.axis
A_ko, A_ki = s[s_A].split(A_k, factor=BLOCK_SIZE)

# s_A 需要 compute_at 到 C 的 ko level
# 这样 A 在每个 ko 迭代时加载一次，供所有 ki 迭代复用
s[s_A].compute_at(s[C.op], ko)

# 拆分 s_A 的 axis 以匹配 C 的结构
s[s_A].split(A_k, factor=BLOCK_SIZE)
s[s_A].split(A_i, factor=THREAD_SIZE)

# 绑定 A 的读取到 shared memory
# (thread 级别需要 cooperative load)
s[s_A].bind(A_i, te.thread_axis("threadIdx.x"))
s[s_A].bind(A_k, te.thread_axis("threadIdx.y"))

# 同样的操作给 B
s_B = s.cache_read(B, "shared", [C.op])
s[s_B].compute_at(s[C.op], ko)
B_j, B_k = s[s_B].op.axis
s[s_B].split(B_k, factor=BLOCK_SIZE)
s[s_B].split(B_j, factor=THREAD_SIZE)
s[s_B].bind(B_j, te.thread_axis("threadIdx.x"))
s[s_B].bind(B_k, te.thread_axis("threadIdx.y"))

# ============================================================
# 第四步：其他优化
# ============================================================

# 展开 inner reduce axis
s[C.op].unroll(ki)

# 或者 vectorize inner axis (如果有的话)
# s[C.op].vectorize(...)

# ============================================================
# 第五步：Build 和验证
# ============================================================

# Lower schedule 生成低级 IR
print(tvm.lower(s, [A, B, C], simple_mode=True))

# Build to CUDA
# dev = tvm.cuda(0)
# f = tvm.build(s, [A, B, C], "cuda", target_host="llvm")
# ctx = dev
# a = tvm.nd.array(np.random.rand(N, K).astype("float32"), ctx)
# b = tvm.nd.array(np.random.rand(K, N).astype("float32"), ctx)
# c = tvm.nd.array(np.zeros((N, N)).astype("float32"), ctx)
# f(a, b, c)
# print(c.numpy()[:4, :4])

6.2 Schedule 原语效果对照表

原语	作用	典型参数	效果
`split(factor)`	拆分 axis	factor=32-128	增加并行度，允许 tiling
`split(nparts)`	按数量拆分	nparts=grid_size	控制并行度
`reorder`	重排 loop 嵌套	-	改变内存访问模式，优化 locality
`tile`	二维分块	x_factor=32, y_factor=32	等价于 split×2 + reorder
`fuse`	合并 axis	-	简化索引，减少维度
`bind`	绑定到 thread/block	"threadIdx.x", "blockIdx.x"	GPU 并行化
`vectorize`	向量化	axis	生成 SIMD 指令
`unroll`	循环展开	-	减少分支，提高 ILP
`compute_at`	移动计算位置	target, axis	inline 优化，减少内存访问
`cache_read`	缓存读取	"shared"	把 global memory 读取缓存到 smem
`cache_write`	缓存写入	"shared"	把结果先写到 smem 再写回 gmem

第7节 TE vs TOPI——高层 Schedule 库

7.1 TOPI 是什么

python

# TOPI = TVM Operator Inventory
# TOPI 提供了常用算子的预定义 schedule

from topi import nn
from topi.util import get_const_int

# TOPI 提供了简单的 GEMM 实现
# topi.matmul 已经做了合理的默认 schedule

C = topi.matmul(A, B)  # 自动 schedule

# TOPI 的优势:
# 1. 开箱即用，不需要手写 schedule
# 2. 跨平台 (CUDA/OpenCL/AArch64)
# 3. 包含常见算子: conv2d, pool, softmax, batch_matmul

# TOPI 的劣势:
# 1. 不够灵活，无法定制复杂 schedule
# 2. 对新硬件架构支持滞后

第8节调试 Schedule 错误

8.1 常见错误和解决方案

错误现象	原因	解决方案
segmentation fault	非法的 schedule（如 bind 到不存在的 axis）	检查 bind 的 target
性能反而下降	错误的 reorder 破坏了 locality	检查内存访问模式
编译失败	axis 拆分参数不合理	检查 factor 是否整除 extent
结果错误	reduce axis 顺序错误	确保 reduce axis 在最内层
内存占用过大	shared memory 申请太大	减少 tile size 或 block size

8.2 调试技巧

python

# 技巧 1: 打印 lower 后的 IR，检查 schedule 是否生效
print(tvm.lower(s, [A, B, C], simple_mode=True))

# 技巧 2: 使用 verbose 输出查看 schedule 应用
with tvm.transform.PassContext(opt_level=3):
    lib = tvm.build(s, [A, B, C], "cuda")

# 技巧 3: 使用 Nsight 验证生成的代码
# tvm 会生成临时 .cu 文件，路径在编译输出中

# 技巧 4: 逐步应用 schedule，每次验证正确性
s = te.create_schedule(C.op)

# 第一步: 只 split i
i_outer, i_inner = s[C.op].split(i, factor=64)

# 验证
print(tvm.lower(s, [A, B, C], simple_mode=True))

# 第二步: 再加 reorder
s[C.op].reorder(i_outer, j, i_inner, k)

# 验证
print(tvm.lower(s, [A, B, C], simple_mode=True))

升华

┌────────────────────────────────────────────────────────────────────────────┐
│                     🚀 TVM Tensor Expression 核心原则                      │
├────────────────────────────────────────────────────────────────────────────┤
│                                                                            │
│  ① Compute-Schedule 分离是核心：先写正确的 compute，再用 schedule 优化       │
│     不要在 compute 阶段就想着优化                                          │
│                                                                            │
│  ② split 是最重要的原语：几乎所有优化都从 split 开始                        │
│     split 把一维 axis 变成二维（block/thread 或 tile 内/外）                 │
│                                                                            │
│  ③ reorder 决定内存访问模式：循环顺序决定数据局部性                         │
│     关键是让内层循环访问连续内存                                           │
│                                                                            │
│  ④ 遵循 Roofline 指导优化方向：Memory-Bound 用 tiling + cache_read           │
│     Compute-Bound 用 unroll + vectorize                                   │
│                                                                            │
│  ⑤ 测量验证：TVM lower 后打印 IR，用 Nsight 验证生成的代码                   │
│     猜测不能替代测量                                                       │
│                                                                            │
└────────────────────────────────────────────────────────────────────────────┘

"AI 可查 vs 必须理解"清单

必须理解（不理解就等于不会）：

🔴 Compute-Schedule 分离的意义：为什么分离是 TVM 的核心价值
🔴 split 的物理意义：split 把一维循环变成两级，这两级分别映射到什么
🔴 reorder 对内存访问的影响：最内层循环必须是连续内存访问
🔴 reduce_axis 的特殊性：reduce axis 必须放在最内层还是外层？为什么
🔴 bind 的作用：threadIdx.x/y 和 blockIdx.x/y 分别控制什么

AI 可查（知道去哪查就行）：

✅ 具体算子的 TOPI 模板代码（TOPI 源码）
✅ te.create_schedule 的完整 API（TVM 文档）
✅ 特定 GPU 的 shared memory 大小限制（NVIDIA 硬件文档）
✅ TVM Pass 的具体实现（TVM 源码或文档）
✅ 特定 schedule 配置的性能 benchmark 数据（社区 benchmark）

学习状态：🟡 开始学习

TVM Tensor Expression 与计算调度分离 / TVM Tensor Expressions and Compute-Schedule Separation ​

第1节 Compute-Schedule 分离——描述"算什么"和"怎么算"分开 ​

1.1 传统编程模型 vs TVM TE 编程模型 ​

1.2 为什么分离很重要 ​

1.3 TVM TE 的两个核心 API ​

第2节 Tensor Expression 基础 ​

2.1 placeholder——声明输入张量 ​

2.2 compute——声明计算 ​

2.3 te.compute 的高级用法 ​

第3节 Schedule 原语 ​

3.1 Schedule 基础——create_schedule ​

3.2 split——拆分 axis ​

3.3 reorder——改变 loop 嵌套顺序 ​

3.4 tile——二维分块 ​

3.5 fuse——合并相邻 axis ​

3.6 vectorize——向量化 ​

3.7 unroll——循环展开 ​

3.8 bind——绑定到 GPU thread/block ​

第4节 Compute at 策略——inline vs. compute_at ​

4.1 compute_at 的概念 ​

第5节 Cache Read/Write——手动管理缓存 ​

5.1 cache_read / cache_write ​

第6节 完整的 TE GEMM 示例 ​

6.1 带 Shared Memory 的 GEMM Schedule ​

6.2 Schedule 原语效果对照表 ​

第7节 TE vs TOPI——高层 Schedule 库 ​

7.1 TOPI 是什么 ​

第8节 调试 Schedule 错误 ​

8.1 常见错误和解决方案 ​

8.2 调试技巧 ​

升华 ​

"AI 可查 vs 必须理解"清单 ​

TVM Tensor Expression 与计算调度分离 / TVM Tensor Expressions and Compute-Schedule Separation

第1节 Compute-Schedule 分离——描述"算什么"和"怎么算"分开

1.1 传统编程模型 vs TVM TE 编程模型

1.2 为什么分离很重要

1.3 TVM TE 的两个核心 API

第2节 Tensor Expression 基础

2.1 placeholder——声明输入张量

2.2 compute——声明计算

2.3 te.compute 的高级用法

第3节 Schedule 原语

3.1 Schedule 基础——create_schedule

3.2 split——拆分 axis

3.3 reorder——改变 loop 嵌套顺序

3.4 tile——二维分块

3.5 fuse——合并相邻 axis

3.6 vectorize——向量化

3.7 unroll——循环展开

3.8 bind——绑定到 GPU thread/block

第4节 Compute at 策略——inline vs. compute_at

4.1 compute_at 的概念

第5节 Cache Read/Write——手动管理缓存

5.1 cache_read / cache_write

第6节完整的 TE GEMM 示例

6.1 带 Shared Memory 的 GEMM Schedule

6.2 Schedule 原语效果对照表

第7节 TE vs TOPI——高层 Schedule 库

7.1 TOPI 是什么

第8节调试 Schedule 错误

8.1 常见错误和解决方案

8.2 调试技巧

升华

"AI 可查 vs 必须理解"清单