📚 课程目录第2章 NPU内存带宽测试

2. NPU内存带宽测试

任务目标

编写一个运行于NPU上的内存带宽测试程序，分别测量DDR和SRAM的读写带宽，并通过一个矩阵乘法示例说明带宽对模型推理性能的约束关系。

环境依赖

算能SDK（含libbmcpu.so运行时库）
Python 3.8+，安装sophon包（pip install sophon）
TPU-MLIR编译器（用于将测试代码编译为bmodel）

实现原理

NPU通过独立的内存控制器访问DDR和片上SRAM。测试程序使用DMA引擎在DDR/SRAM与运算单元之间搬运数据块，通过计时搬运固定大小数据所需时间计算带宽。矩阵乘法示例用于展示：当算力充足时，实际执行时间由内存带宽决定（memory-bound）。

关键完整代码

# bandwidth_test.py
import sophon.sail as sail
import numpy as np
import time

def measure_bandwidth(handle, src_dev, dst_dev, size_mb, is_sram):
    """测量设备间拷贝带宽"""
    bytes_total = size_mb * 1024 * 1024
    src = sail.Tensor(handle, [bytes_total // 4], sail.Dtype.BM_FLOAT32, True, src_dev)
    dst = sail.Tensor(handle, [bytes_total // 4], sail.Dtype.BM_FLOAT32, True, dst_dev)
    
    # 初始化源数据
    src_data = np.random.rand(bytes_total // 4).astype(np.float32)
    src.sync_from_np(src_data)
    
    # 预热
    for _ in range(3):
        dst.sync_from_tensor(src)
    handle.sync()
    
    # 正式计时
    start = time.perf_counter()
    iterations = 10
    for _ in range(iterations):
        dst.sync_from_tensor(src)
    handle.sync()
    elapsed = time.perf_counter() - start
    
    bandwidth = (bytes_total * iterations) / (elapsed * 1e9)  # GB/s
    return bandwidth

def matrix_multiply_bandwidth_limited(handle, size):
    """矩阵乘法：验证带宽瓶颈"""
    a = sail.Tensor(handle, [size, size], sail.Dtype.BM_FLOAT32, True, sail.Device.DDR)
    b = sail.Tensor(handle, [size, size], sail.Dtype.BM_FLOAT32, True, sail.Device.DDR)
    c = sail.Tensor(handle, [size, size], sail.Dtype.BM_FLOAT32, True, sail.Device.DDR)
    
    a_data = np.random.rand(size, size).astype(np.float32)
    b_data = np.random.rand(size, size).astype(np.float32)
    a.sync_from_np(a_data)
    b.sync_from_np(b_data)
    
    # 使用sail的矩阵乘接口（底层调用NPU的tpu_kernel）
    handle.matmul(a, b, c, False, False)
    handle.sync()
    
    # 计算理论最小时间
    ops = 2 * size**3
    # 假设NPU算力为32 TOPS（需根据实际芯片调整）
    compute_time = ops / (32e12)
    # 数据搬运量：读A+B，写C
    data_bytes = 3 * size * size * 4
    ddr_bw = measure_bandwidth(handle, sail.Device.DDR, sail.Device.DDR, 64, False)
    mem_time = data_bytes / (ddr_bw * 1e9)
    
    return compute_time, mem_time

if __name__ == "__main__":
    handle = sail.Handle(0)  # 使用NPU设备0
    
    # 1. DDR带宽测试
    ddr_bw = measure_bandwidth(handle, sail.Device.DDR, sail.Device.DDR, 128, False)
    print(f"DDR带宽: {ddr_bw:.2f} GB/s")
    
    # 2. SRAM带宽测试（通过DDR到SRAM的拷贝近似）
    sram_bw = measure_bandwidth(handle, sail.Device.DDR, sail.Device.SRAM, 16, True)
    print(f"SRAM带宽: {sram_bw:.2f} GB/s")
    
    # 3. 矩阵乘带宽影响分析
    size = 2048
    compute_t, mem_t = matrix_multiply_bandwidth_limited(handle, size)
    print(f"矩阵乘({size}x{size}):")
    print(f"  理论计算时间: {compute_t*1000:.2f} ms")
    print(f"  理论内存时间: {mem_t*1000:.2f} ms")
    print(f"  结论: {'内存瓶颈' if mem_t > compute_t else '算力瓶颈'}")

编译与运行

# 安装依赖
pip install sophon

# 直接运行（sophon包已包含NPU运行时）
python bandwidth_test.py

验证方法

带宽数值合理性：DDR带宽应在芯片规格的60%~85%范围内（例如BM1684X理论DDR带宽约102.4GB/s，实测应大于60GB/s）；SRAM带宽通常为DDR的3~5倍。
矩阵乘瓶颈判断：当mem_t > compute_t时输出"内存瓶颈"，说明当前矩阵乘受限于DDR带宽；若输出"算力瓶颈"则说明计算单元已饱和。
可重复性：连续运行3次，带宽波动应小于5%。

常见问题

设备初始化失败：确认sail.Handle(0)中的设备号与bm-smi显示的NPU编号一致。
SRAM容量不足：SRAM通常只有几十MB，测试时size_mb不要超过实际SRAM大小（BM1684X为32MB），否则会回退到DDR。
带宽偏低：检查是否开启了NPU的DVFS节能模式，可通过bm-smi --set-dvfs 0关闭。