昇腾CANN ops-nn MatMul 算子：一条调用链背后的五层仓库
这篇文章深入解析了PyTorch中torch.matmul在昇腾NPU上的完整执行流程。调用链分为五层架构：1)框架适配层通过torch_npu将PyTorch调用转换为AscendCL接口；2)C API层进行参数校验；3)ops-nn算子库实现算子融合（如MatMul+Bias+Activation）；4)ops-blas库负责GEMM分块计算和缓存优化；5)catlass模板库生成底层硬件指
·
前言
你写一行 torch.matmul(x, y)，背后发生了什么？
从 PyTorch 的一次 torch.matmul，到昇腾 NPU 真正执行矩阵乘法，中间经过了 ops-nn、ops-blas、catlass 三层仓库的接力。这篇文章用费曼科普的方式，把这条调用链拆干净，让你能说清楚每层在干什么。
从一次 torch.matmul 说起
Python 代码到 NPU 执行
import torch
# 这行代码触发的调用链：
x = torch.randn(16, 512)
y = torch.randn(512, 2048)
# 方案1：直接 matmul
z1 = torch.matmul(x, y)
# 方案2：通过 torch_npu（昇腾适配）
z2 = torch_npu.matmul(x, y)
无论哪种写法，最终都要落到昇腾 NPU 的硬件上执行。
调用链的五层
整体架构图
┌─────────────────────────────────────────────┐
│ PyTorch / torch_npu │ ← 第一层：框架适配
└────────────────────┬──────────────────────┘
│
┌────────────────────▼──────────────────────┐
│ AscendCL（aclnnMatMul） │ ← 第二层：C API
└────────────────────┬──────────────────────┘
│
┌────────────────────▼──────────────────────┐
│ ops-nn（MatMul + 融合） │ ← 第三层：算子库
└────────────────────┬──────────────────────┘
│
┌────────────────────▼──────────────────────┐
│ ops-blas（GEMM 分块） │ ← 第四层：BLAS 库
└────────────────────┬──────────────────────┘
│
┌────────────────────▼──────────────────────┐
│ catlass（GEMM 模板） │ ← 第五层：硬件模板
└────────────────────┬──────────────────────┘
│
┌────────────────────▼──────────────────────┐
│ 昇腾达芬奇 NPU（CUBE 单元执行） │ ← 硬件层
└─────────────────────────────────────────────┘
每层干什么
| 层级 | 仓库 | 职责 | 关键产物 |
|---|---|---|---|
| 框架适配 | torch_npu | PyTorch → AscendCL | 算子注册 |
| C API | AscendCL | 统一接口，参数校验 | aclnnMatMul |
| 算子库 | ops-nn | 融合、融合策略 | MatMul+Bias+Act |
| BLAS 库 | ops-blas | 分块调度、缓存优化 | GEMM |
| 硬件模板 | catlass | Cube 指令生成 | 汇编 |
| 硬件 | 达芬奇 NPU | 矩阵乘法 | 矩阵结果 |
第三层：ops-nn 的 MatMul
ops-nn 里的 MatMul 是什么
ops-nn 是"神经网络层算子库"，它的 MatMul 比 ops-blas 的 GEMM 多了一些东西：
ops-nn MatMul = GEMM + BiasAdd + Activation
# ops-nn MatMul 的融合模式
# 原始调用（3次 NPU 调用）：
# 1. MatMul
# 2. BiasAdd
# 3. Activation（GELU/SiLU/ReLU）
# ops-nn 融合后（1次 NPU 调用）：
# 一个 kernel 跑完 MatMul + Bias + GELU
为什么 ops-nn 要融合
| 调用方式 | Launch 开销 | 数据搬运 | 总耗时 |
|---|---|---|---|
| 分离调用 | 3×1ms | 3×HBM | 5ms |
| 融合调用 | 1×1ms | 1×HBM | 2ms |
融合后，数据只搬一次，kernel 只 launch 一次。
ops-nn MatMul 的融合配置
import cann
# 创建融合 MatMul（MatMul + BiasAdd + GELU）
matmul_op = cann.ops.nn.MatMul(
transpose_a=False,
transpose_b=False,
activation="GELU", # 激活函数
output_dtype="float16"
)
# 调用
output = matmul_op(input_tensor, weight_tensor, bias_tensor)
第四层：ops-blas 的 GEMM
ops-blas 是什么
ops-blas 是"基础线性代数算子库"，专注矩阵乘法（GEMM = GEneral Matrix Multiply）。
# ops-blas GEMM 的核心参数
def gemm(a, b, c=None,
alpha=1.0, beta=0.0,
transpose_a=False, transpose_b=False,
m=None, n=None, k=None):
"""
c = alpha * op(A) @ op(B) + beta * c
A: (m, k) 或 (k, m) 如果转置
B: (k, n) 或 (n, k) 如果转置
C: (m, n) 输出
"""
GEMM 的分块策略
为什么要分块？因为 Cube 单元的寄存器有限，一次算不完整个矩阵。
原始矩阵乘法：
A (m×k) × B (k×n) = C (m×n)
┌────────┐ ┌────────┐
│ │ │ │
│ A │ × │ B │
│ │ │ │
└────────┘ └────────┘
分块后：
A 被切成 m/M 个 tile（每块 M×k）
B 被切成 k/K 个 tile（每块 k×n）
C = Σ (A_tile_i × B_tile_i)
每个 tile 刚好放得进 Cube 单元的缓存
ops-blas 的缓存优化
# ops-blas GEMM 的 L1 Cache 优化策略
class GEMMTiler:
"""
把大矩阵切成小 tile，充分利用 L1 Cache
L1 Cache 大小：昇腾 910B 是 64KB
"""
def __init__(self, m, n, k, l1_size=64*1024):
# 计算最优分块大小
# 每个 tile 需要存 A 的 M×K 和 B 的 K×N
# 2 × M × K × sizeof(half) ≤ 64KB
self.block_m = min(512, m) # 输出 tile 大小
self.block_k = min(256, k) # 中间维度
self.block_n = min(512, n)
def tile(self, A, B, C):
"""分块执行"""
for m_start in range(0, A.shape[0], self.block_m):
for n_start in range(0, B.shape[1], self.block_n):
# 加载当前 tile 到 L1
A_tile = A[m_start : m_start+self.block_m, :]
B_tile = B[:, n_start : n_start+self.block_n]
# Cube 计算
C_tile = self.cube_matmul(A_tile, B_tile)
# 累加到 C
C[m_start : m_start+self.block_m,
n_start : n_start+self.block_n] += C_tile
第五层：catlass 模板
catlass 是什么
catlass 是"昇腾算子模板库"，它给 ops-blas 提供 GEMM 的底层实现。
# catlass GEMM 模板的参数化接口
# 来自 catlass/gemm/template.h
class GemmTemplate {
// 矩阵形状
uint32_t m; // A 的行数，C 的行数
uint32_t n; // B 的列数，C 的列数
uint32_t k; // A 的列数，B 的行数
// 数据类型
DataType dtype_a; // A 的数据类型
DataType dtype_b; // B 的数据类型
DataType dtype_c; // C 的数据类型
// 分块参数
uint32_t block_m; // M 方向分块
uint32_t block_n; // N 方向分块
uint32_t block_k; // K 方向分块
// 硬件配置
CubeUnit cube_unit;
VectorUnit vector_unit;
};
catlass 生成的是什么
catlass 模板实例化后，生成的是昇腾达芬奇 NPU 的汇编代码：
┌──────────────────────────────────────┐
│ catlass 模板 │
│ ↓ 实例化 │
│ 生成的汇编（Vector 指令 + Cube 指令）│
│ ↓ 编译 │
│ NPU 可执行 kernel │
└──────────────────────────────────────┘
完整调用链代码
# 完整调用链示例
import torch
import torch_npu
# 1. PyTorch 调用（框架层）
x = torch.randn(16, 512).npu() # 移到 NPU
w = torch.randn(2048, 512).npu() # 移到 NPU
# 2. torch_npu 拦截，调用 AscendCL（API 层）
# torch_npu 内部：aclnnMatMul(x, w, &output)
# 3. AscendCL 调用 ops-nn MatMul（算子层）
# ops-nn 内部：matmul = ops.blas.gemm()
# ops-nn 内部：output = matmul(x, w)
# 4. ops-nn 调用 ops-blas GEMM（BLAS 层）
# ops-blas 内部：tiler = GEMMTiler(m=16, k=512, n=2048)
# ops-blas 内部：result = tiler.tile()
# 5. ops-blas 调用 catlass 模板（模板层）
# catlass 内部：生成汇编，提交 Cube 指令
# 6. 达芬奇 NPU 执行（CUBE 单元）
# 最终计算在硬件上完成
output = torch_npu.matmul(x, w.t()) # 最终输出
ops-nn MatMul vs ops-blas GEMM：什么时候用哪个
| 场景 | 推荐 | 理由 |
|---|---|---|
| Transformer 前向推理 | ops-nn MatMul | 融合 Bias+Activation，省一次数据搬运 |
| BERT / GPT 推理 | ops-nn MatMul | 多层堆叠，融合节省明显 |
| 裸矩阵乘法（无激活） | ops-blas GEMM | 跳过融合开销，更直接 |
| 自定义融合模式 | ops-blas + catlass | 从底层自己拼 |
| 性能 profiling | 两层都要看 | 看 ops-nn 的融合是否生效 |
附录：ops-nn vs ops-blas vs catlass 的关系图
ops-nn
├── MatMul（融合版）
│ ├── 内部调用 ops-blas
│ └── 融合 Bias + Activation
├── Conv2d
├── LayerNorm
└── Softmax
ops-blas
├── GEMM（基础版）
│ ├── 内部调用 catlass
│ └── 分块 + 缓存优化
├── Batch GEMM
└── Strided GEMM
catlass
├── GemmTemplate
├── ConvTemplate
└── AttentionTemplate
选择建议：
-
只需要矩阵乘 + 激活 → ops-nn
-
需要裸 GEMM → ops-blas
-
需要深度定制 → catlass
常见问题 FAQ
Q1: ops-nn 的融合会自动开启吗？
是的，默认开启。但只有 MatMul+Bias+Activation 连续调用时才会融合。
Q2: 融合后精度掉了？
检查是否开启了错误的融合模式。用精度对比工具验证。
Q3: 为什么有时候 ops-blas 比 ops-nn 更快？
融合有 overhead。如果只有 MatMul 没有后续操作，ops-nn 的融合反而浪费。
性能调优实践
profiling 工具
import cann.profiler as profiler
# profiling MatMul 调用
with profiler.Profile("matmul_profile.pb") as p:
for _ in range(1000):
result = matmul_op(x, w)
report = p.get_report()
# 输出：
# - MatMul kernel 耗时
# - BiasAdd kernel 耗时
# - Fusion 融合效果
# - L1/L2 Cache 命中率
# 如果 Cache 命中率 < 80%，ops-blas 分块参数需要调
分块参数调优
# catlass 手动分块（高级用法）
tiler = cann.ops.blas.GEMMTiler(
m=512, k=512, n=2048,
block_m=64, # M 方向分块
block_k=128, # K 方向分块
block_n=64 # N 方向分块
)
# L1 Cache 命中率测试
for bm in [32, 64, 128, 256]:
for bk in [64, 128, 256]:
for bn in [32, 64, 128]:
tiler.block_m = bm
tiler.block_k = bk
tiler.block_n = bn
# 测试性能
elapsed = benchmark(tiler)
cache_hit = tiler.get_l1_hit_rate()
if cache_hit > 0.85 and elapsed < best_time:
best = (bm, bk, bn)
精度 vs 性能权衡
# 融合后精度对比
def check_fusion_accuracy():
# FP32 基准
result_fp32 = model_fp32(input)
# FP16 + 融合
result_fp16 = model_fp16_fused(input)
# 计算相对误差
diff = torch.abs(result_fp16 - result_fp32) / (torch.abs(result_fp32) + 1e-8)
max_rel_err = diff.max()
print(f"Max relative error: {max_rel_err:.6f}")
# 误差阈值：< 1% 通常可接受
if max_rel_err > 0.01:
print("⚠️ Fusion accuracy degradation detected")
else:
print("✅ Fusion accuracy OK")
总结
从 torch.matmul 到 NPU 执行：
- PyTorch → torch_npu 拦截
- AscendCL → 参数校验、API 分发
- ops-nn → 融合 MatMul+Bias+Act
- ops-blas → GEMM 分块、缓存优化
- catlass → 生成 Cube 指令
- 达芬奇 NPU → 硬件执行
理解调用链的价值：遇到性能问题，你知道该优化哪一层。融合不够去 ops-nn，缓存不命中去 ops-blas，指令生成有瓶颈去 catlass。
仓库地址：https://atomgit.com/cann/ops-nn
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。
更多推荐
0
0- 0
已为社区贡献19条内容
所有评论(0)