第一章：AI芯片软件栈全景

大家好，我是你们这堂课的主讲人。做AI芯片软件栈这行快十年了，踩过的坑比写过的代码还多。今天咱们先不急着深入细节，先站在高处，把整个软件栈的全貌看清楚。

说白了，AI芯片软件栈就是一座桥。桥的一头是算法工程师写的Python模型，另一头是芯片上跑的二进制指令。没有这座桥，再牛的算法也只是纸上谈兵。

1.1 从AI算法到芯片硬件的全链路抽象

我们先看一个最简单的流程。你写了一段PyTorch代码，定义了一个卷积神经网络。然后呢？

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(64*30*30, 10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

这段代码，从算法视角看，是一个网络结构。但从芯片视角看，它只是一堆待执行的数学运算。中间差了十万八千里。

我个人习惯把整个链路分成四个抽象层次：

• 算法层：PyTorch/TensorFlow，定义网络结构和训练逻辑
• 框架层：把算子调用翻译成中间表示（IR）
• 编译器层：做图优化、算子融合、内存分配
• 硬件层：生成芯片能执行的指令序列

你想想看，每一层都在做一件事——把上一层的抽象，翻译成下一层能理解的语言。翻译过程中，信息会丢失，也会被重构。

1.2 软件栈的分层架构

我在项目中遇到过很多次，团队里算法和硬件工程师互相抱怨。算法说芯片跑得慢，硬件说算法写得烂。其实问题出在中间层——软件栈没搭好。

一个成熟的AI芯片软件栈，通常包含这几层：

层级	核心组件	主要职责
应用层	PyTorch/TensorFlow/PaddlePaddle	算法开发、模型训练
框架适配层	ONNX、TorchScript、自定义前端	模型导入、算子映射
图优化层	计算图优化、算子融合、常量折叠	减少计算量、提升访存效率
算子调度层	算子选择、内存规划、任务调度	决定谁先执行、数据放哪里
代码生成层	指令生成、汇编优化、二进制输出	生成芯片能跑的机器码
运行时层	驱动、Runtime、内存管理	加载模型、管理资源、执行推理

嗯，这里要注意。这张表看起来是顺序的，但实际工程中，各层之间是互相影响的。比如你在图优化层做了算子融合，算子调度层就得重新规划内存。牵一发而动全身。

1.3 核心组件概览

我们挑几个关键组件，简单聊一聊。后面每个组件都会单独开一章细讲。

1.3.1 前端解析器

前端解析器负责把PyTorch或TensorFlow的模型，转换成芯片软件栈能理解的中间表示。我曾经遇到过一个坑：PyTorch的dynamic shape（动态形状）在前端解析时没处理好，导致后面所有优化都失效了。从那以后，我要求团队必须在前端做形状推导，把动态的变成静态的。

1.3.2 图优化器

图优化器是整个软件栈的「大脑」。它做的事情包括：

• 算子融合：把多个小算子合并成一个大算子，减少访存开销
• 常量折叠：把编译时就能算出来的结果提前算好
• 死代码消除：删掉那些永远不会执行的节点

说白了，图优化就是在「偷懒」——能少算就少算，能提前算就提前算。

1.3.3 内存分配器

这个组件经常被忽视，但往往是性能瓶颈。AI芯片的片上存储（SRAM）通常很小，可能只有几MB。你要把几十MB的模型参数和中间结果，合理地塞进去。我见过一个项目，因为内存分配策略不对，导致芯片利用率从80%掉到20%。

1.3.4 指令生成器

这是软件栈的最后一公里。它把优化后的计算图，翻译成芯片的指令序列。不同的芯片架构，指令集完全不同。有的芯片用VLIW（超长指令字），有的用SIMT（单指令多线程），还有的用脉动阵列。

指令生成器要做的事情很琐碎：寄存器分配、指令调度、流水线安排。每一步都影响最终性能。

1.4 软件栈的全景图

说了这么多，不如一张图来得直观。下面是我手绘的软件栈全景图，展示了从算法到硬件的完整链路。

这张图我画了很多遍。每次给新同事讲软件栈，我都会指着这张图说：「你看，每一层都在做翻译。翻译得好，性能就好；翻译得差，再好的芯片也白搭。」

1.5 为什么软件栈这么重要？

你可能觉得，芯片性能好不好，主要看硬件设计。这话对，但不全对。

我举个例子。同样是英伟达的A100芯片，你用PyTorch默认配置跑ResNet-50，和用TensorRT优化后跑，性能能差3-5倍。芯片没变，变的只是软件栈。

为什么会这样？因为硬件只是提供了「可能性」，软件栈才是把可能性变成「现实」的关键。一个好的软件栈，能把芯片的算力利用率从30%提升到90%。

1.6 我的学习建议

学软件栈，我建议你从两个方向入手：

1. 从上往下学：从PyTorch/TensorFlow开始，理解模型是怎么定义的，然后一步步往下看，看它怎么被翻译成芯片指令。
2. 从下往上学：从芯片的指令集和硬件架构开始，理解芯片能做什么、不能做什么，然后看软件栈怎么把算法映射上去。

我个人更推荐第二种。因为只有理解了硬件的「脾气」，你才能写出高效的软件栈。就像你要给一个人安排工作，你得先知道他能干什么、不能干什么。

好了，第一章就到这里。这一章我们搭好了框架，后面每一章都会深入一个组件，把细节掰开揉碎了讲清楚。