先把背景说清楚。几个月前我在做一个小型推理运行时核心目标是让同一个量化模型的计算图既能跑到 CPU 上也能不经过重写就迁移到 GPU。最初的做法很朴素把量化卷积、量化矩阵乘、量化 ReLU 这些算子一个一个硬编码出来调度器这边用一个巨大的match分发到具体函数。模型少的时候这套东西确实很爽加算子就是加一个分支但当我把模型从三个增加到十几个再开始接触边缘检测、拉普拉斯算子这类图像算子时麻烦就来了——新后端要把所有算子重写一遍算子融合根本做不下去一个简单的算子组合就要写一堆重复的 tensor 生命周期代码。于是我把这一版推翻重来把量化算子拆成更底层的“执行原语”再在 Rust 里搭了一套可复用的执行模型。这篇博客就把整个过程完整记录下来包括设计取舍、trait 接口怎么定、如何从具体算子里“折叠”出通用原语以及实操中踩过的坑。适合正在用 Rust 写推理引擎、AI 组件或者想理解执行模型设计的开发者参考即使你是 Rust 新手只要跟着步骤走也能复现一个最小版本。1. 为什么要把量化算子拆成执行原语1.1 一开始的硬编码算子问题出在哪假设你写了一个 int8 量化卷积算子输入形状是[N, C, H, W]权重是[K, C, R, S]计算流程一般是先做 im2col 展开再做矩阵乘累加器用 int32最后乘上scale并加上zero_point截断回 int8。这段逻辑本身不复杂但要写得性能好就需要处理内存布局、缓存命中、多线程切分。我刚写成第一版时代码只有三个算子量化卷积、量化全连接、量化 ReLU。每个算子内部都自己处理输入校验、shape 推导、内存分配。后来新增了一个DepthwiseConv复制粘贴了前一个算子的大部分代码再后来要支持 GPU我打开文件就意识到事情不对——同样的 im2col、同样的 gemm 核心CPU 用 rayon 写并行循环GPU 要用 shader 或者矩阵运算库重写一遍。模型一多维护成本就失控了。这里最核心的问题不是“算子太多”而是我一开始就把算子的业务逻辑比如“这是不是量化卷积”和计算执行逻辑比如“怎么高效做乘加累加”焊死在了同一个函数里。那么换一个后端或者换一种组合方式就必须复制整个函数。这种耦合也是很多小型推理引擎发展到后期特别难重构的原因。1.2 原语是什么比算子更小、更通用的执行单元我后来参考了硬件厂商和几个成熟推理框架的思路把算子继续向下拆了一层算子可以看成一组更小执行单元的“组合”这些执行单元我统称为执行原语。执行原语不是新概念它本质上是一类不能也不需要再被语义化拆分的计算动作。举个例子Gemm矩阵乘累加参数包含左右张量、偏置、转置标志、量化系数。Im2col把卷积窗口展开成矩阵。Elementwise逐元素运算参数是一个函数指针或者枚举比如relu、clamp、add、mul。Reduce沿指定维度做归约比如sum、max。DataMove拷贝、转置、拼接、切分。量化卷积这个算子在语义上是“对图像做卷积特征提取”但它落到执行原语层面就是Im2col - Gemm - BiasAdd - Requantize的流水线。如果把量化参数作为每个原语带上的元数据那么算子就变成了一张有向无环图而不是一个硬编码函数。这样做的好处前面提到了但真正让我决定重构的原因是模型组合的灵活度和后端的迁移成本都变得可控了。1.3 这个抽象解决了三个实际问题第一是可迁移。新后端只需要实现一套原语的数量而不是针对每个模型算子做适配。比如大部分算子里的矩阵乘最终都落到Gemm原语那么 GPU 端只要把Gemm映射到对应的高性能库或者 shader 即可卷积算子原语图不需要改动。第二是可融合。原语越细越容易在构图阶段把相邻原语合并成一个融合包。比如Elementwise(relu)Requantize这两个原语就可以合成一个FusedReLURequantize省掉一次中间张量读写。这是量化模型推理性能提升非常重要的一环。第三是可测试。原来一个算子内部步骤太多出问题很难定位。拆成原语后每个原语可以单独做单元测试和性能基准组合逻辑再单独测试。出现量化误差可以直接对比Gemm底层结果和参考实现的中间结果定位问题不需要再拆代码。2. 执行模型的整体骨架分层与核心数据结构2.1 从模型算子到硬件内核的四层结构重构时我把整个运行时分成四层算子层、原语层、后端层、运行时层。算子层负责描述训练/导出模型中“卷积”“全连接”“归一化”这类语义节点它不关心具体执行原语层负责把算子翻译成一组标准原语后端层只负责一件事——把原语发射到具体硬件运行时层则承载张量池、节点调度和形状推断。这样分层之后大家各管一段。四层结构大概是层级职责代表类型算子层描述模型语义保存权重和量化参数ConvOp,GemmOp,ReluOp原语层抽象硬件无关的计算动作可组合Primitive,PrimitiveGraph后端层内存分配、内核启动、实际并行计算CpuBackend,GpuBackend运行时层拓扑排序、张量生命周期、上下文管理ExecutionContext,GraphRunner这个分层最关键的点在于原语层不依赖任何后端它只描述“做什么”而“怎么做”由后端决定。这样原语列表本身就成了整个运行时最大的接口面原语越稳定后端的适配成本就越低。2.2 用 trait 表达原语和执行上下文在 Rust 里我直接用了 trait 来表达执行原语。第一个版本pub trait Primitive: Send Sync { fn name(self) - str; fn execute(self, ctx: mut ExecContext) - Result(); }看起来很简单但这里有几个细节值得展开。第一我要求Send Sync因为图执行阶段通常会尝试并行执行没有依赖关系的多个原语如果原语内部不允许跨线程共享就没法做并行调度。第二execute走的是可变上下文原语内部不持有张量数据所有中间结果都通过ExecContext存取。这样调度器可以掌握哪些张量已经被消费能及时回收内存。后来我加上了形状推断接口原语的 trait 变成了这样pub trait Primitive: Send Sync { fn name(self) - str; fn infer_shape(self, inputs: [Shape]) - ResultVecShape; fn execute(self, ctx: mut ExecContext) - Result(); }形状推断放在执行之前好处是可以提前为所有中间张量分配内存池空间执行阶段减少分配次数。像 Rust 这种所有权语言如果在执行最终阶段才动态分配张量不仅慢而且容易出现借用冲突。提前推断形状后ExecContext会持有张量池原语之间通过TensorId引用而不是传递ArcTensor这样的重量级句柄。2.3 张量、形状和内存池怎么组织张量的抽象我用了非常朴素的结构pub struct Tensor { pub id: TensorId, pub shape: Shape, pub dtype: DType, pub buffer: BufferId, }BufferId是后端分配的一块显式缓冲区它不直接暴露裸指针或Vecu8给原语使用。原语执行时通过后端提供的视图接口拿到缓冲区切片最多加一些生命周期限制。比如 CPU 后端pub struct CpuTensorViewa { pub data: a [u8], pub shape: Shape, pub dtype: DType, }这里一开始我犯过蠢——让原语直接拿mut [u8]新手看起来很方便但很快会遇到一个问题很多原语需要同时读两个输入并写一个输出。如果两个输入引用了同一块底层缓冲区Rust 的借用检查器会直接拒绝编译。与其到处.clone()不如在ExecContext里做一层张量版本管理每次写入时判断是否存在别名再决定是原地修改还是分配新缓冲区。内存池方面CPU 后端用了Vecu8池加 freelist分配大块内存后按需切块回收时不需要释放直接标记为空闲。GPU 后端相对复杂但缓冲区句柄的抽象方式是一样的。整个执行模型靠ExecContext串起来它对外暴露接口可以注册算子、构建原语图、分配张量、执行整个图。3. 核心实现量化算子到执行原语的折叠3.1 量化卷积算子里到底藏着什么要拆好一个量化卷积算子先得把它内部计算过程完整列出来。以 int8 对称量化为例一个典型的量化卷积步骤如下输入x为 int8权重w为 int8。对输入做 im2col得到N, K, C*R*S的矩阵。矩阵乘并把累加器放到 int32。加偏置偏置一般是 int32。乘权重 scale 和输入 scale 的乘积并加上输出 zero point截断到 int8。如果你把这些步骤写进一个函数里函数会很长但拆开看其实每个步骤都是一个非常通用的原语。我的折叠原则是只要能通过参数表达差异化的原语就不用单独算子。量化卷积和量化全连接的区别本质上是全连接不需要 im2col其余几乎一样。因此我只需要两个原语级别的东西一个Im2col原语一个Gemm原语再加上BiasAdd和Requantize就能组合出卷积和全连接。3.2 拆解把卷积拆成 Im2col、Gemm、偏置和量化这里以一次 3x3 卷积为例。原语图是这样构建的let mut graph PrimitiveGraph::new(); let im2col graph.add_node(Im2col { kernel_size: [3, 3], stride: [1, 1], padding: [1, 1], dilations: [1, 1], }); let gemm graph.add_node(Gemm { transpose_a: false, transpose_b: true, beta: 0.0, quant: QuantParams { input_scale, weight_scale, output_scale, zero_point }, }); let bias graph.add_node(Elementwise { op: ElementwiseOp::Add, dtype: DType::I32 }); let requant graph.add_node(Requantize { output_scale, zero_point });Gemm原语里的QuantParams是整个设计里很关键的一笔。它记录的是原语实例化的量化参数而不是全局状态。这样同一个Gemm原语实现无论被量化卷积引用还是量化全连接引用都能按各自的 scale 和 zero point 工作。这个思路就是“折叠”把不同算子之间的差异收敛成参数而不是复制代码。3.3 用组合方式实现一个量化全连接算子全连接算子就简单多了不需要Im2col直接Gemm - BiasAdd - Requantize。在 Rust 里的执行代码大概是这样let g Gemm { quant: QuantParams { input_scale: 0.01, weight_scale: 0.02, output_scale: 0.1, zero_point: 0 }, transpose_b: true, ..Default::default() }; g.execute(mut ctx)?;当你用原语组合算子时真正的工作其实是图构建和参数校验。原语本身没有“全连接”这个概念它只认识“一个矩阵乘和一个逐元素操作”。这种设计会迫使你把模型配置和计算内核分离而不是在代码里写满if is_conv { ... } else if is_fc { ... }。有一位朋友看我代码后说这套东西更像“指令集”而不是“算子库”我觉得这个比喻挺贴切。3.4 泛化拉普拉斯算子怎么搭进来热搜词里一直有人提“拉普拉斯算子”其实它在图像处理里就是一个固定系数的卷积核。在传统图像处理场景拉普拉斯边缘检测用的是浮点卷积核[0,-1,0; -1,4,-1; 0,-1,0]。如果你的原语层支持浮点Gemm或者更通用的Convolution那么这个算子直接复用一个卷积原语唯一不同的只是权重值。我在重构后做过一个小小的边缘检测 demo把拉普拉斯算子的权重直接塞进已实现的量化卷积原语图里然后跑到 GPU 上全程新写的代码不超过一百行。这就是通用执行原语的直接价值——业务层只需要关心参数不用再关心计算结构。4. 可迁移后端CPU/GPU 的差异如何被原语吞掉4.1 后端接口设计要让原语在不同硬件上执行后端抽象需要暴露足够的基础能力。我设计的Backendtrait 核心方法如下pub trait Backend: Send Sync { fn alloc(mut self, size: usize) - ResultBufferId; fn dealloc(mut self, id: BufferId) - Result(); fn read(self, id: BufferId) - ResultVecu8; fn write(mut self, id: BufferId, data: [u8]) - Result(); fn execute_primitive(mut self, prim: dyn Primitive, ctx: mut ExecContext) - Result(); }初看我以为execute_primitive会变成一个很长的分支这是后端收到原语后需要分派给具体实现。但后来我把Primitivetrait 改成了“自带后端实现”也就是每个原语实例知道自己是跑在哪个后端上的于是execute内部会调用后端提供的小型 kernel 注册表。这样做的原因很简单一个原语在不同后端上的实现差异太大如果在后端里统一分派原语数量一多分支就爆炸。4.2 CPU 后端的瓶颈与优化CPU 后端首先要解决的是访存密度和并行粒度问题。大量使用算子时性能挑战往往不在算子本身而在于中间张量写回内存再重新读出来的次数。举个例子卷积原语执行完输出是一个 int32 张量如果下一步是Requantize那么这个 int32 张量完整写一遍内存、再读一遍对于大特征图来说浪费非常明显。我实测下来单是“融合Gemm和Requantize”这一项在一个 1x1 量化卷积上就能提高约 20% 到 30% 的端到端速度。融合方法也很简单给Gemm原语加一个可选字段fusion: OptionFusedElementwise如果该字段存在则在每个输出通道的 int32 累加完成后直接做一次乘加截断不落内存。CPU 后端这里我用 rayon 做行级并行每个线程负责一部分输出行融合逻辑在线程内部完成不会有竞争。另一个容易踩的坑是内存对齐。int8 数据如果希望用 SIMD 加速最好按 64 字节对齐分配。我在Vecu8池上做了一层对齐包装分配器返回的缓冲区都保证对齐。如果你一开始就忽略对齐将来要接 AVX/NEON 时代码可能得整个重写。4.3 GPU 后端的低成本接入方案GPU 后端我选了 wgpu因为它是纯 Rust 生态下覆盖面较广的跨平台图形/计算接口可以同时覆盖 Vulkan/Metal/DX12省去维护多套 API 的麻烦。接入过程验证了原语抽象拆得好不好。我在 CPU 端已经实现了Gemm、Im2col、Elementwise三种原语GPU 端要做的事情很简单alloc/dealloc/read/write替换成 wgpu buffer 操作。Gemm映射成基于 wgpu compute shader 的分块矩阵乘。Im2col映射成texture上的并行内存复制 kernel。Elementwise映射成一个通用的逐元素 compute shader。算子层、原语图、调度器全部不用改。从零开始写 GPU 后端到把 LeNet 量化模型跑通大概花了两周业余时间。如果把旧版硬编码算子适配到 GPU我的估算是最少一个月因为每个算子都要重新写一遍 shader 和内存流转逻辑。5. 实操记录从零搭建最小执行模型5.1 环境与依赖如果是新环境先装 Rust。我统一用 rustupcurl --proto https --tlsv1.2 -sSf https://sh.rustup.rs | sh rustup update stable项目本身我用 cargo workspace核心库命名为exec-core依赖很克制rayon做 CPU 并行half处理半精度数据thiserror做错误类型。性能基准用的是criterion这部分你可以后续再加。不推荐一开始就堆一堆依赖否则连自己都无法判断性能瓶颈到底在哪一层。5.2 定义原语注册表原语注册表是整个执行模型和外部解释器之间的桥梁。我希望模型解析层只认识字符串.op_name”不直接依赖具体原语类型。注册表大致是这样pub struct Registry { factories: HashMapString, Boxdyn Fn(serde_json::Value) - ResultBoxdyn Primitive, } impl Registry { pub fn parse_node(self, node: GraphNode) - ResultBoxdyn Primitive { let factory self.factories.get(node.op_type) .ok_or_else(|| anyhow!(unknown op: {}, node.op_type))?; factory(node.attributes) } }这里面有一个经验工厂函数的入参最好直接是平台的配置对象而不要传原始 JSON。我一开始图省事直接传 JSON 字符串结果每个原语创建时都要重复解析、字段校验报错信息还很差。后来改成强类型配置对象易用性和可维护性都提升了一个档次。5.3 前向执行流程执行一个图的流程比较固定但每一步都决定上层体验遍历图节点做拓扑排序。对每一个节点用注册表把节点信息解析成Boxdyn Primitive。调用infer_shape把中间张量形状推出来为所有中间张量预分配BufferId。按拓扑顺序执行execute每执行完一个节点就把引用计数清掉能立刻释放缓冲区。执行器用了一个有向无环图调度器对有多个输入的节点要等所有上游完成。如果两个分支之间没有数据依赖就可以用rayon::join并行。这里初始实现踩了一个很大的坑直接用ThreadPool嵌套并行原语遇到 GPU 后端时缓冲区同步开销反而更大。所以后来我加了一个开关CPU 后端默认开并行GPU 后端默认顺序执行因为 GPU 端 kernel 启动本身是异步提交的CPU 侧并行调度的收益远小于复杂度代价。5.4 基准和结果我用三个模型验证执行模型量化 LeNet、量化 MLP、一个 3x3 量化卷积的边缘检测。结果没有特别惊艳但验证了抽象方向是对的量化 MLP 在原语组合模式下端到端性能和手写硬编码版本几乎一致差异在 2% 以内。量化卷积在做了Gemm Requantize融合后比未融合版本快约 22%。迁移到 GPU 后MLP 小模型因为 kernel 启动开销较大CPU 反而更快但大卷积模型 GPU 明显占优。这个结果也说明执行原语抽象不会白白让你损失性能关键是原语要支持融合而硬件选型则要结合实际模型大小不要盲目 GPU。6. 常见问题与排错实录6.1 泛型还是 trait object调度器的快乐与痛苦一开始我很想把原语设计成泛型这样 CPU 端可以让每个原语的execute内联获得更好的优化。但原语图里放的具体类型数量不定泛型会迅速导致类型爆炸。所以最终我放弃了泛型统一接口改用Boxdyn Primitive。每次执行多一次虚函数调用在深度学习计算里占的时间可以忽略不计。真正让我头疼的是三种灾难组合闭包捕获的上下文、trait object 上带泛型方法、以及原语之间存在共享状态。这几个问题加起来编译期报错能把自己绕晕。我的经验是如果某个原语非常依赖编译期优化就把它限定在 CPU 后端内部作为一个 specialization 实现不要塞进公共原语 trait公共 trait 只保留最稳定、最通用的执行动作。6.2 内存飘移Buffer 明明还活着怎么被释放了有段时间 CPU 后端上出现了很诡异的结果两个不同模型跑同一个原语图第一个结果正确第二个结果随机错。排查了很久发现问题出在内存池复用逻辑上。我的调度器在节点执行完成后会立刻回收输入缓冲区到 freelist但没考虑“某些原语内部还握着临时视图”。CPU 端我修的方法是给每个原语增加一个post_verify阶段在执行返回后检查它仍引用的BufferId是否已经被回收如果被回收就报一个清晰的内存生命周期错误而不是让数据静默被覆盖。6.3 量化误差忽大忽小不是算子的问题量化算子的误差问题排查起来最绕。同样的模型CPU 上单线程结果和 GPU 结果偶尔对不上一开始怀疑后端实现有 bug后来发现是 int32 累加顺序不同导致的截断差异。GPU 端我在分块 GEMM 里用的是每块局部累加再合并和 CPU 端逐行累加的舍入顺序不一样。这个问题无法完全消除只能通过统计误差上界来验证可接受。我把量化参数加了容差检查如果 GPU 和 CPU 的最终输出误差超过某一个阈值就提示输出重新校准。这件事也说明原语抽象再干净量化原语的实现还是得针对不同后端标注舍入策略。6.4 融合后张量形状对不上融合原语最容易犯的错是把两个原语合并后假设中间结果的形状不变但实际位移了通道。比如某些归一化层需要把通道从[N, C, H, W]重排成[N, H, W, C]融合时如果忽略布局转换形状推导就会出错。我的解决方案是在图层面插入着名“布局转换原语”而不允许融合原语内部偷偷做这种重排。牺牲了一点融合粒度但换来的是形状系统非常容易验证。写在最后的体会这套执行模型从设计到落地大概花了两个月中间推倒重来了两次。我最大的体会是在 Rust 里做这类系统真正的复杂度往往不是“怎么写某个算子”而是“如何让一组原语在没有隐式全局状态的前提下可靠协作”。trait 和所有权语义会逼你把边界画清楚这是好事但画得太细会让代码率整体变慢画得太粗又会把后端耦合进来。我自己后来的标准是原语只描述计算动作不持有张量数据不直接管理后端资源。如果你也在做类似的东西建议先画一张执行原语清单把每个算子的计算步骤拆到最小不可再拆的单元再考虑怎么写代码。拆完之后你会发现很多看起来很高深的算子最后也就是几个原语来回组合而已。