##前言

Cast、Abs、Add、Mul 这些数学算子不 起眼，但几乎每个模型都在调。ops-math 仓是 CANN 算子库的基础，位于第二层 AOL 算子库，提供 conversion、math、random 类的 basic 算子。这篇文章把它拆开一组，看看它到底管哪些事。

ops-math 在 CANN 架构中的位置

CANN 是昇腾异构计算架构，分五层。ops-math 在第二层——昇腾计算服务层的 AOL 算子库。AOL 底下是 NN（神经网络）、BLAS（线性代数）、DVPP（数字视觉）、AIPP（AI 预处理）、HCCL（集合通信）等算子库，ops-math 属于 NN 同级的基础算子库。

所有 ops-* 仓库都依赖 opbase 这个基础组件。opbase 提供了算子开发的基础架构，包括统一的内存管理、tiling 策略、kernel 模板等。ops-math自己的定位是给其他算子仓库提供基础的 element-wise 操作。换个角度看：无论是 ops-nn 的 MatMul 还是 ops-blas 的 GEMM，底层都要用到 Cast 来做数据类型转换，用 Add 做结果累加，用 Mul 做缩放。

算子分类

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

第一类是 Cast 类，负责数据类型转换。包括 Cast（FP16 ↔ FP32 ↔ INT8 等）、Floor（向下取整）、Ceil（向上取整）、Round（四舍五入）、Truncate（截断）。数据在 GPU 上跑的时候经常要转换精度——训练用 FP32，推理用 FP16 省钱，量化用 INT8 进一步压缩。

第二类是 Math 类，负责基本的数学运算。包括 Add（element-wise 加）、Mul（element-wise 乘）、Sub（element-wise 减）、Div（element-wise 除）、Abs（绝对值）、Sign（符号函数）、Neg（取负）、Pow（幂运算）、Sqrt（开方）、Exp、Log、Rsqrt（倒数开方）、Square（平方）、 Reciprocal（倒数）等。这一类算是最常用的，量大管饱。

第三类是 Reduce 类，负责归约操作。包括 ReduceSum、ReduceMax、ReduceMin、ReduceMean、ReduceProd、ReduceAll、ReduceAny。这些操作用于在某个维度上做汇聚，比如 Softmax 里的 reduce_sum。

第四类是 Compare 类，负责比较操作。包括 Less、Greater、Equal、NotEqual、LessEqual、GreaterEqual、IsInf、IsNaN、IsFinite。这一类常用于 mask 生成，比如 attention 里的 causal mask。

第五类是 Logical 类，负责布尔操作。包括 LogicalAnd、LogicalOr、LogicalNot、LogicalXor。

第六类是 Random 类，负责随机数生成。包括 Uniform（均匀分布）、Normal（正态分布）、Poisson（泊松分布）、Multinomial（多项分布）。这在大模型里用于 Dropout 的随机 mask 生成。

第七类是 Tensor 类，负责张量形状操作。包括 Reshape、Transpose、Slice、Gather、Scatter、Concat、Split、Tile、Repeat。从名字就能看出来这一类跟张量组织有关，不太像传统意义上的“数学”算子，但确实划在 ops-math 里。

关键算子代码示例

看几个关键算子的 Ascend C 实现，对理解 CANN 编程有帮助。先看最简单的 Cast 算子，做数据类型转换：

// Cast 算子：FP32 转 FP16
// 这个算子看起来简单，但背后有精度损失的风险
extern "C" __global__ __aicore__ void cast_fp32_to_fp16(
    GM_ADDR input, GM_ADDR output, int64_t size)
{
    TPipe pipe;
    TQue<QuePosition::VECIN, 1> in_q;
    TQue<QuePosition::VECOUT, 1> out_q;
    // 单 buffer 够用了，不搞双缓冲
    pipe.InitBuffer(in_q, size * sizeof(float));
    pipe.InitBuffer(out_q, size * sizeof(half));

    // 从 HBM 搬到 L1
    LocalTensor<float> in_local = in_q.AllocTensor<float>();
    DataCopy(in_local, input, size * sizeof(float));

    // Vector 单元做类型转换
    // 每个 lane 处理一个元素，并行度很高
    LocalTensor<half> out_local = out_q.AllocTensor<half>();
    // 向量转型的核心是精度处理
    // 这里的 trunc 和 round 策略会影响精度
    // 默认用 round，数值更稳定
    vec_cast_half_to_float(out_local, in_local, size, 1);

    // 写回 HBM
    DataCopy(output, out_local, size * sizeof(half));
}

再看一个稍微复杂的 Add 算子，支持广播的 element-wise 加法：

// Add 算子，支持广播的 element-wise 加法
// 广播的意思是不同 shape 的 tensor 可以做运算
extern "C" __global__ __aicore__ void add_element_wise(
    GM_ADDR a, GM_ADDR b, GM_ADDR o,
    int64_t a_shape, int64_t b_shape, int64_t o_shape)
{
    TPipe pipe;
    // 这个算子在 vector 单元上跑，很适合 element-wise 操作
    // tile 策略取决于数据量和 L1 容量
    int64_t total = o_shape;
    int64_t tile_size = 256 * 1024 / sizeof(half);  // 256KB 的 L1 空间
    
    for (int64_t i = 0; i < total; i += tile_size) {
        int64_t cur = min(tile_size, total - i);
        
        // 两路输入要各自处理广播
        // 如果 a_shape != o_shape，说明 a 要广播
        // 广播策略：把小的维度复制到大的维度上
        LocalTensor<half> a_local, b_local, o_local;
        // ... 广播逻辑省略
        
        // Vector 单元批量做加法
        // 这是一条指令同时做 N 个元素的加法
        vec_add(a_local, b_local, o_local, cur, 1);
    }
}

Reduce 类的算子稍微复杂一些，因为它涉及跨维度的聚合计数。以 ReduceSum 为例：

// ReduceSum 算子：在某个维度上求和
// axis=0 表示按行聚合计数，axis=1 表示按列
extern "C" __global__ __aicore__ void reduce_sum(
    GM_ADDR input, GM_ADDR output,
    int64_t in_h, int64_t in_w, int64_t axis)
{
    TPipe pipe;
    // Reduce 算子通常分两阶段：第一阶段在 Core 内做局部聚合并将结果写回 L1
    // 第二阶段在另一个核上做最终的全局聚合，最后汇总到主核
}

跟 ops-nn、ops-blas 的协同关系

ops-math 是一个基础库，真正的业务逻辑要靠上面的 ops-nn 和 ops-blas 来承接。一个典型的调用链是这样的：

用户代码 (PyTorch / AscendCL)
    ↓
ops-nn.Linear (调用 MatMul + BiasAdd)
    ↓
ops-blas.GEMM (调用底层 Matrix Multiplication)
    ↓
catlass (调用 GEMM 模板)
    ↓
ops-math.ReduceSum / ops-math.Add (底层操作)

举一个实际的例子：你在 PyTorch 里写一个 Linear 层

linear = nn.Linear(4096, 11008, bias=True).npu()
output = linear(input)

这行代码在 CANN 里的执行流程是：

1. AscendCL 收到调用请求，构造 OpKernel
2. 调度到 ops-nn 仓的 MatMul 算子
3. ops-nn.MatMul 底层调用 ops-blas.GEMM
4. ops-blas.GEMM 内部调用 catlass 模板
5. catlass 模板里会有大量的 ops-math 算子调用
   - Cast: 输入数据类型转换
   - Mul: 矩阵乘的 scale
   - Add: 偏置相加
   - ReduceSum: softmax 里的归约

整个链路看下来，ops-math 相当于是地基。上层的算子仓库建在这个地基之上，而用户一般感知不到它的存在。

什么时候直接调 ops-math

大多数情况下你不会直接调 ops-math，而是通过上层的 nn.Linear、nn.Conv 这类高层 API 间接用。但有些场景需要直接调用：

第一种是自定义算子开发。如果你要给昇腾NPU 写一个特殊的算子，底层免不了要调用 ops-math 里的 element-wise 操作。常见做法是用 Ascend C 写 kernel，内部嵌入 ops-math 的等价操作。

第二种是手工计算图优化。有些算子级别的优化框架会探知整个计算图，然后手动把连续的 Cast + Add + Cast 这种 pattern 手��� fuse 成单个算子。这种情况下会直接操作用 ops-math。

第三种是精度调试。当怀疑某个算子有精度问题时，直接调 ops-math 对比 input 和 output，能更快定位问题出在哪个环节。比如你怀疑某个 layer norm 后面的数值有问题，可以拆成单独的 Cast、Add、Reduce 操作逐个排查。

# 直接用 Ascend C API 调用 ops-math 的例子
import torch_npu
from torch_npu.contrib import npu_ops

# 调用 Cast 将 FP32 转为 FP16
input_fp32 = torch.randn(1, 512, dtype=torch.float32).npu()
output_fp16 = npu_ops.cast(input_fp32, npu_ops.DATA_TYPE_HALF)
print(output_fp16.dtype)  # torch.float16

ops-math 这个仓的存在价值就在于它是 CANN 算子库的最小公约数。所有上层算子——不管你是做卷积还是做 attention——底层都免不了要跟这些基本的数学运算打交道。知道它管哪些事儿，关键时刻能帮你快速定位问题出在哪个环节。

仓库地址：https://atomgit.com/cann/ops-math