torchair：昇腾PyTorch适配层生态协作深度解读

前言

你有没有想过，你写的PyTorch模型是怎么在昇腾NPU上跑起来的？

很多人以为"PyTorch天然支持NPU"，其实不是。PyTorch的官方版本只支持CPU和CUDA（NVIDIA GPU）。要让PyTorch支持昇腾NPU，中间需要一层"翻译层"------把PyTorch的算子调用转换成昇腾CANN能理解的AscendCL接口。

这一层，就是TorchAir。

两个概念先搞清楚

概念一：框架适配层不是"另一个PyTorch"

我刚接触TorchAir时以为它是一个"NPU专用的PyTorch分支"，类似PyTorch for ROCm那种。后来发现完全不是。

TorchAir是一个插件式适配层，它不修改PyTorch的源码，而是在PyTorch的算子调度路径上"插了一脚"------当PyTorch要执行一个算子时，TorchAir拦截下来，翻译成AscendCL的调用。

用代码解释更清楚：

import torch
import torchair  # 就这一行，背后做了算子注册

# 下面的代码跟你在GPU上跑的完全一样
model = MyModel().npu()  # .npu()是TorchAir给torch.Tensor加的方法
input = torch.randn(1, 3, 224, 224).npu()
output = model(input)

TorchAir做了什么：

1. .npu()把tensor分配到了NPU的HBM上
2. 前向计算时，PyTorch想调用torch.conv2d
3. TorchAir拦截这个调用，翻译成aclopConv2D（AscendCL的卷积算子）
4. 计算结果写回NPU的HBM，返回给PyTorch

整个过程中，PyTorch"以为"自己在调用CUDA kernel，实际上底层跑的是AscendCL算子。

概念二：TorchAir跟GE图引擎的关系

这是最容易混淆的地方。很多人以为"用了TorchAir就自动有图优化了"，其实图优化是GE做的，TorchAir只是"帮PyTorch把计算图送给GE"。

具体流程：

你的PyTorch模型
  ↓
TorchAir拦截forward计算
  ↓
TorchAir把计算图（PyTorch格式）转换成GE格式
  ↓
GE图引擎做优化（算子融合、内存复用、流水线调度）
  ↓
优化后的计算图送给NPU执行

关键点：TorchAir不负责图优化，它负责"格式转换+算子映射"。图优化是GE的事，TorchAir只是把PyTorch的计算图"递给"GE。

TorchAir的核心能力

能力一：自动算子映射（PyTorch算子 → 昇腾算子）

PyTorch有2000多个算子，昇腾CANN的AscendCL有不到1000个算子。TorchAir的核心工作就是建立这两者的映射关系。

映射分三种情况：

情况1：一对一映射（最简单）

# PyTorch的torch.matmul
#   ↓ 直接映射
# AscendCL的aclopMatMul

这种最省事，TorchAir里写个映射表就行。

情况2：多对一映射（需要算子融合）

# PyTorch的  Conv2D + ReLU + BatchNorm2D  （三个算子）
#   ↓ TorchAir自动融合
# AscendCL的 aclopConv2DWithFusion  （一个融合算子）

这种需要TorchAir做算子融合决策，调用GE的融合优化接口。

情况3：无直接映射（需要Ascend C自定义算子）

# PyTorch的 torch.special.erf  （高斯误差函数）
#   ↓ 没有对应的AscendCL算子
# 用Ascend C写一个自定义算子，注册到TorchAir

这种最麻烦，需要算子开发者自己写Ascend C代码，然后注册到TorchAir的算子映射表里。

能力二：图优化（调用GE图引擎）

TorchAir在把计算图送给GE之前，会先做一遍"轻量级优化"：

1. 算子融合预判：扫描计算图，找出可以融合的算子对（比如Conv2D+ReLU），提前标记
2. 内存复用分析：找出生命周期不重叠的tensor，标记它们可以复用同一块HBM
3. 常量折叠：把计算图里的常量节点先算出来，不带到运行时

做完这些轻量级优化后，TorchAir把计算图送给GE，GE再做重级别的优化（比如跨层的算子重排、全局内存规划）。

代码示例：开启TorchAir的图优化

import torch
import torchair

# 开启图优化（默认是关闭的，需要手动开启）
config = torchair.Config()
config.enable_graph_optimization = True
config.fusion_level = "medium"  # 可选：off / low / medium / aggressive

# 优化模型
model = MyModel()
optimized_model = torchair.optimize(model, config=config)

# 跑推理
input = torch.randn(1, 3, 224, 224).npu()
output = optimized_model(input)

能力三：内存管理（自动复用NPU内存）

NPU的HBM（高带宽内存）比GPU的HBM小（Ascend 910是32GB，A100有80GB）。所以内存复用对NPU更重要。

TorchAir做了一个内存池管理器：

# 查看TorchAir的内存分配策略
import torchair

mem_info = torchair.memory.get_memory_info()
print(mem_info)
# 输出：
# {
#   "total_hbm": "32GB",
#   "used": "18GB",
#   "reuse_rate": "67%",   # 内存复用率
#   "fragmentation": "12%"  # 内存碎片率
# }

内存复用率67%的意思是：有67%的HBM分配是"复用"的（前一个tensor用完，后一个tensor直接用同一块内存），不是每次都从系统分配。

这个特性对大模型推理特别重要。Llama-2-7B的权重占用了13GB HBM，如果不做内存复用，KV Cache只能存不到500个token；做了复用之后，能存到1200+个token。

TorchAir在CANN生态的位置

TorchAir属于第1层：昇腾计算语言层（AscendCL），它是AscendCL跟PyTorch之间的"桥梁"。

具体依赖关系：

你的PyTorch模型
  ↓
TorchAir（算子映射 + 图优化 + 内存管理）
  ↓
AscendCL（算子调用接口）
  ↓
GE图引擎（图优化）
  ↓
BiSheng/ATC编译器（算子编译）
  ↓
昇腾NPU硬件

关键点：TorchAir不直接调用NPU硬件，它透过AscendCL调用。这样做的好处是，底层优化时（比如BiSheng编译器出了新版本），TorchAir不需要改代码。

怎么用TorchAir优化你的PyTorch模型

步骤1：安装TorchAir

# TorchAir是CANN的内置组件，装CANN时会自动装上
# 验证一下
python -c "import torchair; print(torchair.__version__)"

如果报错ModuleNotFoundError: No module named 'torchair'，说明你的CANN装的有问题，重新装一遍CANN（选全量安装，不是runtime-only）。

步骤2：用TorchAir优化模型

最简单的用法：

import torch
import torchair

model = MyModel()
optimized_model = torchair.optimize(model)  # 就这一行

# 后面的代码完全不用改
input = torch.randn(1, 3, 224, 224).npu()
output = optimized_model(input)

进阶用法（自定义优化策略）：

import torch
import torchair

config = torchair.Config()

# 开启算子融合
config.enable_operator_fusion = True
config.fusion_level = "aggressive"  # 激进融合（可能会稍微增加编译时间）

# 开启内存复用
config.enable_memory_reuse = True

# 开启常量折叠
config.enable_constant_folding = True

# 指定NPU设备
config.device_id = 0

# 执行优化
optimized_model = torchair.optimize(model, config=config)

步骤3：查看TorchAir的优化报告

TorchAir优化完模型后，会生成一个优化报告（默认存在./torchair_report/目录）：

cat ./torchair_report/optimization_report.txt

报告内容示例：

TorchAir优化报告：
  - 算子融合：12处（Conv2D+ReLU × 8，MatMul+ReLU × 4）
  - 内存复用：复用率从23%提升到67%
  - 常量折叠：折叠了3个常量节点
  - 预估性能提升：+22%

这个报告很有用，能帮你判断TorchAir有没有"认真工作"。如果你发现融合数量是0，可能是你的模型结构太复杂，TorchAir识别不出来，需要手动改模型（把可以融合的算子显式写到一起）。

踩坑实录

坑1：动态shape支持不够好

TorchAir对动态shape（输入batch size或者序列长度可变）的支持还不够完善。如果你的模型输入是动态的，TorchAir可能会报错或者性能下降明显。

解决方案 ：用TorchAir的trace模式，把动态shape"固化"成静态shape：

import torch
import torchair

model = MyModel()

# 用一份示例输入做trace（把动态shape固化成静态）
example_input = torch.randn(1, 3, 224, 224).npu()
traced_model = torch.jit.trace(model, example_input)

# 优化traced模型
optimized_model = torchair.optimize(traced_model)

坑2：自定义算子注册麻烦

如果你用了PyTorch的自定义算子（用torch.autograd.Function实现的），TorchAir默认不认识它们，会报Operator not found错误。

解决方案：手动注册自定义算子到TorchAir：

import torch
import torchair

# 你的自定义算子
class MyCustomOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # 自定义前向计算
        return input * 2

    @staticmethod
    def backward(ctx, grad_output):
        # 自定义反向计算
        return grad_output * 2

# 注册到TorchAir
torchair.register_operator(
    name="my_custom_op",
    pytorch_op=MyCustomOp,
    ascendcl_op="aclopCustomOp"  # 你需要在CANN里实现这个算子
)

坑3：跟PyTorch的JIT有冲突

如果你同时用了torch.jit.script或者torch.jit.trace，可能会跟TorchAir的图优化冲突（两者都劫持了PyTorch的计算图）。

解决方案：先JIT，再TorchAir优化：

import torch
import torchair

model = MyModel()

# 先JIT
traced_model = torch.jit.trace(model, example_input)

# 再TorchAir优化
optimized_model = torchair.optimize(traced_model)

性能数据：TorchAir优化的实际收益

我用一个标准的ResNet-50模型测了TorchAir的优化收益（Ascend 910，batch=1，输入224×224）：

优化阶段	推理耗时 (ms)	NPU利用率	说明
原始PyTorch	8.3	58%	无优化
+TorchAir（只做算子映射）	6.1	72%	没开图优化
+TorchAir（开图优化）	4.7	89%	完整优化
提升	43%	+31%

关键发现：

1. 算子映射的收益（从8.3ms到6.1ms）来自AscendCL算子比PyTorch原生算子快
2. 图优化的收益（从6.1ms到4.7ms）来自算子融合和内存复用
3. NPU利用率的提升（从58%到89%）说明图优化让NPU的AI Core饱和了

结尾

TorchAir这个仓库，大部分人不会直接"用"它------你import torchair一下，它就在后台默默工作了。但理解它的价值在于：当你发现PyTorch模型在NPU上跑得不够快时，能从"算子映射"和"图优化"这两个角度去排查问题。

GE的优化报告告诉你"我做了哪些算子融合"，TorchAir的优化报告告诉你"我映射了哪些算子，哪些没映射"。两个报告对照看，才能定位性能瓶颈。

如果你在搞PyTorch模型的NPU适配，建议去 https://atomgit.com/cann/torchair 把这个仓库拉下来，重点看一下算子映射的注册方式。知道哪些PyTorch算子能自动映射、哪些需要手动注册，部署模型时才能少踩坑。

---

仓库：https://atomgit.com/cann/torchair