📚 课程目录第3章 NPU矩阵乘法加速

3. NPU矩阵乘法加速

任务目标

在NPU上实现一个固定尺寸（64×64）的矩阵乘法 C = A × B，并与CPU上的同等实现进行耗时对比，验证NPU加速效果。本案例聚焦于NPU算子调用与数据搬运流程，不涉及自定义算子开发。

环境与依赖

硬件：搭载NPU的AI加速卡（如昇腾310P/910）
软件：CANN Toolkit 6.0+，Python 3.8+
依赖库：numpy（用于CPU对比与数据生成）、acl（CANN Python接口）

关键完整代码

import numpy as np
import time
import acl

# 矩阵尺寸
N = 64

def cpu_matmul(a, b):
    """CPU参考实现"""
    return np.dot(a, b)

def npu_matmul(a, b):
    """NPU实现：使用CANN的MatMul算子"""
    # 初始化
    ret = acl.init()
    assert ret == 0, f"acl.init failed: {ret}"
    ret = acl.rt.set_device(0)
    assert ret == 0, f"set_device failed: {ret}"
    context, ret = acl.rt.create_context(0)
    assert ret == 0, f"create_context failed: {ret}"

    # 准备数据（转为fp16以匹配NPU常用精度）
    a_fp16 = a.astype(np.float16)
    b_fp16 = b.astype(np.float16)

    # 分配NPU设备内存
    a_buffer, ret = acl.rt.malloc(a_fp16.nbytes, acl.const.MEM_MALLOC_NORMAL_ONLY)
    b_buffer, ret = acl.rt.malloc(b_fp16.nbytes, acl.const.MEM_MALLOC_NORMAL_ONLY)
    c_buffer, ret = acl.rt.malloc(N * N * 2, acl.const.MEM_MALLOC_NORMAL_ONLY)  # fp16输出

    # 数据拷贝到设备
    a_np = np.frombuffer(a_buffer, dtype=np.float16).reshape(N, N)
    b_np = np.frombuffer(b_buffer, dtype=np.float16).reshape(N, N)
    a_np[:] = a_fp16
    b_np[:] = b_fp16

    # 创建MatMul算子描述
    matmul_desc = acl.ops.matmul.create_desc()
    acl.ops.matmul.set_desc(matmul_desc, acl.ops.matmul.DESC_INPUT_A, 0, a_fp16.shape, acl.ops.matmul.FORMAT_ND, acl.ops.matmul.DTYPE_FP16)
    acl.ops.matmul.set_desc(matmul_desc, acl.ops.matmul.DESC_INPUT_B, 0, b_fp16.shape, acl.ops.matmul.FORMAT_ND, acl.ops.matmul.DTYPE_FP16)
    acl.ops.matmul.set_desc(matmul_desc, acl.ops.matmul.DESC_OUTPUT_C, 0, (N, N), acl.ops.matmul.FORMAT_ND, acl.ops.matmul.DTYPE_FP16)

    # 执行MatMul
    stream, ret = acl.rt.create_stream()
    ret = acl.ops.matmul.exec(stream, matmul_desc, [a_buffer], [b_buffer], [c_buffer])
    assert ret == 0, f"matmul exec failed: {ret}"
    acl.rt.synchronize_stream(stream)

    # 拷贝回主机
    c_result = np.zeros((N, N), dtype=np.float16)
    c_np = np.frombuffer(c_buffer, dtype=np.float16).reshape(N, N)
    c_result[:] = c_np

    # 清理资源
    acl.rt.free(a_buffer)
    acl.rt.free(b_buffer)
    acl.rt.free(c_buffer)
    acl.ops.matmul.destroy_desc(matmul_desc)
    acl.rt.destroy_stream(stream)
    acl.rt.destroy_context(context)
    acl.rt.reset_device(0)
    acl.finalize()

    return c_result.astype(np.float32)

if __name__ == "__main__":
    # 生成随机输入
    np.random.seed(42)
    A = np.random.rand(N, N).astype(np.float32)
    B = np.random.rand(N, N).astype(np.float32)

    # CPU执行与计时
    t0 = time.time()
    c_cpu = cpu_matmul(A, B)
    t_cpu = time.time() - t0

    # NPU执行与计时（首次调用含初始化，取第二次执行时间）
    npu_matmul(A, B)  # 预热
    t0 = time.time()
    c_npu = npu_matmul(A, B)
    t_npu = time.time() - t0

    # 验证正确性（容差1e-2，因fp16精度）
    max_diff = np.max(np.abs(c_cpu - c_npu))
    print(f"CPU耗时: {t_cpu*1000:.2f} ms")
    print(f"NPU耗时: {t_npu*1000:.2f} ms")
    print(f"最大误差: {max_diff:.4f}")
    print(f"加速比: {t_cpu/t_npu:.2f}x")
    assert max_diff < 1e-2, "NPU结果与CPU差异过大"
    print("验证通过")

运行方式

# 在已配置CANN环境的NPU服务器上执行
python3 npu_matmul_demo.py

验证方法与判定标准

正确性验证：程序自动比较NPU输出与CPU参考结果，最大误差小于0.01即判定计算正确（fp16精度下合理容差）。
性能验证：观察输出的CPU耗时与NPU耗时，加速比大于1即表明NPU对该尺寸矩阵乘法有加速效果。对于64×64矩阵，通常可观察到2～10倍加速。
通过条件：程序打印“验证通过”且加速比>1，说明NPU矩阵乘法实现正确且达到加速目的。

常见问题

初始化失败：确认已安装CANN且当前用户有NPU设备访问权限（通常需root或加入HwHiAiUser用户组）。
内存不足：64×64矩阵仅需约24KB设备内存，若失败请检查NPU显存占用或驱动状态。
精度误差偏大：若误差超过0.01，可改用fp32执行NPU计算（需确认硬件支持），或增大容差至0.05。