muBLAS 开发者指南

muBLAS 是基于 MUSA 开发的基础线性代数运算库，在 MTGPU 上经过深度优化，在 AI 和 HPC 场景下被广泛使用。

概述

什么是 muBLAS

muBLAS（MUSA Basic Linear Algebra Subprograms，基础线性代数算子库）提供 BLAS 标准风格的向量、矩阵与矩阵乘法接口，开发者可以通过统一的 handle、stream 和设备内存模型调用 MTGPU 上优化后的线性代数算子。

关键特性

• 覆盖 BLAS Level 1、Level 2 和 Level 3 常用计算。
• 支持单精度、双精度、复数等基础数据类型。
• 可与 MUSA stream 配合使用，支持异步执行和多 stream 并发。
• 针对 MTGPU 架构优化，适用于 AI、HPC 和科学计算场景。

功能分类

按照计算复杂性，muBLAS 函数可分为三类：

类别	说明	示例函数
第一类	标量、向量和向量与向量间的运算	axpy, dot, nrm2, scal
第二类	向量与矩阵之间的运算	gemv, ger
第三类	矩阵与矩阵间的运算	gemm, trsm

数据类型支持

muBLAS 支持以下数据类型：

类型前缀	数据类型
s	单精度浮点 (float)
d	双精度浮点 (double)
c	单精度复数 (complex float)
z	双精度复数 (complex double)

快速开始

完整示例：SAXPY

SAXPY: $y = \alpha x + y$ (单精度 AX 加 Y)

#include <cstdio>
#include <cstdlib>
#include <musa_runtime.h>
#include <mublas.h>
#include <vector>

int main(int argc, char* argv[]) {
    mublasHandle_t mublasH = NULL;
    musaStream_t   stream  = NULL;

    // 初始化数据
    // A = [1.0, 2.0, 3.0, 4.0]
    // B = [5.0, 6.0, 7.0, 8.0]
    const std::vector<float> A = {1.0, 2.0, 3.0, 4.0};
    std::vector<float>       B = {5.0, 6.0, 7.0, 8.0};
    const float              alpha = 2.1;
    const int                incx = 1;
    const int                incy = 1;

    float* d_A = nullptr;
    float* d_B = nullptr;

    // Step 1: 创建 mublas handle，绑定 stream
    mublasCreate(&mublasH);
    musaStreamCreateWithFlags(&stream, musaStreamNonBlocking);
    mublasSetStream(mublasH, stream);

    // Step 2: 复制数据到设备
    musaMalloc(reinterpret_cast<void**>(&d_A), sizeof(float) * A.size());
    musaMalloc(reinterpret_cast<void**>(&d_B), sizeof(float) * B.size());

    musaMemcpyAsync(d_A, A.data(), sizeof(float) * A.size(), 
                    musaMemcpyHostToDevice, stream);
    musaMemcpyAsync(d_B, B.data(), sizeof(float) * B.size(), 
                    musaMemcpyHostToDevice, stream);

    // Step 3: 计算 B = alpha * A + B
    // B = 2.1 * [1, 2, 3, 4] + [5, 6, 7, 8]
    // B = [7.1, 10.2, 13.3, 16.4]
    mublasSaxpy(mublasH, A.size(), &alpha, d_A, incx, d_B, incy);

    // Step 4: 复制数据回主机
    musaMemcpyAsync(B.data(), d_B, sizeof(float) * B.size(), 
                    musaMemcpyDeviceToHost, stream);
    musaStreamSynchronize(stream);

    // 输出结果
    printf("B = [");
    for(int i = 0; i < B.size(); i++) {
        printf("%f ", B[i]);
    }
    printf("]\n");
    // 预期输出：B = [7.100000 10.200000 13.300000 16.400000]

    // 释放资源
    musaFree(d_A);
    musaFree(d_B);
    mublasDestroy(mublasH);
    musaStreamDestroy(stream);
    musaDeviceReset();

    return EXIT_SUCCESS;
}

编译和运行

# 使用 mcc 编译
mcc saxpy_example.cpp \
    -L/usr/local/musa/lib \
    -I/usr/local/musa/include \
    -lmublas -lmusart \
    -o saxpy_example

# 运行
./saxpy_example

或使用 CMake:

cmake_minimum_required(VERSION 3.10)
project(SAXPY_Example LANGUAGES CXX)
add_executable(saxpy_example saxpy_example.cpp)
target_include_directories(saxpy_example PRIVATE /usr/local/musa/include)
target_link_directories(saxpy_example PRIVATE /usr/local/musa/lib)
target_link_libraries(saxpy_example PRIVATE mublas musart)
set_property(TARGET saxpy_example PROPERTY CXX_STANDARD 17)

编译步骤：

# 1. 创建构建目录并进入
mkdir build && cd build

# 2. 配置 CMake（需在包含 CMakeLists.txt 的目录下执行）
cmake ..

# 3. 编译
cmake --build .
# 或使用 make / ninja，与项目约定一致即可

注意：若 MUSA 未安装在默认路径 /usr/local/musa，需将 target_include_directories 和 target_link_directories 中的路径改为实际安装路径。

---

API 参考

向量与向量运算 (Level 1)

axpy - 向量线性组合

mublasStatus mublasSaxpy(mublasHandle_t handle,
                         int n,
                         const float* alpha,
                         const float* x,
                         int incx,
                         float* y,
                         int incy);

功能: $y = \alpha x + y$

参数:

参数	说明
handle	muBLAS 句柄
n	向量元素数量
alpha	标量乘数
x	输入向量 x
incx	x 的步长
y	输入/输出向量 y
incy	y 的步长

dot - 向量点积

mublasStatus mublasSdot(mublasHandle_t handle,
                        int n,
                        const float* x,
                        int incx,
                        const float* y,
                        int incy,
                        float* result);

功能: $result = x \cdot y = \sum_i x_i y_i$

nrm2 - 向量 L2 范数

mublasStatus mublasSnrm2(mublasHandle_t handle,
                         int n,
                         const float* x,
                         int incx,
                         float* result);

功能: $result = \lVert x \rVert_2 = \sqrt{\sum_i \lvert x_i \rvert^2}$

scal - 向量缩放

mublasStatus mublasSscal(mublasHandle_t handle,
                         int n,
                         const float* alpha,
                         float* x,
                         int incx);

功能: $x = \alpha x$

矩阵与向量运算 (Level 2)

gemv - 矩阵 - 向量乘法

mublasStatus mublasSgemv(mublasHandle_t handle,
                         mublasOperation_t trans,
                         int m,
                         int n,
                         const float* alpha,
                         const float* A,
                         int lda,
                         const float* x,
                         int incx,
                         const float* beta,
                         float* y,
                         int incy);

功能: $y = \alpha \operatorname{op}(A)x + \beta y$

其中 op(A) 可以是：

• MUBLAS_OP_N: A (不转置)
• MUBLAS_OP_T: A^T (转置)
• MUBLAS_OP_C: A^H (共轭转置)

说明

共轭转置仅适用于 c 和 z 复数精度接口。对于 s 和 d 实数精度接口，通常使用 MUBLAS_OP_N 或 MUBLAS_OP_T。

矩阵与矩阵运算 (Level 3)

gemm - 矩阵 - 矩阵乘法

mublasStatus mublasSgemm(mublasHandle_t handle,
                         mublasOperation_t transA,
                         mublasOperation_t transB,
                         int m,
                         int n,
                         int k,
                         const float* alpha,
                         const float* A,
                         int lda,
                         const float* B,
                         int ldb,
                         const float* beta,
                         float* C,
                         int ldc);

功能: $C = \alpha \operatorname{op}(A)\operatorname{op}(B) + \beta C$

参数:

参数	说明
transA, transB	A 和 B 的转置选项
m	op(A) 的行数，C 的行数
n	op(B) 的列数，C 的列数
k	op(A) 的列数，op(B) 的行数
alpha	标量乘数
A, B	输入矩阵
lda, ldb	A 和 B 的前导维度
beta	标量乘数用于 C
C	输出矩阵
ldc	C 的前导维度

---

使用模式

多 Stream 并发

mublasHandle_t handle1, handle2;
musaStream_t stream1, stream2;

// 创建多个 handle 和 stream
mublasCreate(&handle1);
mublasCreate(&handle2);
musaStreamCreate(&stream1);
musaStreamCreate(&stream2);

mublasSetStream(handle1, stream1);
mublasSetStream(handle2, stream2);

// 并发执行多个操作
mublasSaxpy(handle1, n, &alpha, d_x, 1, d_y, 1);
mublasSgemv(handle2, MUBLAS_OP_N, m, n, &alpha, d_A, lda, d_x, 1, &beta, d_z, 1);

// 等待完成
musaStreamSynchronize(stream1);
musaStreamSynchronize(stream2);

mublasDestroy(handle1);
mublasDestroy(handle2);

批处理操作

// 使用连续存储的批处理矩阵乘法
int batchSize = 32;
long long int strideA = static_cast<long long int>(lda) * k;
long long int strideB = static_cast<long long int>(ldb) * n;
long long int strideC = static_cast<long long int>(ldc) * n;

mublasSgemmStridedBatched(handle,
                          MUBLAS_OP_N, MUBLAS_OP_N,
                          m, n, k,
                          &alpha,
                          d_A, lda, strideA,
                          d_B, ldb, strideB,
                          &beta,
                          d_C, ldc, strideC,
                          batchSize);

如果每个批次的矩阵地址不连续，可使用指针数组形式的批处理接口：

int batchSize = 32;
const float** d_A_array = nullptr;
const float** d_B_array = nullptr;
float** d_C_array = nullptr;

// d_A_array、d_B_array 和 d_C_array 为设备端指针数组，
// 数组元素分别指向每个批次的设备端矩阵。

mublasSgemmBatched(handle,
                   MUBLAS_OP_N, MUBLAS_OP_N,
                   m, n, k,
                   &alpha,
                   d_A_array, lda,
                   d_B_array, ldb,
                   &beta,
                   d_C_array, ldc,
                   batchSize);

---

错误处理

muBLAS 接口返回 mublasStatus，应用程序应检查返回值并根据错误类型定位问题。

状态码	说明
MUBLAS_STATUS_SUCCESS	操作执行成功。
MUBLAS_STATUS_NOT_INITIALIZED	muBLAS handle 未初始化、无效或为空。
MUBLAS_STATUS_ALLOC_FAILED	内部内存分配、复制或释放失败。
MUBLAS_STATUS_INVALID_VALUE	传入的参数无效。
MUBLAS_STATUS_ARCH_MISMATCH	当前设备架构缺少该函数所需的能力。
MUBLAS_STATUS_MAPPING_ERROR	访问 GPU 内存空间失败，通常与资源绑定或内存映射异常有关。
MUBLAS_STATUS_EXECUTION_FAILED	GPU 程序执行失败，常见原因包括 kernel 启动失败或运行时错误。
MUBLAS_STATUS_INTERNAL_ERROR	muBLAS 库内部错误。
MUBLAS_STATUS_NOT_SUPPORTED	当前函数或参数组合暂不支持。
MUBLAS_STATUS_LICENSE_ERROR	许可证检查失败，当前操作无法继续执行。

---

完整示例：矩阵乘法

#include <cstdio>
#include <cstdlib>
#include <musa_runtime.h>
#include <mublas.h>
#include <vector>

int main() {
    const int M = 256, N = 256, K = 256;
    
    // 初始化主机数据
    std::vector<float> h_A(M * K), h_B(K * N), h_C(M * N);
    for (int i = 0; i < M * K; i++) h_A[i] = 1.0f;
    for (int i = 0; i < K * N; i++) h_B[i] = 2.0f;
    
    // 创建 handle
    mublasHandle_t handle;
    mublasCreate(&handle);
    
    // 分配设备内存
    float *d_A, *d_B, *d_C;
    musaMalloc(&d_A, M * K * sizeof(float));
    musaMalloc(&d_B, K * N * sizeof(float));
    musaMalloc(&d_C, M * N * sizeof(float));
    
    // 复制数据到设备
    musaMemcpy(d_A, h_A.data(), M * K * sizeof(float), musaMemcpyHostToDevice);
    musaMemcpy(d_B, h_B.data(), K * N * sizeof(float), musaMemcpyHostToDevice);
    
    // 矩阵乘法参数
    const float alpha = 1.0f;
    const float beta = 0.0f;
    
    // C = A * B
    mublasSgemm(handle,
                MUBLAS_OP_N, MUBLAS_OP_N,
                M, N, K,
                &alpha,
                d_A, M,
                d_B, K,
                &beta,
                d_C, M);
    
    // 复制结果回主机
    musaMemcpy(h_C.data(), d_C, M * N * sizeof(float), musaMemcpyDeviceToHost);
    
    // 验证结果 (预期所有元素都是 2*K = 512)
    printf("C[0] = %f (expected: %f)\n", h_C[0], K * 2.0f);
    
    // 释放资源
    musaFree(d_A);
    musaFree(d_B);
    musaFree(d_C);
    mublasDestroy(handle);
    
    return 0;
}

编译和运行：编译与链接方式与上文 SAXPY 示例相同（同样的 mcc、-lmublas -lmusart、-L/-I 路径）。

# 编译（假设将代码保存为 matrix_mul.cpp）
mcc matrix_mul.cpp \
    -L/usr/local/musa/lib \
    -I/usr/local/musa/include \
    -lmublas -lmusart \
    -o matrix_mul

# 运行
./matrix_mul

---

相关文档

• muBLAS API 参考