Skip to content
Gains Summary
Main Navigation 首页 / Home
C++ 编程 / C++ Programming
系统与高性能 / Systems & Performance
Web 开发 / Web Development
人工智能 / Artificial Intelligence
工业软件 / Industrial Software
其他内容 / Other Topics
C++ 编程 / C++系统与性能 / SystemsWeb 开发 / Web人工智能 / AI工业软件 / Industrial

外观

Sidebar Navigation

← 人工智能 / Artificial Intelligence

AI 编译器 / AI Compilers

1. AI 编译器全景——为什么模型需要编译器 / The AI Compiler Landscape and Why Models Need Compilers

2. 编译原理速通——面向 ML 工程师的核心概念 / Compiler Fundamentals for Machine Learning Engineers

3. 中间表示基础——理解 IR 层级与 lowering 链路 / Intermediate Representation Levels and Lowering Pipelines

4. 计算图的构建与表示 / Building and Representing Computational Graphs

5. MLIR 架构、方言与渐进式降级 / MLIR Architecture, Dialects, and Progressive Lowering

6. 算子语义、广播、归约与形状推导 / Operator Semantics, Broadcasting, Reduction, and Shape Inference

7. 模型前端格式:ONNX、TFLite、HLO 与 SavedModel / Model Frontend Formats: ONNX, TFLite, HLO, and SavedModel

8. 图优化 Pass——经典优化在 ML 中的应用 / Graph Optimization Passes for Machine Learning

9. 算子融合——编译器最重要的性能优化 / Operator Fusion as a Core Compiler Optimization

10. 内存规划——Buffer 分配与生命周期管理 / Memory Planning, Buffer Allocation, and Lifetime Management

11. Layout 优化——数据排布转换与内存效率 / Layout Optimization for Data Movement and Memory Efficiency

12. 动态 Shape——符号分析与形状处理 / Dynamic Shapes, Symbolic Analysis, and Shape Processing

13. 硬件约束下的操作调度 / Operation Scheduling Under Hardware Constraints

14. 从模板、DSL 到 IR 降级的代码生成架构 / Code Generation Architectures from Templates and DSLs to IR Lowering

15. CPU 后端:SIMD、分块与多线程 / CPU Backends with SIMD, Tiling, and Multithreading

16. CUDA 后端:合并访存与 Tensor Core / CUDA Backends, Memory Coalescing, and Tensor Cores

17. NPU 后端:脉动阵列与端侧 AI 生态 / NPU Backends, Systolic Arrays, and Edge AI Ecosystems

18. Kernel 性能基础:Roofline 与 Occupancy / Kernel Performance Fundamentals with Roofline and Occupancy

19. CUTLASS 与分层 GEMM 模板 / CUTLASS and Hierarchical GEMM Templates

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

21. 使用 Triton 编写高性能 GPU Kernel / Triton for High-Performance GPU Kernels in Python

22. 基于成本模型与实测搜索的自动调度 / Automatic Scheduling with Cost Models and Measurement-Based Search

23. XLA 内部机制:HLO、融合与 SPMD / XLA Internals, HLO, Fusion, and SPMD

24. Torch-MLIR:从 PyTorch 算子到 MLIR 方言 / Torch-MLIR from PyTorch Operators to MLIR Dialects

25. torch.compile:Dynamo、AOTAutograd、Inductor 与 Triton / Torch Compile with Dynamo, AOTAutograd, Inductor, and Triton

26. 从 MLIR 经 LLVM 降级到机器码 / Lowering from MLIR Through LLVM to Machine Code

27. 量化——低精度推理的工程实践 / Engineering Low-Precision Inference with Quantization

28. 分布式编译与训练——多设备编排的编译器支持 / Compiler Support for Distributed Training and Multi-Device Orchestration

29. 生产调试——真实问题的编译器视角排查 / Production Debugging from the Compiler Perspective

30. 未来方向——AI 编译器的新挑战与机遇 / Future Challenges and Opportunities for AI Compilers

本页目录

📅 创建时间:2026-06-03 🏷️ 标签:#AutoScheduling #AutoTVM #Ansor #CostModel #Tuning #MeasurementBased #Tensorization 📚 前置知识:[[/04-ai/01-llm-engineering/07-llm-evolution]](LLM 发展脉络)[[19-tvm-te]](TVM TE)[[17-kernel-primer]](Kernel 开发入门) 📚 相关知识:[[20-triton]](Triton)[[18-cutlass]](CUTLASS)


基于成本模型与实测搜索的自动调度 / Automatic Scheduling with Cost Models and Measurement-Based Search ​

┌──────────────────────────────────────────────────────────────────────────────┐ │ 📖 场景:AutoScheduler 调优后上线性能抖动 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ 你用 TVM 的 AutoScheduler(Ansor)给一个自定义算子调优: │ │ - 跑了 8 小时,2000 个配置,找到了一个比默认快 5 倍的 kernel │ │ - 但上线后发现性能抖动很大:有时候快,有时候慢 │ │ - 你怀疑 Ansor 的成本模型有偏差,在特定输入大小上失真了 │ └──────────────────────────────────────────────────────────────────────────────┘

第1节 为什么需要自动调度——手工 schedule 参数空间太大 ​

1.1 Schedule 参数空间有多大 ​

TVM TE 的 schedule 原语看似简单,但组合起来参数空间是爆炸的:

python
# 一个 GEMM 的 schedule 需要选择:
# 
# 1. Tile 大小:
#    - block_m: 16, 32, 64, 128, 256  (5个选项)
#    - block_n: 16, 32, 64, 128, 256  (5个选项)
#    - block_k: 16, 32, 64             (3个选项)
#
# 2. Thread 绑定:
#    - thread_m: 4, 8, 16, 32          (4个选项)
#    - thread_n: 4, 8, 16, 32          (4个选项)
#
# 3. Loop order (每个 tile 后的 reorder):
#    - 可能的排列数: 5! = 120 (但很多是非法的)
#
# 4. Unroll 策略:
#    - 哪些 axis unroll: 2^5 = 32
#    - unroll 阈值: 4, 8, 16, 32, 64
#
# 5. Vectorize:
#    - 哪些 axis vectorize: 2^5 = 32
#    - vector size: 2, 4, 8, 16
#
# 总参数空间: 5×5×3×4×4×32×5×32×32 = 估计 > 10^9 种配置!

# 手工测试?不可能
# Ansor 的价值:在巨大参数空间中自动找到最优配置
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

1.2 手工调优 vs 自动调度 ​

┌─────────────────────────────────────────────────────────────────────────┐
│  手工调优的局限                                                    │
│                                                                   │
│  1. 依赖专家经验:需要理解硬件细节 + schedule 原语                  │
│  2. 难以穷举:10^9 种配置,经验只能覆盖一小部分                   │
│  3. 难以迁移:不同硬件需要重新调优                                  │
│  4. 难以维护:手动找到的参数难以复现和分析                          │
│                                                                   │
│  自动调度的价值                                                    │
│                                                                   │
│  1. 系统化搜索:可以探索手工难以触及的配置                          │
│  2. 可复现:所有配置和结果都有记录                                 │
│  3. 可迁移:相同算子在不同硬件上可以自动重新搜索                   │
│  4. 可积累:历史调优数据可以用于 cost model 训练                   │
└─────────────────────────────────────────────────────────────────────────┘
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

第2节 AutoTVM 架构——调度空间、代价模型、剪枝、测量 ​

2.1 AutoTVM 的四个核心组件 ​

┌─────────────────────────────────────────────────────────────────────┐
│                      AutoTVM 系统架构                               │
│                                                                     │
│   ┌─────────────────────────────────────────────────────────────┐   │
│   │                    1. Search Space (搜索空间)                │   │
│   │                                                          │   │
│   │   你用 TVM TE 定义算子                                      │   │
│   │   AutoTVM 提供 ConfigSpace 描述所有可能的 schedule 配置      │   │
│   │   例如: tile_sizes ∈ {16, 32, 64}, unroll ∈ {0, 1}       │   │
│   └─────────────────────────────────────────────────────────────┘   │
│                              │                                     │
│                              ▼                                     │
│   ┌─────────────────────────────────────────────────────────────┐   │
│   │                    2. Cost Model (代价模型)                  │   │
│   │                                                          │   │
│   │   给定一个 schedule 配置,预测执行时间                      │   │
│   │   两种方式:                                              │   │
│   │   - 静态模型: Roofline model, analytical model            │   │
│   │   - ML 模型: GBDT, Neural Network                          │   │
│   │                                                          │   │
│   │   CostModel.predict(config) → 预测时间 (ms)               │   │
│   └─────────────────────────────────────────────────────────────┘   │
│                              │                                     │
│                              ▼                                     │
│   ┌─────────────────────────────────────────────────────────────┐   │
│   │                 3. Search Policy (搜索策略)                 │   │
│   │                                                          │   │
│   │   如何在巨大空间中高效搜索?                               │   │
│   │   - Random Search: 随机采样                               │   │
│   │   - Evolutionary: 遗传算法,变异+交叉                      │   │
│   │   - RL-based: 强化学习(PPO 等)                          │   │
│   │   - UCB: Upper Confidence Bound,平衡探索与利用             │   │
│   │                                                          │   │
│   │   policy.select() → 选择下一个要评估的配置                 │   │
│   └─────────────────────────────────────────────────────────────┘   │
│                              │                                     │
│                              ▼                                     │
│   ┌─────────────────────────────────────────────────────────────┐   │
│   │                4. Measurement (实际测量)                    │   │
│   │                                                          │   │
│   │   Cost Model 有偏差 → 需要真实测量                         │   │
│   │   在目标硬件上实际运行 kernel,记录执行时间                 │   │
│   │   测量结果反馈给 Cost Model,用于下次预测                   │   │
│   │                                                          │   │
│   │   measure(builder, runner, config) → 实测时间 (ms)         │   │
│   └─────────────────────────────────────────────────────────────┘   │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘
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

2.2 Search Space 定义 ​

python
import tvm
from tvm import te, autotvm

# 定义算子 (同 TVM TE)
N, F, K = 1024, 512, 3

# placeholder
data = te.placeholder((N, F), name="data", dtype="float32")
weight = te.placeholder((F, F), name="weight", dtype="float32")
bias = te.placeholder((F,), name="bias", dtype="float32")

# Conv2d 简化为 GEMM
# 实际使用 Conv2d 更复杂,这里用全连接层演示
k = te.reduce_axis((0, F), name="k")
out = te.compute(
    (N, F),
    lambda i, j: te.sum(data[i, k] * weight[k, j], axis=k) + bias[j],
    name="out"
)

# ============================================================
# 定义 Search Space (ConfigSpace)
# ============================================================

# 使用 @autotvm.config 不同参数化的装饰器
@autotvm.config 不同参数化的装饰器
@autotvm.tuner "grid"  # 使用网格搜索调优器
def tune_gemm(cfg, M, N, K):
    # cfg: 配置对象,调优器会修改这个对象
    # cfg.space: 搜索空间定义
    
    s = te.create_schedule(out.op)
    
    # ============================================================
    # 定义可调参数
    # ============================================================
    
    # 方法 1: cfg.space 定义离散取值
    # cfg.define_split("tile_x", N, num_outputs=2)
    # cfg.define_knob("tile_y", [16, 32, 64])
    
    # 方法 2: cfg.define_float_range / cfg.define_int_range
    # cfg.define_knob("unroll_k", [0, 1])  # 0 = 不 unroll, 1 = unroll
    
    # 在 schedule 中使用配置
    i, j = s[out].op.axis
    k = s[out].op.reduce_axis[0]
    
    # cfg.define_split 生成一个可调参数
    # 第二个参数是 axis 的长度
    # num_outputs=2 表示拆成 outer 和 inner
    ko, ki = s[out].split(k, cfg="tile_k")  # 使用名为 "tile_k" 的配置
    
    # cfg.define_knob 生成一个离散的可调参数
    # 可以直接用 cfg["xxx"] 获取当前值
    if cfg["unroll_k"].value == 1:  # 访问 knob 的当前值
        s[out].unroll(ki)
    
    # 定义 thread binding
    i_outer, i_inner = s[out].split(i, factor=32)
    j_outer, j_inner = s[out].split(j, factor=32)
    
    s[out].reorder(i_outer, j_outer, ko, i_inner, j_inner, ki)
    
    # 绑定到 GPU
    s[out].bind(i_outer, te.thread_axis("blockIdx.x"))
    s[out].bind(j_outer, te.thread_axis("blockIdx.y"))
    s[out].bind(i_inner, te.thread_axis("threadIdx.x"))
    s[out].bind(j_inner, te.thread_axis("threadIdx.y"))
    
    return s, [data, weight, bias, out]
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
69
70
71

2.3 Cost Model ​

python
# AutoTVM 支持的 Cost Model 类型

from tvm import autotvm

# 1. 静态模型 (Roofline-based)
# 适用于: 结构简单的算子,算术强度容易估计
model_static = autotvm.tensorize_cost_model.RooflineModel()

# 2. GBDT 模型 (Gradient Boosting)
# 适用于: 复杂算子,需要从大量数据中学习模式
model_gbdt = autotvm.model.GBDTTuner(
    feature_type="knob",      # 使用 schedule knob 作为特征
    loss_type="rank:listMLE", # 排序损失
)

# 3. Neural Network 模型
# 适用于: 大规模搜索,需要更好的泛化
model_nn = autotvm.model.NNTTuner(
    model="mlp",             # 多层感知机
    loss_type="rank:cosine",  # 余弦相似度排序
)

# Cost Model 训练数据
# (schedule_config, measured_latency) pairs
training_data = [
    (config_1, 1.2),   # config_1 实测 1.2ms
    (config_2, 0.8),   # config_2 实测 0.8ms
    (config_3, 2.1),   # config_3 实测 2.1ms
    # ...
]

# 训练
model_static.fit(training_data)
model_gbdt.fit(training_data, valid_data=validation_data)

# 预测
predicted_latency = model_gbdt.predict(unseen_config)
# 返回: 预测的 latency (ms)
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

2.4 Search Policy ​

python
# AutoTVM 的搜索策略

from tvm import autotvm

# 1. Random Search
# 最简单,但效率低
tuner_random = autotvm.tuner.RandomTuner(task)

# 2. Grid Search
# 适用于: 搜索空间小,参数离散
tuner_grid = autotvm.tuner.GridSearchTuner(task)

# 3. GBDT-based Tuner
# 使用 GBDT 模型指导搜索
tuner_gbdt = autotvm.tuner.GBDTTuner(
    task,
    loss_type="rank:pairwise",  # pairwise ranking loss
    num_feature_type="knob",
)

# 4. XGBoost Tuner
# 效果通常比 GBDT 好
tuner_xgb = autotvm.tuner.XGBTuner(
    task,
    plan_size=64,           # UCB 中 exploration 的数量
    feature_type="knob",
)

# 5. RL-based Tuner
# 使用强化学习策略
tuner_rl = autotvm.tuner.RLTuner(task, num_agents=4)

# 运行调优
tuner = tuner_xgb
tuner.tune(
    n_trial=1000,              # 最大试验次数
    early_stopping=100,         # 提前停止(100次无改进后)
    measure_option=autotvm.measure_option(
        builder=autotvm.LocalBuilder(build_func="cuda"),  # 本地编译
        runner=autotvm.LocalRunner(  # 本地运行
            number=10,              # 每个配置跑10次取平均
            repeat=3,               # 重复3轮
            min_repeat_ms=100,      # 最少运行100ms
        )
    ),
)
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

第3节 Ansor(AutoScheduler)——自动程序合成 ​

3.1 Ansor vs AutoTVM 的区别 ​

┌────────────────────────────────────────────────────────────────────────┐
│                     AutoTVM vs Ansor                                   │
│                                                                       │
│  AutoTVM:                                                              │
│    - 你手动定义 search space (define_split, define_knob)              │
│    - 搜索空间有限(你定义了什么,就搜什么)                             │
│    - 优点:精确控制,适合专家                                         │
│    - 缺点:容易遗漏好配置,搜索空间爆炸时难定义                        │
│                                                                       │
│  Ansor (AutoScheduler):                                                │
│    - 自动推导 search space(从 compute graph)                         │
│    - 搜索空间更大、更完整                                             │
│    - 程序合成器 (Program Synthesizer) 可以探索程序变体                  │
│    - 更好的覆盖率                                                     │
│                                                                       │
└────────────────────────────────────────────────────────────────────────┘
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

3.2 Ansor 的核心概念 ​

python
# Ansor 使用流程

from tvm import auto_scheduler

# ============================================================
# 步骤 1: 定义算子 (同 TVM TE)
# ============================================================

N, F = 1024, 512

data = te.placeholder((N, F), name="data", dtype="float32")
weight = te.placeholder((F, F), name="weight", dtype="float32")
bias = te.placeholder((F,), name="bias", dtype="float32")

k = te.reduce_axis((0, F), name="k")
out = te.compute(
    (N, F),
    lambda i, j: te.sum(data[i, k] * weight[k, j], axis=k) + bias[j],
    name="out"
)

# ============================================================
# 步骤 2: 创建 Task
# ============================================================

# Task: 一个需要调优的算子
# Ansor 会自动为这个 task 生成 search space
task = auto_scheduler.create_task(
    func=lambda: (out.op, [data, weight, bias, out]),  # compute + I/O tensors
    target="cuda",                                        # 目标硬件
)

# 打印 task 信息
print(task)

# ============================================================
# 步骤 3: 配置 Search Policy
# ============================================================

# AutoScheduler 使用的搜索策略
# 默认使用 evolutionary search + cost model
tune_option = auto_scheduler.TuningOptions(
    num_measure_trials=2000,     # 最大测量次数 (耗时长时减少)
    num_measures_per_round=64,   # 每轮测量多少个配置
    early_stopping=100,          # 提前停止
    num_search_iterations=100,   # 搜索迭代次数
    min_sample_interval_ms=50,   # 两次采样之间的最小间隔
    
    # Builder: 在目标硬件上编译 kernel
    builder=auto_scheduler.LocalBuilder(build_func="cuda"),
    
    # Runner: 在目标硬件上运行 kernel 并测量时间
    runner=auto_scheduler.LocalRunner(
        number=10,
        repeat=3,
        min_repeat_ms=100,
    ),
)

# ============================================================
# 步骤 4: 运行调优
# ============================================================

# 方法 1: 自动调优
policy = auto_scheduler.SketchPolicy(
    task,
    verbose=1,
    search_policy="evolutionary",  # 使用 evolutionary search
)
sch, args = auto_scheduler.auto_schedule(task, policy=policy, tuning_options=tune_option)

# 方法 2: 手动调优循环
# 给你更多控制
policy = auto_scheduler.SketchPolicy(task, verbose=1)
tune_option.runner = auto_scheduler.LocalRunner(number=10, repeat=3)

# 迭代搜索
for trial in range(1000):
    # 采样一个配置
    configs = policy.sample_valid_config(num=64)
    
    # 测量配置
    results = []
    for config in configs:
        latency = auto_scheduler.measure_one_config(task, config, tune_option)
        results.append((config, latency))
    
    # 更新 cost model
    policy.update(results)
    
    # 检查是否收敛
    if trial > 100 and best_latency_stable:
        break

# ============================================================
# 步骤 5: 提取最优 schedule
# ============================================================

sch, args = auto_scheduler.auto_schedule(task)
print(sch)
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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

3.3 Program Synthesizer(程序合成器) ​

Ansor 的 Program Synthesizer:

AutoTVM 的问题: 你需要手动定义 search space
  - 你定义了 tile_x ∈ {16, 32, 64}
  - 但如果最优值是 48 呢?你错过了

Ansor 的解决方案: Program Synthesizer
  - 自动从 compute graph 推导出所有可能的程序变体
  - 不需要手动定义 search space
  - 能够发现手工难以想到的程序结构

程序合成的工作原理:

1. 从 compute graph 提取 "sketch"
   sketch = 程序骨架(不完整,有空白需要填充)

2. 对 sketch 的空白进行填充
   例如: tile factor, loop order, inline/compute_at 决策

3. 通过 mutation (变异) 探索程序变体
   - Change tile size
   - Change loop order
   - Change compute_at position
   - Fuse/split axis
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

第4节 Cost Model——预测配置性能 ​

4.1 Cost Model 的类型 ​

模型类型原理优点缺点
静态 (Roofline)基于算术强度和硬件参数计算可解释,不需要训练数据精度低,不考虑具体硬件特性
GBDT/XGBoost梯度提升树训练快,能处理离散特征容易过拟合,泛化差
Neural NetworkMLP/Transformer 预测 latency泛化能力强需要更多训练数据
Hybrid结合静态 + ML 模型平衡可解释性和精度实现复杂

4.2 静态 Cost Model(Roofline-based) ​

python
# Roofline Cost Model 的实现思路

class RooflineCostModel:
    """
    基于 Roofline Model 的静态代价模型
    
    步骤:
    1. 从 schedule 推导出算术强度 (Flops/Bytes)
    2. 使用 Roofline 公式计算理论上限
    3. 结合 occupancy 估算实际性能
    """
    
    def __init__(self, hardware_params):
        self.bandwidth = hardware_params["bandwidth_gbps"]
        self.peak_tflops = hardware_params["peak_tflops"]
        
    def predict(self, schedule, compute):
        """
        预测给定 schedule 的性能
        """
        # 步骤 1: 分析 schedule 的内存访问模式
        memory_access = self.analyze_memory_access(schedule)
        flops = self.count_flops(compute)
        
        # 步骤 2: 计算算术强度
        ai = flops / memory_access
        
        # 步骤 3: Roofline 上限
        roofline_peak = min(
            self.bandwidth * ai,  # Memory-bound 上限
            self.peak_tflops      # Compute-bound 上限
        )
        
        # 步骤 4: 考虑 occupancy 和其他因素
        estimated_occupancy = self.estimate_occupancy(schedule)
        predicted_latency = roofline_peak / estimated_occupancy
        
        return predicted_latency
    
    def analyze_memory_access(self, schedule):
        """
        分析 schedule 的内存访问次数
        """
        # 从 schedule IR 分析:
        # - global memory loads/stores
        # - shared memory 使用
        # - register pressure
        pass
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

4.3 ML-based Cost Model ​

python
# GBDT Cost Model 的特征工程

class GBDTCostModel:
    """
    使用 GBDT 的代价模型
    
    特征类型:
    1. Knob Features: 直接从 schedule config 提取
       - tile sizes
       - unroll flags
       - vectorize factors
    
    2. Schedule Features: 从 schedule IR 提取
       - loop depth
       - memory access pattern
       - shared memory 使用量
    
    3. Compute Features: 从 compute 定义提取
       - 算子类型
       - 数据大小
       - reduce axis
    """
    
    def extract_features(self, schedule, compute):
        features = {}
        
        # Knob Features
        features["tile_m"] = schedule.config["tile_m"]
        features["tile_n"] = schedule.config["tile_n"]
        features["tile_k"] = schedule.config["tile_k"]
        features["unroll_k"] = schedule.config["unroll_k"]
        
        # Schedule Features
        features["num_loops"] = self.count_loops(schedule)
        features["has_shared_memory"] = self.check_smem(schedule)
        features["has_vectorize"] = self.check_vectorize(schedule)
        
        # Compute Features
        features["flops"] = self.count_flops(compute)
        features["memory_bytes"] = self.count_bytes(compute)
        features["arithmetic_intensity"] = features["flops"] / max(features["memory_bytes"], 1)
        
        return features
    
    def train(self, X, y):
        """
        X: 特征矩阵 (N × num_features)
        y: 实测 latency (N,)
        """
        import xgboost as xgb
        
        self.model = xgb.XGBRegressor(
            n_estimators=100,
            max_depth=6,
            learning_rate=0.1,
        )
        self.model.fit(X, y)
    
    def predict(self, features):
        return self.model.predict([features])[0]
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

第5节 测量-based Tuning 的陷阱 ​

5.1 硬件干扰 ​

硬件干扰是测量结果抖动的主要原因:

1. Thermal Throttling (热节流)
   - GPU 温度过高时自动降频
   - 表现为: 长时间运行后性能下降
   - 解决: 测量前让 GPU 冷却,监控温度

2. Background Processes (后台进程)
   - 其他进程占用 GPU
   - 表现为: 随机的高 latency
   - 解决: 关闭其他进程,使用独占 GPU

3. Power Capping (功耗限制)
   - GPU 功耗达到上限时降频
   - 表现为: 多个 kernel 同时运行时性能下降
   - 解决: 使用 nvidia-smi 监控功耗

4. Memory Allocation Overhead (内存分配开销)
   - 第一次运行需要分配显存
   - 表现为: 第一次测量比后续慢很多
   - 解决: warm-up run,忽略第一次测量

5. JIT Compilation Overhead (JIT 编译开销)
   - TVM/Ansor 首次编译需要时间
   - 表现为: 第一次运行慢,后续快
   - 解决: 预热编译,结果中排除首次
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

5.2 输入数据差异 ​

python
# 测量时输入数据的差异导致性能差异

# 问题 1: 不同 batch size
batch_sizes = [1, 4, 16, 64]
for batch in batch_sizes:
    # 不同 batch 的算术强度不同
    # 小 batch: 算术强度低,memory-bound
    # 大 batch: 算术强度高,compute-bound
    latency = measure(batch=batch)
    print(f"batch={batch}, latency={latency}ms")

# 问题 2: 不同 tensor shape
shapes = [(1024, 1024), (2048, 2048), (4096, 4096)]
for shape in shapes:
    # 不同 shape 的 tiling 效果不同
    # 小的 shape 可能无法充分利用 GPU
    latency = measure(M=shape[0], N=shape[1], K=shape[1])
    print(f"shape={shape}, latency={latency}ms")

# 问题 3: 不同 data layout
layouts = ["NCHW", "NHWC", "CHWN"]
for layout in layouts:
    # 不同 layout 的内存访问模式不同
    # NCHW: 连续的 channel
    # NHWC: 连续的 spatial
    latency = measure(layout=layout)
    print(f"layout={layout}, latency={latency}ms")

# 最佳实践: 
# 1. 测量时要固定输入 shape 和 layout
# 2. 在目标部署场景的 shape 上进行测量
# 3. 使用与生产环境相同的数据进行测量
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

5.3 JIT 编译开销 ​

python
# JIT 编译开销的处理

# 问题: TVM 会动态编译 kernel
# 第一次运行慢,后续运行快

# 方案 1: Warm-up (预热)
for _ in range(5):
    # 运行一个 dummy kernel,让 TVM 完成编译
    dummy_kernel(data)
    # 不记录这次的时间

# 方案 2: 记录编译时间
import time
start_compile = time.time()
kernel = tvm.build(s, args, target="cuda")
compile_time = time.time() - start_compile

# 测量时排除编译时间
for i in range(10):
    start_run = time.time()
    kernel(*args)
    run_time = time.time() - start_run
    
    # 如果 compile_time 被计入,需要减去
    corrected_time = run_time - compile_time if i == 0 else run_time
    
    record(corrected_time)

# 方案 3: 使用 AOT (Ahead-of-Time) 编译
# 在测量前预先编译所有可能的配置
configs = [...]  # 所有候选配置
compiled_kernels = {}
for config in configs:
    s = apply_config(original_compute, config)
    compiled_kernels[config] = tvm.build(s, args, target="cuda")

# 测量时直接使用预编译的 kernel
for config in configs:
    kernel = compiled_kernels[config]
    latency = benchmark_kernel(kernel, args)
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

第6节 Tensorization——映射到硬件指令 ​

6.1 什么是 Tensorization ​

Tensorization = 把特定计算 pattern 映射到硬件的矩阵乘指令

传统 GEMM:
  for i in range(M):
      for j in range(N):
          for k in range(K):
              C[i,j] += A[i,k] * B[k,j]
  
  每步: 1次乘 + 1次加 = 2 FLOPs
  效率低,没有利用 SIMD/Tensor Core

Tensorized GEMM:
  使用 WMMA/Tensor Core 指令:
  - WMMA: Warp-level Matrix Multiply Accumulate
  - 每条指令: 16×16×16 矩阵乘 = 8192 FLOPs
  
  效率 = 8192 FLOPs / 1 instruction = 高得多
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

6.2 在 TVM 中使用 Tensorization ​

python
# TVM TE 中的 tensorization

from tvm import te
from tvm import topi

# 定义 GEMM (同之前)
A = te.placeholder((1024, 1024), name="A")
B = te.placeholder((1024, 1024), name="B")
k = te.reduce_axis((0, 1024), name="k")
C = te.compute((1024, 1024), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k))

# 创建 schedule
s = te.create_schedule(C.op)

# 定义 tiling
i, j = s[C].op.axis
k = s[C].op.reduce_axis[0]

# 使用 tensorization 需要指定 tensor intrinsic
# tensor intrinsic 定义了如何使用硬件指令

# 例如: 使用 wmma tensorization (Ampere Tensor Core)
# 定义 tensor intrinsic
from tvm.topi import tensorize

# 设置 tensorization 引用的 intrin
# 需要定义:
# 1. input shape: 16×16×16
# 2. output shape: 16×16
# 3. compute: C = A @ B + C

# 这部分需要定义 TensorIntrin 对象
# TVM 提供了预定义的 tensorization intrinsic
# 例如: topi.nn.mma_intrin_group

# 配置 schedule 使用 tensorization
# s[C].tensorize(j, wmma_intrin_group)

# 注意事项:
# 1. Tile size 必须匹配 tensor intrinsic 的要求 (通常是 16 的倍数)
# 2. 需要正确设置 memory layout
# 3. 累加类型和输入类型可能不同 (fp16 input + fp32 accumulator)
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

第7节 实际调优工作流 ​

7.1 从定义 Search Space 到 Production ​

完整调优工作流:

┌─────────────────────────────────────────────────────────────────────┐
│  Phase 1: 准备                                                     │
│                                                                     │
│  1. 确定目标硬件                                                   │
│     - GPU: A100, H100, RTX 4090...                                 │
│     - 指定 target="cuda" 或 target="opencl"                        │
│                                                                     │
│  2. 确定目标 workload                                              │
│     - 具体的 shape: (B, H, S, D)                                  │
│     - 具体的 dtype: float16, bfloat16, int8...                   │
│     - 实际的 batch size 和序列长度                                 │
│                                                                     │
│  3. 准备测量环境                                                   │
│     - 关闭其他进程                                                 │
│     - 设置 GPU 为独占模式                                          │
│     - 确保散热良好                                                 │
└─────────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────────┐
│  Phase 2: 调优                                                     │
│                                                                     │
│  1. 定义算子和 search space                                        │
│     - 使用 TVM TE 定义 compute                                     │
│     - 使用 AutoTVM 或 Ansor 定义 search space                      │
│                                                                     │
│  2. 选择调优器                                                     │
│     - 小搜索空间: GridSearch                                       │
│     - 中等搜索空间: XGBoost / GBDT                                 │
│     - 大搜索空间: RL-based / Evolutionary                          │
│                                                                     │
│  3. 运行调优                                                       │
│     - 设置 num_measure_trials (取决于时间预算)                     │
│     - 监控测量结果                                                  │
│     - 注意排除 warm-up 和 outlier                                  │
│                                                                     │
│  4. 验证结果                                                       │
│     - 在多个输入 shape 上验证                                       │
│     - 检查性能稳定性 (多次运行的方差)                              │
│     - 与 baseline (cuBLAS 等) 对比                                 │
└─────────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────────┐
│  Phase 3: 部署                                                     │
│                                                                     │
│  1. 导出最优配置                                                   │
│     - 保存 schedule 和 config 到文件                               │
│     - 包括: schedule JSON, tuning records                          │
│                                                                     │
│  2. 集成到生产代码                                                 │
│     - 加载预调优的 schedule                                        │
│     - 或使用 compiled runtime                                      │
│                                                                     │
│  3. 监控和回归测试                                                 │
│     - 上线后监控实际性能                                           │
│     - 定期重新调优 (硬件变化、驱动更新)                            │
└─────────────────────────────────────────────────────────────────────┘
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

7.2 AutoTVM vs Ansor vs RL Schedulers 对比 ​

维度AutoTVMAnsor (AutoScheduler)RL-based
搜索空间定义手动定义 (define_split, define_knob)自动推导手动定义
覆盖率受限于手动定义高(程序变体搜索)受限于定义
需要专家知识高低高
调优时间短(搜索空间小)长(搜索空间大)中等
泛化能力低(只覆盖定义的配置)高(能发现新变体)中等
成本模型可选内置可选
Tensorization 支持✅ 手动指定✅ 自动推断❌ 不支持
适用场景专家调优、小搜索空间大搜索空间、自动搜索研究、复杂策略

7.3 调优脚本示例 ​

python
#!/usr/bin/env python3
"""
AutoTVM / Ansor 调优脚本模板
"""

import tvm
from tvm import te, autotvm, auto_scheduler
import numpy as np

def tune_with_autotvm(task, num_trials=500):
    """使用 AutoTVM 调优"""
    
    # 创建调优器
    tuner = autotvm.tuner.XGBTuner(task)
    
    # 配置测量选项
    measure_option = autotvm.measure_option(
        builder=autotvm.LocalBuilder(build_func="cuda"),
        runner=autotvm.LocalRunner(
            number=10,
            repeat=3,
            min_repeat_ms=100,
            enable_cpu_affinity=False,
        ),
    )
    
    # 运行调优
    tuner.tune(
        n_trial=num_trials,
        early_stopping=50,
        measure_option=measure_option,
        callbacks=[
            autotvm.callback.progress_bar(num_trials),
            autotvm.callback.log_to_file("tuning_records.log"),
        ],
    )
    
    # 加载最优配置
    with autotvm.apply_history_best("tuning_records.log"):
        with tvm.target.Target("cuda"):
            sch, args = tvm.build(task.compute_dag, task.config_space)
    
    return sch, args


def tune_with_ansor(task, num_trials=1000):
    """使用 Ansor 调优"""
    
    # 配置搜索策略
    policy = auto_scheduler.SketchPolicy(
        task,
        verbose=1,
        search_policy="evolutionary",
    )
    
    # 配置测量选项
    tune_option = auto_scheduler.TuningOptions(
        num_measure_trials=num_trials,
        num_measures_per_round=64,
        early_stopping=100,
        builder=auto_scheduler.LocalBuilder(build_func="cuda"),
        runner=auto_scheduler.LocalRunner(
            number=10,
            repeat=3,
            min_repeat_ms=100,
        ),
    )
    
    # 运行调优
    sch, args = auto_scheduler.auto_schedule(task, policy=policy, tuning_options=tune_option)
    
    return sch, args


def main():
    # 定义算子
    N, F = 1024, 512
    
    data = te.placeholder((N, F), name="data", dtype="float32")
    weight = te.placeholder((F, F), name="weight", dtype="float32")
    k = te.reduce_axis((0, F), name="k")
    out = te.compute(
        (N, F),
        lambda i, j: te.sum(data[i, k] * weight[k, j], axis=k),
        name="out"
    )
    
    # 创建 task
    task = tvm.auto_scheduler.create_task(
        func=lambda: (out.op, [data, weight, out]),
        target="cuda",
    )
    
    # 选择调优方法
    use_ansor = True  # 设置为 False 使用 AutoTVM
    
    if use_ansor:
        sch, args = tune_with_ansor(task, num_trials=1000)
    else:
        sch, args = tune_with_autotvm(task, num_trials=500)
    
    # 编译
    with tvm.target.Target("cuda"):
        module = tvm.build(sch, args, "cuda")
    
    # 保存编译结果
    module.export_library("tuned_gemm.tar")
    
    print("调优完成!结果已保存。")


if __name__ == "__main__":
    main()
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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

升华 ​

┌────────────────────────────────────────────────────────────────────────────┐
│                          🚀 自动调度核心原则                               │
├────────────────────────────────────────────────────────────────────────────┤
│                                                                            │
│  ① 先测量再搜索:代价模型有偏差,测量验证是必须的                           │
│     搜索的结果只是候选,必须在真实硬件上验证                                │
│                                                                            │
│  ② 搜索空间决定上限:Ansor 自动推导优于手动定义                            │
│     手动定义的搜索空间可能遗漏最优配置                                     │
│                                                                            │
│  ③ 测量环境要干净:硬件干扰是抖动的主要原因                                │
│     热节流、后台进程、JIT 开销都会影响测量准确性                           │
│                                                                            │
│  ④ 测量要有代表性:只在目标 shape 上调优                                   │
│     小 batch 调优的结果可能不适用于大 batch                                │
│                                                                            │
│  ⑤ 调优不是一次性:上线后要监控,定期重新调优                             │
│     驱动更新、硬件老化都可能导致性能变化                                   │
│                                                                            │
└────────────────────────────────────────────────────────────────────────────┘
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

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

必须理解(不理解就等于不会):

  • 🔴 为什么需要自动调度:手工 schedule 参数空间爆炸问题
  • 🔴 AutoTVM vs Ansor 的区别:搜索空间是如何定义的
  • 🔴 Cost Model 的作用和局限:为什么 cost model 预测不等于真实性能
  • 🔴 测量-based tuning 的陷阱:热节流、JIT 开销、数据差异如何影响结果
  • 🔴 测量环境的重要性:为什么要在与生产相同的环境中测量

AI 可查(知道去哪查就行):

  • ✅ AutoTVM 的具体 API(tuner.tune、measure_option 等)
  • ✅ Ansor 的具体参数(SketchPolicy、TuningOptions)
  • ✅ XGBoost/GBDT 的特征工程方法(特征类型和提取代码)
  • ✅ 特定硬件的 optimal 参数(benchmark 数据)
  • ✅ TVM tensorization intrinsics 的具体定义(TVM 源码)

学习状态:🟡 开始学习

最后更新于:

Pager
上一篇21. 使用 Triton 编写高性能 GPU Kernel / Triton for High-Performance GPU Kernels in Python
下一篇23. XLA 内部机制:HLO、融合与 SPMD / XLA Internals, HLO, Fusion, and SPMD

持续记录,持续成长

Copyright © Tidenflow