之前带实习生,他问我:“哥,我看代码里调用 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 4096x4096450ms45ms10x
GEMV 4096x115ms2ms7.5x
DOT 40960.5ms0.1ms5x
AXPY 40960.3ms0.05ms6x
NRM2 40960.4ms0.08ms5x
TRMM 2048x2048120ms15ms8x
TRSM 2048x2048150ms18ms8.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开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐