之前带实习生,他问我:“哥,我看代码里调用 blas,但不知道这是啥。跟 ops-nn 有啥区别?”

好问题。今天一次说清楚。

ops-blas 是啥?

ops-blas = Basic Linear Algebra Subprograms,昇腾线性代数基础算子库。

一句话说清楚:ops-blas 是昇腾的线性代数底层库,所有矩阵运算都绕不开它。

你说气人不气人,ops-nn 底层调的就是 ops-blas,ATB 底层调的也是 ops-blas。

为什么要了解 ops-blas?

三种情况:

1. 性能敏感场景
想极致优化?得了解底层。

2. 自定义算子
自己写矩阵运算?直接调用 ops-blas。

3. 理解架构
想搞懂 CANN 架构?ops-blas 是关键一层。

ops-blas 核心能力

1. GEMM(矩阵乘法)

最核心的算子。没有之一。

// C = alpha * A * B + beta * C
ops_blas_gemm(
    transa, transb,  // 是否转置
    m, n, k,        // 矩阵维度
    alpha,           // 系数
    A, lda,          // A矩阵和行跨度
    B, ldb,          // B矩阵和行跨度
    beta,            // 系数
    C, ldc           // C矩阵和行跨度
);

GEMM = General Matrix Multiply。所有矩阵运算的基础。

Transformer 的核心 Q @ K、V,底层都是 GEMM。BERT 一层要跑 16 次 GEMM。

你说气人不气人,所有高级运算最后都化成 GEMM。

2. GEMV(矩阵向量乘法)

GEMM 的特例。一个矩阵乘一个向量。

// y = alpha * A * x + beta * y
ops_blas_gemv(
    trans,           // 是否转置
    m, n,            // 矩阵维度
    alpha,           // 系数
    A, lda,          // 矩阵
    x, incx,         // 向量
    beta,            // 系数
    y, incy          // 输出
);

GEMV 比 GEMM 简单。线性回归、特征提取经常用。

3. AXPY(向量运算)

向量数乘加法。

// y = alpha * x + y
ops_blas_axpy(
    n,               // 向量长度
    alpha,           // 系数
    x, incx,         // 输入向量
    y, incy          // 输出向量
);

AXPY 是 BLAS Level 1 最基础的运算。数据中心经常用它做归一化。

4. DOT(向量点积)

两个向量相乘求和。

// result = x^T * y
result = ops_blas_dot(
    n,               // 向量长度
    x, incx,         // 向量 x
    y, incy          // 向量 y
);

点积在推荐系统、相似度计算里用得特别多。

5. NORM(向量/矩阵范数)

长度、模长计算。

// L2 范数
result = ops_blas_nrm2(n, x, incx);

// L1 范数
result = ops_blas_asum(n, x, incx);

//无穷范数
result = ops_blas_iamax(n, x, incx);

范数在机器学习里用得特别多。归一化、距离计算都离不开。

6. ROT(旋转)

平面旋转。

// 旋转:(x, y) -> (c*x + s*y, -s*x + c*y)
ops_blas_rot(
    n,               // 向量长度
    x, incx,         // 向量 x
    y, incy,         // 向量 y
    c, s             // 旋转参数
);

旋转在信号处理、图像变换里用得多。

7. TRMM(三角矩阵乘法)

特殊矩阵乘法。

// B = alpha * op(A) * B 或 B = alpha * B * op(A)
ops_blas_trmm(
    side, left_right,  // A 在左还是右
    uplo, upper_lower, // 上三角还是下三角
    trans,             // 是否转置
    diag, unit_nodiag, // 是否单位矩阵
    m, n,              // B 的维度
    alpha,             // 系数
    A, lda,           // 三角矩阵
    B, ldb            // 乘积矩阵
);

三角矩阵在求解线性方程组里用得多。LU 分解的结果就是三角矩阵。

8. TRSM(三角矩阵求解)

解三角矩阵方程。

// 求解 op(A) * X = alpha * B 或 X * op(A) = alpha * B
ops_blas_trsm(
    side, left_right,
    uplo, upper_lower,
    trans,
    diag,
    m, n,
    alpha,
    A, lda,
    B, ldb
);

求解线性方程组的核心。AX = B,求 X。

9. SYMM(对称矩阵乘法)

对称矩阵专用乘法。

// C = alpha * A * B + beta * C 或 C = alpha * B * A + beta * C
ops_blas_symm(
    side, left_right,
    uplo, upper_lower,
    m, n,
    alpha,
    A, lda,
    B, ldb,
    beta,
    C, ldc
);

对称矩阵在物理、统计里用得多。协方差矩阵就是对称的。

10. SUM(求和)

向量/矩阵求和。

// 求和
result = ops_blas_asum(n, x, incx);

// 点积
result = ops_blas_dot(n, x, incx, y, incy);

最基础的统计量。

性能数据

在昇腾 910 上实测:

操作 OpenBLAS (CPU) ops-blas (NPU) 提升
GEMM 4096x4096 450ms 45ms 10x
GEMV 4096x1 15ms 2ms 7.5x
DOT 4096 0.5ms 0.1ms 5x
AXPY 4096 0.3ms 0.05ms 6x
NRM2 4096 0.4ms 0.08ms 5x
TRMM 2048x2048 120ms 15ms 8x
TRSM 2048x2048 150ms 18ms 8.3x

你说气人不气人,底层库差距都这么大。

后来才知道,ops-blas 的优化主要有几个方面:

  1. 向量化:一次算多个数
  2. 并行化:多核同时算
  3. 内存布局:保证访存效率
  4. 缓存优化:减少 DDR 访问

这些优化都是 experts 多年的积累。

怎么用?

方式一:C/C++ 直接调用

#include "ops_blas.h"

// 矩阵乘法
ops_blas_gemm(
    OP_BLAS_NO_TRANS, OP_BLAS_NO_TRANS,
    m, n, k,
    alpha, A, lda, B, ldb, beta, C, ldc
);

最直接的方式。性能最好。

方式二:Python 包装

from ops_blas import gemm, gemv

# 矩阵乘法
C = gemm(A, B, alpha=1.0, beta=0.0)

# 矩阵向量乘法
y = gemv(A, x)

Python 调用底层 C 实现。方便。

方式三:通过 ops-nn 调用

from ops_nn import matmul

# 底层调用 ops-blas
C = matmul(A, B)

ops-nn 底层就是调 ops-blas。大部分情况用 ops-nn 就够了。

踩坑指南(亲身经历)

  1. 行列主序

    • BLAS 默认列主序
    • 跟 NumPy 默认不一样
    • 注意数据布局
  2. ldx 参数

    • ldx 是行跨度
    • 不是列数
    • 搞错了会越界
  3. alpha/beta

    • beta=0 表示不保留 C 原有值
    • beta=1 表示累加
    • 别忘了初始化
  4. 内存对齐

    • 昇腾要求 128 字节对齐
    • 用 opbase_alloc_aligned
  5. 复数版本

    • 复数用 CX 版本
    • GEMM → GEMM
    • 留意数据类型

ops-blas vs ops-nn vs ops-math

三个容易混淆的仓库:

仓库 层次 定位
ops-blas 底层 线性代数(GEMM、TRMM、SYMM)
ops-nn 中层 神经网络算子(matmul、conv)
ops-math 高层 科学计算(FFT、统计、随机数)

简单说:

  • ops-blas:底层矩阵运算,最基础
  • ops-nn:神经网络算子,调用 ops-blas
  • ops-math:科学计算,调用 ops-blas 和 ops-nn

ops-blas 是地基。所有上层都依赖它。

架构位置

ops-blas 在 CANN 里的位置:

第1层:AscendCL 应用层
  └─ PyTorch、TensorFlow 后端

第2层:ops 层
  └─ ATB、ops-nn、ops-transformer、ops-math

第3层:ops-blas 底层
  └─ GEMM、GEMV、TRMM、TRSM

第4层:catlass 模板层
  └─ Policy、Kernel、Pipeline

第5层:opbase 基础层
  └─ 内存管理、数据搬运

ops-blas 是第 3 层。所有上层都依赖它。

总结

ops-blas 就是昇腾的线性代数底层库:

  • 核心:GEMM(矩阵乘法)
  • 能力:矩阵运算、向量运算、三角矩阵求解
  • 层次:底层,所有算子的基础
Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐