AMD GPU 内存优化实战：从显存管理到性能调优

1次阅读

共计 2569 个字符，预计需要花费 7 分钟才能阅读完成。

在深度学习训练和科学计算中，AMD GPU 以其出色的性价比和并行计算能力逐渐受到开发者青睐。然而，随着模型规模扩大和计算复杂度提升，显存管理问题逐渐成为性能瓶颈。本文将分享我在 AMD GPU 显存优化方面的实战经验，希望能帮助开发者更好地驾驭 ROCm 生态。

显存碎片化问题 ：频繁的小块内存分配释放会导致显存碎片化，最终可能因找不到足够大的连续显存空间而触发 OOM（Out Of Memory）。
PCIe 带宽竞争 ：当主机内存与设备显存之间需要频繁传输数据时，有限的 PCIe 带宽会成为性能瓶颈，导致计算单元等待数据。
内存访问延迟 ：不合理的访存模式会导致缓存命中率下降，显著增加内存访问延迟。

CUDA 和 ROCm 都提供了统一内存管理机制，但在实现细节上有所不同：

// CUDA Unified Memory 示例
cudaMallocManaged(&data, size);
...
// ROCm HIP 等效实现
hipMallocManaged(&data, size, hipMemAttachGlobal);

关键差异点：

CUDA 的 Unified Memory 默认采用按需页面迁移策略
HIP 的 hipMallocManaged 需要显式指定内存附加域（hipMemAttachGlobal 或 hipMemAttachHost）
ROCm 4.5+ 引入了更细粒度的缓冲区管理能力

// 分配统一内存
float *data;
hipMallocManaged(&data, N*sizeof(float), hipMemAttachGlobal);

// 在主机初始化数据
for(int i=0; i<N; i++) data[i] = 1.0f;

// 内核函数直接访问
__global__ void kernel(float *data) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    data[idx] *= 2.0f;
}

// 分配可分页锁定主机内存
float *h_data;
hipHostMalloc(&h_data, N*sizeof(float), hipHostMallocDefault);

// 设备内存分配
float *d_data;
hipMalloc(&d_data, N*sizeof(float));

// 异步传输优化
hipMemcpyAsync(d_data, h_data, N*sizeof(float), hipMemcpyHostToDevice, stream);

ROCm 4.5+ 引入了细粒度缓冲区支持，可以更精确地控制内存访问：

// 创建细粒度缓冲区
hipMemAllocationType type = hipMemAllocationTypePinned;
hipMemAllocationHandleType handleType = hipMemHandleTypeNone;
hipMemAllocationProp prop{};
prop.type = type;
prop.location.type = hipMemLocationTypeDevice;
prop.location.id = deviceId;

hipMemGenericAllocationHandle_t handle;
hipMemCreate(&handle, size, &prop, handleType);

下面是一个优化后的矩阵乘法实现，展示了多个内存优化技巧：

// 使用 __restrict__ 避免指针别名
__global__ void matmul_kernel(const float * __restrict__ A,
                             const float * __restrict__ B,
                             float * __restrict__ C,
                             int M, int N, int K) {// 分块优化实现...}

int main() {
    // 分配统一内存
    float *A, *B, *C;
    hipMallocManaged(&A, M*K*sizeof(float));
    hipMallocManaged(&B, K*N*sizeof(float));
    hipMallocManaged(&C, M*N*sizeof(float));

    // 设置内存建议
    hipMemAdvise(A, M*K*sizeof(float), hipMemAdviseSetPreferredLocation, deviceId);
    hipMemAdvise(B, K*N*sizeof(float), hipMemAdviseSetAccessedBy, deviceId);

    // 预取数据
    hipMemPrefetchAsync(A, M*K*sizeof(float), deviceId, 0);
    hipMemPrefetchAsync(B, K*N*sizeof(float), deviceId, 0);

    // 执行内核
    dim3 blocks(32, 32);
    dim3 threads(16, 16);
    matmul_kernel<<<blocks, threads>>>(A, B, C, M, N, K);

    // 使用 rocprof 分析
    // rocprof --stats -o profile.csv ./matmul
}