统一的硬件视界深入理解 AI 编译器中的 TOSA 方言在构建高性能 AI 推理引擎与异构计算芯片如 NPU、ASIC、DSP的编译管线时如何将顶层五花八门的深度学习框架高效、稳定地映射到五花八门的底层硬件上一直是整个工业界的痛点。在上文中我们聊到了深度绑定 XLA 生态的MHLO。而在 MLIR 的多级方言Dialect谱系中还有一个由 Arm 领衔主导、旨在成为开放工业标准的关键角色——TOSATensor Operator Set Architecture张量算子集架构。本文将带你深入理解 TOSA 的定位、核心设计哲学以及它与 MHLO 的本质区别。1. 什么是 TOSA它的核心定位TOSA并不是一个高层图优化框架而是一套全功能的、定义明确的张量算子集规范同时也是 MLIR 生态中的原生方言。它的核心愿景是充当 AI 编译器中的“硬件抽象层HAL”。[前端框架 / 高层 IR: MHLO, StableHLO, Torch-MLIR] │ ▼ [ TOSA Dialect ] -- 硬件厂商的统一承接面 (边界清晰、算子原子化) │ ▼ [底等多级降级: Linalg / Vector / TOSA-to-NPU]如果说 MHLO 的视角更偏向于“高层图优化和框架适配”那么 TOSA 的视角则更偏向于“硬件实现与合规性Compliance”。TOSA 规定了算子的精确行为、数据类型支持和量化细节让异构芯片厂商只需要专注于完美实现 TOSA 算子就能自动对接整个上层 AI 生态。2. TOSA 的三大核心设计哲学TOSA 之所以能在异构芯片和边缘端Edge AI编译中广受欢迎得益于其独特的架构设计① 算子原子化Atomicity与极简集合为了减轻底层硬件或后端编译器实现的负担TOSA 拒绝臃肿。它将高层的复杂复合算子拆解为更小的原子操作。例如一个高层的复杂池化或特定激活函数在 TOSA 层会被彻底分解为最基础的算术和几何算子组合。这使得硬件后端的降级匹配Pattern Matching逻辑变得极其简单清晰。② 严格且内置的量化语义Quantization量化Quantization是 AI 推理落地的核心。传统的量化常常作为“补丁”附加在 IR 上导致信息丢失或精度不一致。TOSA 在设计之初就将整数低精度量化INT8/INT16作为一等公民。它在算子规范中显式定义了零点Zero Point、缩放因子Scale、以及截断Clamping和移位Shifting的精确位宽。这意味着无论在什么硬件上跑TOSA 的量化算子都能保证完全相同的数学行为和高精度对齐。③ 明确的数值合规性标准TOSA 为每个算子提供了严格的精度和数学范围限制。这使得它非常像传统 CPU 的 ISA指令集架构有一个明确的测试集Compliance Test Suite来检验硬件厂商的编译器后端是否“合规”。3. MHLO vs TOSA两者的本质区别在实际的 AI 编译器管线中例如 IREE 或 OpenXLA我们经常能看到 MHLO 和 TOSA 同时出现。它们到底有什么区别特性MHLO / StableHLOTOSA出身与背景源自 Google XLA紧跟大模型、分布式训练/推理生态。由 Arm、Linaro 等主导面向全行业异构硬件与边缘端落地。算子粒度粒度较粗。包含复杂的复合算子如特定的高层Convolution或DotGeneral语法。粒度极细。强调原子化拆解减少后端的实现变数。量化支持依赖底层的统一量化方言或类型系统组合实现。原生内置、标准化的量化语义完美适配 INT8/INT16 推理。动态形状具备极强的动态形状表达和推导能力适合大模型。设计更倾向于可预测性虽然支持动态但更擅长静态和确定性内存规划。核心职责极致的**图级融合Fusion**与高层内存优化。极致的硬件一致性与低级方言降级Lowering。4. TOSA 的典型降级管线Lowering Pathways在 MLIR 生态中TOSA 通常作为高层 IR 变换后的落脚点。一旦代码进入 TOSA它就会分化为两条典型的硬件进化路径开源通用软件路径TOSA-Linalg Dialect-Vector Dialect-LLVM IR。这是最标准的 CPU/GPU 通用编译链路利用循环变换和向量化压榨芯片多核性能。专用硬件NPU/ASIC路径TOSA-硬件专用厂商方言如特定 NPU 的后端指令。异构芯片专家直接编写 Pass 捕获 TOSA 的原子算子序列将其无缝映射到 DMA 控制器和矩阵计算阵列Tensor Core上。总结在 AI 编译器日趋模块化的今天TOSA 提供了一张关键的“硬件准入证”。它用严谨的数学行为规范和内置的量化语义在混乱的异构算力与蓬勃的高层框架之间划定了一条坚实的边界。对于 AI 基础设施开发者而言理解 TOSA 就像理解硬件的通用汇编一样——它不负责顶层的天马行空那是 MHLO 和前端框架的事但它确保了任何一段 AI 计算都能在沙砾与芯片上精准、高效地落地。
