使用 gem5 模拟 MI300X，立省 10 万块？

General
1

前段时间朋友发了一篇使用 gem5 模拟 MI300X 的知乎，正好我最近在验证 AMDGPU 的浮点运算精度，就想着对比一下 gem5 MI300X 的 model 浮点精度和真实硬件有没有什么差异。

这里推荐使用服务器或者工作站来运行 gem5，个人电脑资源可能不太够用。

主要参考雾佬的文章，以及官方提供的文档 Full System AMD GPU model 来搭建。

基本思路是:

  1. 使用 qemu-system-x86_64 来制作包含 AMDGPU 驱动和 ROCm 环境的镜像，以及所需的内核；
  2. 使用 gem5 的 kvm 模式，模拟一个 x86 ，将 MI300X 的 model 作为 PCIe 卡集成进去；
  3. 最后使用 gem5 加载制作好的镜像和内核，就可以使用了。

为了方便，我使用 gem5 提供的 docker 环境来构建。

docker pull ghcr.io/gem5/ubuntu-24.04_all-dependencies:v24-0

在本地拉取 gem5 相关仓库，并创建 gem5 的 docker 容器，然后将仓库所在目录共享给容器：

git clone https://github.com/gem5/gem5.git
git clone https://github.com/gem5/gem5-resources.git gem5/gem5-resources
docker run --name gem5-amdgpu \
--device /dev/kvm \
--volume /home/zevorn/gem5:/gem5 \
--privileged \
-it ghcr.io/gem5/ubuntu-24.04_all-dependencies:v24-0

接下来就是编译环节，分两步。

首先在容器外面的 gem5/gem5-resources 仓库里制作需要的镜像和内核：

cd gem5/gem5-resources/src/x86-ubuntu-gpu-ml/
./build.sh

这里需要耐心等待，这个构建脚本会使用 packer 来制作镜像，制作成功以后，会生成以下文件：

├── disk-image
│   └── x86-ubuntu-gpu-ml
├── vmlinux-gpu-ml

然后进入容器内，开始构建 gem5：

cd /gem5
scons build/VEGA_X86/gem5.opt -j$(nproc)

我们需要修改默认的 MI300X 的 python 文件，避免运行完退出：

$ git diff configs/example/gpufs/mi300.py
@@ -153,7 +153,7 @@ def runMI300GPUFS(
         )
         b64file.write(runscriptStr)

-    args.script = tempRunscript
+    # args.script = tempRunscript

     # Defaults for CPU
     args.cpu_type = "X86KvmCPU"

在 gem5 源码路径下新建一个 pytorch_test.py 文件，用于 gem5 运行 mi300 时加载，但实际我们不会执行它：

#!/usr/bin/env python3

import torch

x = torch.rand(5, 3).to('cuda')
y = torch.rand(3, 5).to('cuda')

z = x @ y

我们在容器内安装一个 tmux，方便启动 gem5 以后，通过另一个窗格连接进终端：

apt update
apt install tmux

然后启动试试：

build/VEGA_X86/gem5.opt configs/example/gpufs/mi300.py --disk-image gem5-resources/src/x86-ubuntu-gpu-ml/disk-image/x86-ubuntu-gpu-ml --kernel gem5-resources/src/x86-ubuntu-gpu-ml/vmlinux-gpu-ml --app pytorch_test.py

我们使用tmux，在当前窗口再新建一个终端，然后使用下面的命令连接到 gem5:

./util/term/gem5term localhost 3456
root@gem5:~#

然后我们需要手动加载 AMDGPU 驱动，可以直接在 gem5 终端输入：

export LD_LIBRARY_PATH=/opt/rocm/lib:$LD_LIBRARY_PATH
export HSA_ENABLE_INTERRUPT=0
export HCC_AMDGPU_TARGET=gfx942
dmesg -n8
cat /proc/cpuinfo
dd if=/root/roms/mi300.rom of=/dev/mem bs=1k seek=768 count=128

我们来验证一下是否加载成功：

root@gem5:~# rocminfo
ROCk module version 6.12.12 is loaded
=====================
HSA System Attributes
=====================
Runtime Version:         1.15
Runtime Ext Version:     1.7
System Timestamp Freq.:  0.001000MHz
Sig. Max Wait Duration:  18446744073709551615 (0xFFFFFFFFFFFFFFFF) (timestamp count)
Machine Model:           LARGE
System Endianness:       LITTLE
Mwaitx:                  DISABLED
XNACK enabled:           NO
DMAbuf Support:          YES
VMM Support:             YES
...
*******
Agent 2
*******
  Name:                    gfx942
  Uuid:                    GPU-XX
  Marketing Name:          AMD Instinct MI300X
  Vendor Name:             AMD
  Feature:                 KERNEL_DISPATCH
  Profile:                 BASE_PROFILE
  Float Round Mode:        NEAR
  Max Queue Number:        128(0x80)
  Queue Min Size:          64(0x40)
  Queue Max Size:          131072(0x20000)
  Queue Type:              MULTI
  Node:                    1
  Device Type:             GPU
...

接下来就可以正常使用了。

我在测试 gem5 MI300X 的 fma 指令的运算结果是，和 cpu 以及 MI300X 真实硬件进行 bit 级对比后，发现总是相差 3~4 ulp，于是我阅读 fma 的源码，发现是 gem5 用乘加两步运算来模拟的，这导致了和真实硬件的浮点精度存在误差，于是我修改源码：

$ git diff src/arch/amdgpu/vega/insts/instructions.hh
diff --git a/src/arch/amdgpu/vega/insts/instructions.hh b/src/arch/amdgpu/vega/insts/instructions.hh
index 7a328f9230..76bd96beee 100644
--- a/src/arch/amdgpu/vega/insts/instructions.hh
+++ b/src/arch/amdgpu/vega/insts/instructions.hh
@@ -44352,8 +44352,9 @@ namespace VegaISA
                         int lane_A = i + M * (block + B * (k / K_L));
                         int lane_B = j + N * (block + B * (k / K_L));
                         int item = k % K_L;
-                        result[i][j] +=
-                          src0[item][lane_A] * src1[item][lane_B];
+                        result[i][j] =
+                            std::fma(src0[item][lane_A], src1[item][lane_B],
+                                     result[i][j]);
                     }
                 }
             }

之后测试就正常了，我将这笔 bugfix，贡献到了上游社区，目前已被合并：arch-vega: Improve MFMA precision to match MI300X hardware

本文首发于微信公众号：GTOC

5 个赞
2

这个标题太强了，这个星期横竖我都要试试！

2 个赞