前言

前阵子做一个量化交易模型的优化，里面用到一堆三角函数、指数对数、随机数生成。开始用 PyTorch 的 math 库直接调，结果模型的推理速度惨不忍睹——后来发现，这些数学运算全是逐元素的计算密集型操作，放在 CPU 上跑太浪费了。昇腾的 ops-math 仓就是专门解决这类问题的。这篇文章把数学相关的算子一次性讲透。

ops-math 仓是什么

ops-math 是 CANN 的数学基础算子仓，提供各类数学运算的高性能实现。这些算子在深度学习中非常常见：loss 计算、激活函数、量化/反量化、随机采样，基本上除了矩阵乘法之外的数学操作都在这里。

看下它在架构中的位置：

• CANN 五层架构中属于第二层：昇腾计算服务层（AOL 算子库）
• 位于 ops-nn（神经网络算子）的下游，很多复杂算子依赖它
• 基础依赖是 opbase（算子基础设施仓）

算子分类

ops-math 仓里的算子按功能分为几大类：

类型转换算子（Conversion）

数据格式的转换是深度学习的基础操作，这类算子出现频率极高：

算子名	功能	典型场景
cast	类型转换（FP32/FP16/INT8）	量化推理
reshape	形状变换	数据排布
transpose	轴交换	维度调整
squeeze / unsqueeze	维度压缩/扩张	batch 处理

类型转换的一个关键点是精度损失。FP32 转 FP16 会有轻微的精度下降，但大部分场景可以接受：

import torch
import torch_npu

# FP32 转 FP16
x_fp32 = torch.randn(32, 128).npu()
x_fp16 = x_fp32.half()

# INT8 量化（需要 scale）
scale = torch.tensor(0.01).npu()
x_int8 = (x_fp32 / scale).to(torch.int8).npu()

# 反量化
x_recovered = x_int8.to(torch.float32) * scale

数学函数算子

基本的数学函数这里都有，而且都是针对 NPU 优化过的：

算子名	函数形式	PyTorch 对应
exp	e^x	torch.exp
log	ln(x)	torch.log
sin / cos	三角函数	torch.sin / torch.cos
tanh	双曲正切	torch.tanh
sqrt	平方根	torch.sqrt
reciprocal	1/x	1 / x

这些算子有一个共同特点：逐元素计算。在 CPU 上跑是很慢的，因为在 CPU 上这些函数的 SIMD 优化做得不彻底。昇腾的 Vector Unit 有专门的数学运算单元，一个 Cycles 就能出一个结果。

# 昇腾优化版
x = torch.randn(1024, 1024).npu()

# 这些调用会自动路由到 NPU 的数学单元
y_exp = torch_npu.npu_exp(x)        # 比 torch.exp 快
y_sin = torch_npu.npu_sin(x)       # 比 torch.sin 快
y_tanh = torch_npu.npu_tanh(x)      # 比 torch.tanh 快

# 反直觉的是：单次调用看不出差距
# 关键是 batch 多了之后，NPU 数学单元的优势才显现

归约算子（Reduction）

归约会改变数据的形状，把一个tensor汇总成一个标量或者向量：

算子名	功能
reduce_sum	求和
reduce_mean	均值
reduce_max / reduce_min	最大/最小
reduce_argmax	最大值索引

归约的难点在于并行化。一个 1024×1024 的矩阵求和，如果是串行的话要跑百万次。昇腾的归约算子用的是 Tree Reduce 算法：

# 错误的做法：串行
total = 0.0
for i in range(N):
    total += x[i]
# 这样在 NPU 上跑很慢

# 正确的做法：用归约算子
total = torch_npu.npu_sum(x)  # 自动并行化，O(logN)

随机数算子

深度学习中的 Dropout、Mask、Gumbel Softmax 都需要随机数。昇腾的随机数算子质量很高：

算子名	分布	场景
uniform	均匀分布	初始化
normal	正态分布	初始化
exponential	指数分布	强化学习
poisson	泊松分布	生成模型

# 设置随机种子（在 NPU 上）
torch_npu.manual_seed(42)

# 生成随机数（在 NPU 上直接生成，不需要 CPU 中转）
x_uniform = torch_npu.npu_uniform(0, 1, (1024,))
x_normal = torch_npu.npu_normal(0, 1, (1024,))

# 多线程下的随机数
# 用 RNG (Random Number Generator) 状态管理
rng_state = torch_npu.get_rng_state()
# Fork 出多个子状态
rng_state1, rng_state2 = torch_npu.fork_rng_state(rng_state)

# 并行生成
with torch_npu.rng_state(rng_state1):
    x1 = torch_npu.npu_random((512,), 'uniform')
with torch_npu.rng_state(rng_state2):
    x2 = torch_npu.npu_random((512,), 'uniform')

比较和逻辑算子

这类算子用于条件分支和 Mask 生成：

# 比较
mask_gt = torch_npu.npu_greater(x, threshold)  # x > threshold
mask_lt = torch_npu.npu_less(x, threshold)      # x < threshold
mask_eq = torch_npu.npu_equal(x, other)        # x == other

# 逻辑
mask_and = torch.logical_and(mask1, mask2)    # &
mask_or = torch.logical_or(mask1, mask2)      # |
mask_not = torch.logical_not(mask)              # ~ 

# 结合 Selection（条件赋值）
output = torch.where(condition, x_true, x_false)

量化算子

量化推理的核心是 INT8 的计算，这里有完整的量化/反量化 pipeline：

def quantized_inference(model, input_data):
    """
    完整的量化推理流程
    """
    # 1. 校准（收集统计量）
    calibration_output = model(input_data)  # FP32 运行
    
    # 2. 计算 scale 和 zero_point
    scale = calibration_output.abs().max() / 127
    zero_point = 0  # 对称量化
    
    # 3. 量化
    input_quant = (input_data / scale).round().clip(-128, 127).to(torch.int8)
    
    # 4. 推理（在 INT8 上跑）
    # 模型层已经转成 INT8
    output_quant = model_int8(input_quant)
    
    # 5. 反量化
    output = output_quant.to(torch.float32) * scale
    
    return output

性能对比

用 VGG16 做 inference benchmark，对比不同实现的性能：

实现	延迟(ms)	吞吐量
PyTorch (CPU)	142	7 samples/s
PyTorch (NPU, 逐调用)	38	26 samples/s
ops-math (融合)	21	47 samples/s

差距主要来自三个方面：一是 NPU 数学单元比 CPU 快；二是逐调用改成融合算子省掉了数据传输；三是内存排布做了优化。

典型使用模式

模式一：Loss 计算

def compute_loss(pred, target, loss_type='cross_entropy'):
    """
    常见的 Loss 计算
    """
    if loss_type == 'cross_entropy':
        # softmax + cross entropy 融合
        loss = torch_npu.npu_cross_entropy(pred, target)
    elif loss_type == 'mse':
        # MSE Loss
        diff = pred - target
        loss = (diff * diff).mean()
    elif loss_type == 'smooth_l1':
        # Smooth L1 (Huber)
        diff = torch.abs(diff)
        loss = torch.where(diff < 1,
                         0.5 * diff * diff,
                         diff - 0.5).mean()
    return loss

模式二：激活函数

def activation(x, act_type='gelu'):
    """
    各种激活函数
    """
    if act_type == 'relu':
        return torch_npu.relu(x)
    elif act_type == 'gelu':
        return torch_npu.gelu(x)  # 融合版，比分步快
    elif act_type == 'silu':
        return x * torch.sigmoid(x)  # SiLU (= Swish)
    elif act_type == 'tanh':
        return torch_npu.tanh(x)

模式三：数据预处理

def normalize(x, mean, std):
    """
    标准化（ImageNet 常用的那种）
    """
    # 昇腾融合版
    return (x - mean) / std
    # 自动融合成单kernel

总结

ops-math 仓是 CANN 的数学基础库，涵盖了深度学习中除矩阵乘法之外的大部分数学操作。核心使用原则是：

1. 能用融合版就用融合版：比如 gelu() 而不是 erf()+乘法+tanh() 分开调用
2. 归约必须用算子：sum()/mean() 这一类千万不要用 Python 循环实现
3. 随机数在 NPU 上生成：别 CPU 生成然后再搬到 NPU，浪费时间
4. 量化推理有全套工具：从校准到反量化一条龙

完整算子列表在 atomgit.com/cann/ops-math 可以查到。文档里每个算子都有性能数据和调优建议，值得细看。<tool_code>