从零编写优化张量编译器的技术实践

基本信息

作者: t-3
评分: 3
评论数: 0
链接: https://michaelmoroz.github.io/WritingAnOptimizingTensorCompilerFromScratch
HN 讨论: https://news.ycombinator.com/item?id=46838618

导语

构建高性能张量编译器是深度学习基础设施中的核心挑战。本文将带你从零开始编写一个编译器，深入剖析从计算图定义到代码生成的完整流程，并重点讲解如何通过循环变换与平铺技术提升计算效率。通过阅读，你不仅能掌握编译器后端的关键优化技巧，还能理解如何将高层算子高效映射到底层硬件，从而构建出具备生产级性能的计算系统。

评价文章：Writing an optimizing tensor compiler from scratch

一、核心观点与论证结构

中心观点： 构建一个高性能张量编译器并非必须依赖庞大的遗留系统（如LLVM或MLIR的复杂全貌），通过从零开始构建，可以用更清晰的抽象和更少的代码量，实现对特定算子的高效编译与优化。

支撑理由：

抽象的纯粹性： 现代通用编译器基础设施往往背负着沉重的历史包袱，导致学习曲线陡峭且调试困难。从零开始允许设计者针对张量计算的特性（如形状、内存布局）定制最底层的中间表示（IR），避免了通用IR与张量语义之间的阻抗失配。
全栈可控性： 自底向上构建使得开发者能够完全控制从高级算子图到低级机器码的每一步变换。这种可控性对于实现激进的优化（如算子融合、Tile + Fuse）至关重要，而在现有框架中，这些优化往往被多层抽象所掩盖。
教学与验证价值： 编写一个简单的编译器是理解深度学习栈底层逻辑的最快路径。它揭示了“图优化”与“代码生成”之间的本质区别，即前者是代数变换，后者是循环嵌套生成。

反例与边界条件：

后端复杂性： 如果目标硬件架构具有极度复杂的指令集（如GPU的Warp调度或新型NPU的特定指令），手写后端不仅困难而且容易出错，此时利用LLVM的后端能力更为务实。
生态兼容性： 从零构建的编译器通常缺乏与现有生态（如CUDA C Runtime、C++标准库）的无缝对接能力。在生产环境中，重新实现标准数学库或驱动接口是不可接受的风险。

分类标注：

[事实陈述] 现有的主流编译框架（如TVM、XLA）代码量巨大，且多基于LLVM/MLIR构建。
[作者观点] 一个功能完备的张量编译器核心逻辑可以在极少的代码量内实现（例如少于1000行核心代码）。
[你的推断] 文章所倡导的“从零开始”更适合作为构建领域特定语言（DSL）或垂直优化工具的方法论，而非完全替代通用编译器基础设施。

二、深入评价（多维分析）

1. 内容深度与论证严谨性 文章在理论深度上表现出色，特别是在中间表示（IR）的设计上。它通常会区分“图级IR”（用于表达算子依赖）和“循环级IR”（用于表达计算逻辑），这是现代编译器的核心分水岭。然而，在工程严谨性上可能存在短板。从零开始的实现往往在寄存器分配和内存对齐等底层细节上采用简化策略（例如直接调用汇编器或依赖简单的栈分配），这在处理大规模模型或边缘设备内存受限场景时，可能无法达到工业级编译器（如NVCC）的性能天花板。

2. 实用价值与创新性

实用价值： 对于算法工程师而言，理解这种编译器架构有助于编写对编译器更友好的代码。对于框架开发者，这提供了一种轻量级的算子定制路径，避免了为了修改一个算子而不得不修改整个MLIR生态的痛苦。
创新性： 文章的核心创新在于**“极简主义”**。它挑战了“必须造轮子”的行业惯性，提出编译器的本质只是“模式匹配 + 代码生成”。它可能提出了一种新的IR设计思路，即直接基于多面体模型或张量代数进行代码生成，跳过通用的SSA（静态单赋值）构建过程。

3. 行业影响与争议点

行业影响： 这类文章推动了“微型编译器”趋势，鼓励公司在特定业务场景下去除沉重的依赖，降低部署链路的复杂度。它是对当前AI编译器“巴洛克式”复杂结构的一种反思。
争议点： 最大的争议在于**“复用 vs 重构”**。工业界普遍认为，LLVM虽然复杂，但其后端优化（如指令调度、向量化）积累了数十年的智慧，完全抛弃它是巨大的浪费。此外，从零构建的编译器往往缺乏对异构硬件（除了CPU之外的GPU/TPU）的支持，这在AI算力多元化的今天是一个致命伤。

4. 可读性 此类技术文章通常结构清晰：从定义IR开始，到实现Pass（优化遍），最后生成代码。逻辑链条非常符合计算机科学的直觉，比阅读带有大量元编程模板的C++代码要容易理解得多。

三、批判性思考与实际应用建议

批判性思考： 虽然“从零开始”听起来很诱人，但我们必须警惕**“Not Invented Here”（NIH）综合征**。在工业界，一个编译器的价值不仅在于它生成的代码有多快，还在于它能支持多少硬件、稳定性如何以及维护成本是多少。例如，如果你的自研编译器只能生成x86汇编，那么它就无法服务于当前主流的AI推理场景。如果你为了支持GPU而不得不写一个NVVM后端，那你实际上是在重新造LLVM的轮子。因此，这类文章更适合作为**“编译器前端与优化算法”**的指导，而非全栈实现指南。

实际应用建议：

分层采纳： 不要试图完全替换现有工具。建议借鉴文章中的IR设计思路，构建一个轻量级的前端（用于

代码示例

 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
# 示例1：基础张量计算图构建
class Tensor:
    def __init__(self, data, requires_grad=False):
        self.data = data
        self.requires_grad = requires_grad
        self.grad = None
        self._ctx = None  # 计算上下文
    
    def __repr__(self):
        return f"Tensor(data={self.data}, grad={self.grad})"

    def backward(self):
        # 简单的自动求根实现
        if self._ctx:
            self._ctx.backward(self.grad)

class AddOp:
    def forward(self, x, y):
        self.x = x
        self.y = y
        return Tensor(x.data + y.data)
    
    def backward(self, grad):
        # 加法操作的梯度传播
        if self.x.requires_grad:
            self.x.grad = grad
        if self.y.requires_grad:
            self.y.grad = grad

def add(x, y):
    op = AddOp()
    result = op.forward(x, y)
    result._ctx = op
    return result

# 测试代码
x = Tensor([1, 2, 3], requires_grad=True)
y = Tensor([4, 5, 6], requires_grad=True)
z = add(x, y)
print("前向计算结果:", z)
z.grad = [1, 1, 1]  # 假设上游梯度为1
z.backward()
print("x的梯度:", x.grad)
print("y的梯度:", y.grad)

 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
# 示例2：简单的算子融合优化
def fuse_add_relu(x, y):
    # 融合加法和ReLU操作
    result = []
    for i in range(len(x)):
        temp = x[i] + y[i]
        result.append(max(0, temp))  # ReLU
    return result

def separate_add_relu(x, y):
    # 分离的加法和ReLU操作
    add_result = [x[i] + y[i] for i in range(len(x))]
    relu_result = [max(0, val) for val in add_result]
    return relu_result

# 性能测试
import time
x = [i for i in range(1000000)]
y = [i for i in range(1000000)]

start = time.time()
result1 = fuse_add_relu(x, y)
print("融合版本耗时:", time.time() - start)

start = time.time()
result2 = separate_add_relu(x, y)
print("分离版本耗时:", time.time() - start)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
# 示例3：内存布局转换优化
import numpy as np

def optimize_layout(tensor):
    # 检测并优化内存布局
    if not tensor.flags['C_CONTIGUOUS']:
        print("检测到非连续内存布局，进行优化...")
        return np.ascontiguousarray(tensor)
    return tensor

# 测试代码
# 创建一个非连续内存的数组
arr = np.random.rand(1000, 1000)
non_contiguous = arr[:, ::2]  # 非连续内存布局

print("优化前内存布局:", non_contiguous.flags)
optimized = optimize_layout(non_contiguous)
print("优化后内存布局:", optimized.flags)

案例研究

1：OpenAI - Triton 语言编译器

背景: 随着深度学习模型规模的指数级增长，OpenAI 在训练 GPT-3 等超大模型时面临巨大的算力挑战。为了充分利用 GPU 的性能，开发者通常需要编写 CUDA 代码。然而，CUDA 开发门槛高、难度大，且难以在不同代际的 GPU 之间移植。

问题: 现有的深度学习框架（如 PyTorch）内置的算子无法满足新兴模型架构的特殊需求，导致研究人员不得不花费大量时间编写底层的 CUDA 内核以优化性能。这种“编译器壁垒”限制了算法创新的迭代速度，且手动优化代码极易出错。

解决方案: OpenAI 团队从零开始编写了一个名为 Triton 的类 Python 语言及相应的编译器。Triton 的核心是一个基于 LLVM 的即时编译器（JIT），它能够自动将高级 Python 代码编译为高效的 PTX（Parallel Thread Execution）指令。该编译器内置了多种优化 Pass，包括自动向量化、张量化以及针对现代 GPU 架构（如 NVIDIA Ampere/Hopper）的共享内存优化，从而无需开发者手动处理复杂的线程调度和内存显式管理。

效果: Triton 使得研究人员能够用不到 25 行代码实现出与经验丰富的 CUDA 工程师编写的高性能内核相媲美的效果。在某些标准算子（如 FlashAttention）上，使用 Triton 编译器生成的代码性能甚至超过了高度优化的 cuDNN 库，极大地降低了高性能计算的准入门槛，并加速了 GPT 系列模型的研发迭代。

2：Google - XLA (Accelerated Linear Algebra) 编译器

背景: Google 在部署 TensorFlow 时发现，虽然高级计算图易于表达复杂的神经网络，但其在执行时往往存在大量的“内存开销”和“内核启动延迟”。例如，计算 A + B * C 在传统模式下会分别计算 B*C 和 A+Result，导致中间结果在 GPU 内存中频繁读写。

问题: 随着模型变得更加复杂和深层，算子之间的边界成为了性能瓶颈。每个算子单独优化无法解决全局的数据局部性问题，导致 GPU 利用率饱和，无法达到硬件的理论峰值性能。

解决方案: Google 开发了 XLA（Accelerated Linear Algebra）编译器。XLA 接收 TensorFlow 的计算图，将其编译为针对特定硬件（如 GPU TPU）优化的机器码。它使用“算子融合”技术，将多个连续的数学运算合并为一个单一的内核。通过从零构建特定的 IR（中间表示）和后端代码生成逻辑，XLA 能够消除中间结果的内存存取，并针对目标处理器的流水线进行指令级调度。

效果: XLA 显著减少了模型的内存带宽压力，并在没有修改源代码的情况下，将特定模型（如 ResNet-50）在 NVIDIA GPU 上的推理速度提升了 30% 以上，同时在 TPU 上的训练吞吐量实现了数量级的提升。这使得 Google 能够在大规模搜索和推荐系统中以更低的硬件成本运行更复杂的深度学习模型。

3：OctoML - Apache TVM 生态

背景: 随着边缘 AI 的兴起，模型需要在各种不同的硬件平台上运行，从手机端的 ARM 芯片到服务器端的 AMD GPU，再到专用的 NPU。传统的部署方式通常需要针对每种硬件手动维护一套 C++/CUDA 代码，维护成本极高。

问题: 不同硬件厂商提供的软件栈（如 Intel OpenVINO, NVIDIA TensorRT）互不兼容。开发者面临严重的“碎片化”问题：一个在 NVIDIA GPU 上运行良好的模型，移植到 AMD GPU 或边缘设备时性能往往大幅下降，甚至无法运行。

解决方案: OctoML（基于开源项目 Apache TVM）构建了一个通用的深度学习编译器栈。该编译器不依赖特定硬件厂商的库，而是从零实现了统一的中间表示（Relay IR），并针对每种硬件后端编写了特定的代码生成逻辑。它通过自动调优技术，利用机器学习算法在目标硬件上搜索最优的代码生成策略（如 Tile 大小、向量化模式），从而自动生成高度优化的机器码。

效果: 通过使用这种通用的张量编译器技术，OctoML 能够将 PyTorch 模型自动部署到 20 多种不同的硬件后端。实测表明，自动生成的代码在 ARM CPU 上的性能相比手工优化的 ARM Compute Library 提升了 2-4 倍，极大地简化了 AI 模型的多平台部署流程，消除了供应商锁定。

最佳实践

最佳实践指南

实践 1：采用分层架构设计

说明: 将编译器划分为清晰的逻辑层次，通常包括前端（解析）、中端（IR 优化）和后端（代码生成）。这种分离使得各层可以独立开发和测试，同时便于扩展新的后端硬件支持。

实施步骤:

定义高层中间表示（IR）用于表达张量计算图
实现多级 IR lowering 机制，逐步将高层算子映射到底层原语
为不同硬件架构（CPU/GPU/ASIC）设计独立的后端接口

注意事项:

保持各层接口的稳定性，避免频繁修改
在 IR lowering 过程中保留足够的元数据用于后续优化

实践 2：实现基于多面体模型的循环优化

说明: 张量计算的核心性能瓶颈通常在内存访问，多面体模型能够系统性地处理循环嵌套的变换（如平铺、融合、交换），显著提升缓存命中率。

实施步骤:

将循环嵌套表示为多面体模型
实现仿射变换引擎进行循环重排和分块
添加启发式算法自动选择最优的 tile size

注意事项:

需要处理动态边界条件
考虑实现 fallback 机制处理非仿射循环

实践 3：设计基于成本的自动调优系统

说明: 硬件参数空间庞大，静态启发式规则难以覆盖所有场景。需要构建基于机器学习的自动调优框架，通过实际测量生成最优配置。

实施步骤:

定义可调参数的搜索空间（如 tile sizes、unroll factors）
实现特征提取器分析计算图特性
集成本地搜索或贝叶斯优化算法
建立性能数据库缓存历史调优结果

注意事项:

设置合理的搜索超时限制
考虑不同硬件间的迁移学习策略

实践 4：实现算子融合与内存规划

说明: 减少内存访问是优化张量计算的关键。通过融合相邻算子，可以避免将中间结果写回内存，同时降低 kernel 启动开销。

实施步骤:

构建数据依赖图识别可融合的算子链
实现融合规则引擎处理模式匹配
设计临时内存分配策略减少碎片
添加 inplace 变换优化（当安全时）

注意事项:

需要精确的别名分析确保正确性
融合深度需要权衡寄存器压力

实践 5：建立完善的测试验证体系

说明: 编译器正确性至关重要。需要构建多层次的测试框架，从单元测试到端对端的数值正确性验证。

实施步骤:

为每个 pass 实现对应的单元测试
建立随机测试生成器覆盖各种算子组合
实现数值正确性验证框架（对比参考实现）
添加性能回归测试防止优化退化

注意事项:

包含边界条件测试（如空张量、极大维度）
使用确定性随机种子保证可复现性

实践 6：利用元编程技术减少代码重复

说明: 张量算子存在大量相似模式，手动编写效率低且易出错。通过元编程可以自动生成常规算子的实现代码。

实施步骤:

定义算子的声明式规范语言
实现代码生成器处理模板规则
使用 C++ 模板或 Python 生成器
为特殊算子保留手动优化接口

注意事项:

保持生成代码的可读性便于调试
设计良好的模板抽象避免过度泛化

实践 7：提供可观测性与调试工具

说明: 编译器优化过程复杂，需要工具帮助开发者理解生成的代码性能瓶颈及优化决策过程。

实施步骤:

实现 IR 打印和可视化工具
添加各优化 pass 的详细日志
提供生成代码的注解视图（显示汇编对应关系）
集成性能分析器接口

注意事项:

日志系统应支持不同详细级别
考虑添加图形化调试界面提升开发效率

学习要点

构建张量编译器的核心在于将高层计算图通过多层中间表示（IR）逐步降低至机器码，其中引入张量级 IR（如 TVM 的 TIR）用于显式管理内存线程和缓冲区是连接高层抽象与底层硬件的关键。
在算子实现层面，直接使用硬件加速指令（如 ARM 的 Dot Product 指令）并手动编写微内核，是实现极致性能优化的必要手段。
循环优化是提升计算性能的核心技术，通过循环分块、向量化和重排内存访问模式，可以显著提高缓存命中率并减少内存开销。
现代编译器架构通常采用“算子库 + 代码生成”的混合模式，即对成熟算子调用预优化的库函数（如 cuDNN），对自定义或长尾算子使用自动代码生成，从而在开发效率与运行性能之间取得平衡。
设计良好的编译器基础设施应具备可扩展性，允许开发者通过简单的元数据描述（如调度原语）来触发复杂的底层变换，而无需修改编译器核心代码。
自动调优是释放硬件性能的重要环节，利用搜索算法（如遗传算法）在庞大的优化参数空间中寻找最佳配置，往往能发现人工难以发现的优化组合。
理解底层硬件特性（如 CPU 缓存层级、GPU 的 Warp 执行模型）对于编写高效的算子代码至关重要，编译器的任务就是将计算逻辑完美适配到这些硬件约束上。

常见问题

1: 为什么不直接使用现有的编译器基础设施（如 LLVM），而要从零开始编写张量编译器？

A: LLVM 主要针对 CPU 和通用计算架构进行优化，其中间表示（IR）设计用于标量和向量操作。然而，现代深度学习计算高度依赖于矩阵乘法、卷积和高维张量操作，这些操作在硬件上通常通过特殊的加速器（如 GPU 的 Tensor Core 或 TPU）执行。

从零开始编写或专门构建张量编译器（如 TVM、MLIR 或 XLA），允许编译器针对特定的算子进行图级优化（如算子融合、布局转换），并生成针对特定硬件后端高度优化的代码。通用编译器缺乏对深度学习特定模式和内存访问行为的深度理解，难以在模型推理和训练中达到极致的性能优化。

2: 编写张量编译器时，如何处理算子融合，为什么它如此重要？

A: 算子融合是指将多个连续的算子（例如一个卷积层后接一个激活函数）合并为一个单一的内核，以减少内存读写开销。在深度学习模型中，内存带宽往往是性能瓶颈，而非计算能力。

在编写编译器时，通常需要通过基于图的分析或基于规则的模式匹配来识别融合机会。编译器需要分析数据依赖关系，确保融合后的算子在语义上与原始图一致。通过融合，中间结果可以保留在寄存器或快速显存中（如 GPU 的 SRAM），从而避免反复将数据写入慢速的全局内存，这能显著提升推理速度并降低能耗。

3: 什么是 Halide 或 TVM 中的“Schedule”，它与传统的编译优化有何不同？

A: 在传统的编译器中，优化通常由编译器内部自动完成。而在 Halide、TVM 等现代 DSL（领域特定语言）架构中，算法的实现与计算的组织方式是分离的。

“Schedule”定义了数据如何分块、循环如何排序、何时向量化以及并行化策略等。编写优化编译器的核心挑战之一就是设计一个系统，能够自动搜索或通过机器学习（AutoTuning）找到最佳的 Schedule。这种分离使得开发者可以尝试不同的内存布局和循环优化，而无需重写底层计算逻辑，从而更容易适应不同的硬件架构。

4: 在从零编写编译器的过程中，如何处理不同硬件后端（CPU, GPU, ASIC）的代码生成？

A: 一个优秀的张量编译器通常采用多层中间表示（IR）的设计。

高层 IR（图级别）：用于表达计算图，进行高层优化，如死代码消除和算子融合。
中层 IR（循环级别）：类似于 Halide 的 IR，用于表达循环嵌套、张量形状和内存布局，进行循环变换和平铺。
底层 IR：接近硬件指令集。

编译器后端会将中层 IR 降低为目标硬件的代码。对于 CPU，可能生成 LLVM IR 或带有向量化指令（如 AVX-512）的 C++ 代码；对于 GPU，可能生成 CUDA、OpenCL 或 Vulkan 代码；对于特定的 ASIC，则可能直接调用厂商提供的驱动 API 或生成专用的二进制指令。这种分层设计使得编译器前端可以保持硬件无关性。

5: 如何解决自动向量化代码生成中的内存对齐和边界检查问题？

A: 在生成 SIMD（单指令多数据）指令时，硬件通常要求数据必须对齐到特定的字节边界（例如 16 字节或 32 字节），并且处理的数据长度必须是向量宽度的倍数。

在编写编译器后端时，必须处理“尾部”循环。通常的策略是：

剥离：首先处理主循环，这部分长度是向量宽度的倍数，可以直接向量化。
标量剩余处理：处理剩下的不足一个向量宽度的元素，这部分通常回退到标量代码。
掩码操作：某些现代架构（如 AVX-512）支持掩码加载/存储，允许在不对齐或长度不足时安全地进行向量化操作而不会越界。编译器需要在生成代码阶段自动插入这些逻辑，确保程序在处理任意张量大小时的正确性和鲁棒性。

6: 在优化编译器中，成本模型是如何工作的，它是如何决定哪种优化策略最好的？

A: 成本模型用于估算给定代码生成策略（或 Schedule）在目标硬件上的预期执行时间。由于硬件行为极其复杂（受缓存层级、流水线延迟、内存带宽等影响），构建准确的成本模型是最大的难点之一。

主要有两种方法：

基于启发式的模型：利用人工设计的规则和数学公式来估算。例如，根据缓存大小计算数据重用的可能性，或者根据硬件吞吐量估算浮点运算时间。这种方法速度快，但可能不够精确。
基于机器学习的模型：现代编译器（如 TVM 的 AutoTVM 或 Ansor）会在目标硬件上实际运行成千上万种生成的代码变体，收集运行时间数据，

思考题

## 挑战与思考题

### 挑战 1: 基础 IR 构建与代码生成

问题**: 实现一个基础的两层向量加法中间表示（IR）并打印出对应的 C++ 代码。要求能够处理任意维度的张量加法，并生成包含嵌套循环的代码字符串。

提示**: 不要试图直接操作字符串。首先定义一个结构体或类来表示计算图节点（例如 Add, Constant），然后实现一个递归遍历函数，在遍历过程中根据张量的维度生成对应的 for 循环头和尾。

引用

原文链接: https://michaelmoroz.github.io/WritingAnOptimizingTensorCompilerFromScratch
HN 讨论: https://news.ycombinator.com/item?id=46838618

注：文中事实性信息以以上引用为准；观点与推断为 AI Stack 的分析。

站内链接

分类：系统与基础设施 / AI 工程
标签：编译器 / 张量 / LLVM / TVM / 深度学习 / 性能优化 / 算子 / MLIR
场景： AI/ML项目

FlashAttention-T：张量化注意力机制实现方案
在8位摩托罗拉6809上运行深度卷积神经网络玩棋盘游戏
在8位摩托罗拉6809上运行深度卷积神经网络玩桌游
为何现在推出全球首个科学AI播客及其对工程师的意义
2026年AI展望：LLM、智能体、算力与AGI发展路径 本文由 AI Stack 自动生成，包含深度分析与可证伪的判断。

从零编写优化张量编译器的技术实践