tinygrad AM 驱动深入解析：RDNA3/RDNA4 用户态 GPU 驱动的队列绑定、启动流程与虚拟内存管理

tinygrad AM 驱动深入解析RDNA3/RDNA4 用户态 GPU 驱动的队列绑定、启动流程与虚拟内存管理【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad本篇技术指南围绕 docs/developer/am.md 展开系统讲解 tinygrad 中面向 AMD RDNA3/RDNA4 的 AMAMD Memory/AMD userspace用户态驱动如何在不加载 amdgpu 内核模块的情况下直接用DEVAMD把计算任务提交给 GPU以及其底层的 MEC 计算队列直绑、三段式启动Boot状态机、regSCRATCH_REG7快速启动标志和单 VMID 三级页目录虚拟内存设计。读完本文你将掌握 AM 驱动的运行前提、AM_RESET/AM_DEBUG环境变量的真实语义并能从源码层面理解 tinygrad 是如何绕过 MES 调度器直接驱动 GFX 与 SDMA IP 的。AM 驱动是什么AMAMD驱动是 tinygrad 仓库中一个完全运行在用户态的 GPU 驱动目标硬件为 AMD RDNA3/RDNA4 架构源码中亦包含针对 MI300/MI350 等 GCD 的特殊处理路径见 tinygrad/runtime/support/am/amdev.py 中GC_HWIP (9,5,0)的 aqua 特判。与传统的 amdgpu 内核驱动 KFD 用户态接口不同AM 驱动在用户态直接完成设备枚举、固件加载、寄存器配置、命令队列管理与中断处理其核心价值正如原文档所述You only need tinygrad to send compute tasks to your GPU!——只要 tinygrad 本身就能把计算任务送达 GPU。从源码目录结构看AM 驱动的实现集中在三处tinygrad/runtime/support/am/amdev.py核心设备对象AMDev、固件加载AMFirmware、页表项AMPageTableEntry、内存管理器AMMemoryManagertinygrad/runtime/support/am/ip.py各 IP 块SOC、GMC、IH、PSP、SMU、GFX、SDMA的软件/硬件初始化逻辑tinygrad/runtime/ops_amd.py对外暴露的AMDDevice设备接口负责创建计算/拷贝队列并提交作业。如何运行 AM 驱动原文档给出的运行方式非常简洁确保 amdgpu 内核模块已卸载然后用DEVAMD运行 tinygrad 即可。# 先卸载 amdgpu 模块需 root 权限且确保没有其他进程占用 GPU modprobe -r amdgpu # 运行 tinygrad显式指定 AMD 后端 DEVAMD python your_tinygrad_script.pyDEV变量的完整取值规则可参考 docs/env_vars.mdDEVAMD使用 AMD 设备DEVAMD:LLVM使用 AMD 设备配合 LLVM 渲染器DEVAMD::gfx950显式指定目标 gfx 版本多个后端可以用组合如DEVUSBAMD:LLVM。AM 驱动在 tinygrad 中被注册为AMDDevice其可用的接口类型包括 KFD、PCI 与 USB见 tinygrad/runtime/ops_amd.py 中ifaces列表PCI/USB 两种即对应 AM 的裸设备路径。可选前置条件原文档列出的两条可选需求分别对应两种高级能力可选需求作用无 IOMMU 的系统开启 P2P点对点与 SDMA 支持。有 IOMMU 时系统会隔离 DMA 地址限制跨设备直接访问与 SDMA 直接读写系统内存vfio-pci内核模块用于 IRQ中断处理。将 GPU 设备绑定到 vfio-pci 后tinygrad 可以通过事件机制接收 GPU 产生的中断而不是轮询仓库中 extra/amdpci/setup_vfio.sh 与 extra/amdpci/setup_python_cap.sh 提供了把 GPU 从 amdgpu 剥离并交给 vfio-pci 的辅助脚本可作为配置参考。中断处理在源码中对应 tinygrad/runtime/ops_amd.py 的_collect_interrupts与on_device_hang通过 IHInterrupt HandlerIP 块收集中断若设备进入错误状态则触发recover恢复流程。环境变量详解原文档给出了两个 AM 专属环境变量变量可选值说明AM_RESET[1]执行一次完整的 GPU 复位重新加载全部固件与 IP 块AM_DEBUG[0-4]设置额外的调试信息输出级别结合源码可以更精确地理解它们的语义。AM_DEBUG分级调试输出AM_DEBUG在 tinygrad/runtime/support/am/amdev.py 通过getenv(AM_DEBUG, 0)读取默认值为0。不同级别对应的实际行为AM_DEBUG 1打印固件加载信息文件名 SHA256 哈希见 amdev.py同时在 RLC 网关RLC gateway拒绝寄存器访问时打印错误码见 amdev.pyAM_DEBUG 4打印每一次寄存器读/写操作寄存器地址与写入/读到的值见 amdev.py。这是追踪驱动对硬件寄存器访问轨迹的最细粒度手段。# 查看固件加载细节与 RLC 网关错误 AM_DEBUG1 DEVAMD python your_tinygrad_script.py # 追踪每一次寄存器读写 AM_DEBUG4 DEVAMD python your_tinygrad_script.pyAM_RESET强制完整启动AM_RESET1会强制走完整启动路径而不是快速启动路径。其判断逻辑位于 amdev.pyself.partial_boot (self.reg(regSCRATCH_REG7).read() AMDev.Version) and (getenv(AM_RESET, 0) ! 1)即只有regSCRATCH_REG7中记录的版本号与当前驱动版本AMDev.Version 0xA000000D见 amdev.py一致且未设置AM_RESET时才允许部分启动partial boot。设置AM_RESET1相当于明确告诉驱动忽略上次留下的状态重新完整初始化。此外源码中还存在一个原文档未列出的扩展变量AM_POWER_LIMIT见 amdev.py设置后通过 SMU 限制 GPU 功耗并手动配置时钟档位不设置时默认把时钟打到最高性能档level-1。Compute 与 SDMA 队列直绑 MEC绕过 MES 直接绑定 MEC原文档明确指出AM 将计算队列直接绑定到 MECMicro Engine Controller从而绕过了 MESMicro Engine Scheduler。这是 AM 驱动降低提交延迟的关键设计——跳过 MES 后命令提交路径更短也不需要 MES 固件配合。tinygrad 在 AM 上只使用一条计算队列绑定在pipe0 queue0同样只有一条 SDMA 队列绑定在engine0 queue0。源码中这一映射体现在 ip.py 的setup_ring与_grbm_selectpipe, queue, doorbell idx // 4, idx % 4, am.AMDGPU_NAVI10_DOORBELL_MEC_RING0 for xcc in range(self.xccs if aql else 1): self._grbm_select(me1, pipepipe, queuequeue, instxcc)当idx0时即得到pipe0 queue0并通过regGRBM_GFX_CNTL选中 ME1me1下的对应计算管道。随后驱动把 MQDMemory Queue Descriptor写入显存并激活队列regCP_HQD_ACTIVE.write(0x1)见 ip.py完成硬件队列的挂载。队列创建与提交路径上层创建队列的入口在 ops_amd.py 的create_queueif queue_type kfd.KFD_IOC_QUEUE_TYPE_SDMA: doorbell_index self.dev_impl.sdma.setup_ring(*(rcvr_params:(ring.va_addr, ring.size, gart.va_addrrptr, gart.va_addrwptr, idx))) else: doorbell_index self.dev_impl.gfx.setup_ring(*(rcvr_params:(ring.va_addr, ring.size, gart.va_addrrptr, gart.va_addrwptr, eop_buffer.va_addr, eop_buffer.size, is_aql:(queue_typekfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))计算队列使用AMDComputeQueuePM4 包提交拷贝队列使用AMDCopyQueueSDMA 命令提交两者都继承自HWQueue。SDMA 队列由 ip.py 起的AM_SDMA类管理engine0 queue0对应其pipe0, queue0的寄存器选择逻辑。门铃Doorbell机制队列激活后CPU 侧通过写门铃doorbell通知 GPU 有新任务。门铃地址映射自 PCI BAR2self.doorbell64 self.pci_dev.map_bar(2, fmtQ)见 amdev.py提交作业时signal_doorbell更新写指针以触发硬件消费命令环ring buffer具体实现在 ops_amd.py。Boot 启动流程三种状态与快速启动GPU 的三种可能状态原文档将 AM 驱动接手 GPU 时的状态划分为三种未初始化Not initialized已被 amdgpu 初始化Initialized by amdgpu已被 AM 初始化Initialized by AM第一、二种状态下 GPU 的内部状态未知因此必须执行完整的 GPU 配置其中第二种被 amdgpu 初始化过还需要一次mode1 reset来复位所有组件。第三种状态则可以部分配置以优化启动时间——此时只需要初始化 GFX 与 SDMA 两个 IP 块。源码中这段逻辑以注释形式保留在 amdev.py实现上则是if not self.partial_boot: ... self.smu.mode1_reset() # 状态 2 → 强制 mode1 reset 复位全部组件 ... self.init_hw(self.soc, self.gmc, self.ih, self.psp, self.smu) # 完整初始化 elif not self.is_vf: self.psp._tmr_init() # 无论哪种启动方式GFX 与 SDMA 都会重新初始化 self.init_hw(self.gfx, self.sdma)对应 amdev.py——完整启动时依次初始化 SOC、GMC、IH、PSP、SMU而快速启动时这些 IP 保持原样仅重新初始化 GFX 与 SDMA。独立 Boot 内存与 regSCRATCH_REG7 标志快速启动之所以可行依赖两块关键设计独立的 boot 内存AM 使用一块保证不会被覆盖的独立物理内存boot memory存放所有只在首次 AM 启动时初始化的块PSP 的 TMR、GMC 的 memscratch/dummy page、IH 的 ring 等均在bootTrue条件下分配见 ip.py 与 amdev.py 中boot_size(3 20)。由于这块内存的内容在重启间得以保留第二次启动就不需要重建这些块。regSCRATCH_REG7状态标志AM 用regSCRATCH_REG7记录驱动版本号AMDev.Version作为本 GPU 是否已被 AM 初始化过的标志。启动时读取该寄存器若等于当前版本则认定可走快速启动路径。首次初始化结束时驱动会把版本写入该寄存器见 amdev.py。除此之外regSCRATCH_REG6被用作上次 AM 会话是否正常结束的标志见 amdev.py若上次会话异常终止fini时写入is_err_state见 amdev.py即便regSCRATCH_REG7匹配也会强制走完整复位见 amdev.py 的Malformed state检测。regSCRATCH_REG5则用于保存 PSP TMR 的大小供快速启动复用见 amdev.py。# 场景举例 # 1. 第一次使用 AM 驱动 / 上次会话异常自动完整启动 DEVAMD python your_tinygrad_script.py # 2. 希望强制完整复位重新加载全部固件与 IP 块 AM_RESET1 DEVAMD python your_tinygrad_script.py快速启动时的 MEC 复位在快速启动路径下GFX 的init_hw会提前返回并只执行 MEC 复位reset_mec见 ip.pyif self.adev.partial_boot: return self.reset_mec()reset_mec先出队所有 HQD硬件队列描述符再通过regGRBM_SOFT_RESET软复位 CP/CPC最后重新配置并启用 MEC见 ip.py。这样既能清掉上一次会话遗留的队列状态又避免了整卡复位带来的开销。VM 管理单 VMID 与三级页目录单一 VMID0 与单页目录与多进程内核驱动需要为每个上下文分配独立 VMID 不同AM 驱动每个设备只设置一个VMID0和一个页目录。所有 TLB 刷新都针对vmid0进行例如 ip.py 中self.dev.gmc.flush_tlb(ipGC, vmid0) self.dev.gmc.flush_tlb(ipMM, vmid0)页表项构建在 amdev.py 的AMPageTableEntry中支持大页huge page与系统内存/物理内存地址空间AddrSpace.SYS/AddrSpace.PHYS的区分映射。三级页目录与 512GB 虚拟地址空间原文档说明该页目录为三级结构支持最多 512GB 的虚拟地址。从源码看内存管理器以 amdev.py 的配置构建页表self.mm AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size(3 20), pt_tAMPageTableEntry, va_shifts[12, 21, 30, 39], va_bits48, first_lvam.AMDGPU_VM_PDB2, va_baseAMMemoryManager.va_allocator.base, reserve_ptablenot self.large_bar, ...)first_lvam.AMDGPU_VM_PDB2页表从 PDB2 一级开始构建配合va_shifts[12, 21, 30, 39]的各级偏移布局实现文档所述的三级页目录结构va_bits48虚拟地址位宽为 48 位va_allocator TLSFAllocator((1 44), base0x200000000000)见 amdev.py虚拟地址分配器覆盖 144 字节16TB的地址窗口基址固定为0x200000000000。所有 AM 设备共享一个虚拟地址空间AMMemoryManager.va_allocator的注释明确写着global for all devices——所有 AM 设备的虚拟地址都从同一个分配器中划拨因此所有 AM 设备位于同一个虚拟地址空间。这意味着在多卡含 XGMI 互联的 hive场景下任意设备的虚拟地址可以直接映射到其他设备/本设备的物理显存为跨设备 P2P 访问与统一寻址提供了基础。AMPageTableEntry.set_entry中的paddr2xgmi与xgmi2paddr转换amdev.py正是这一统一地址空间下物理地址与 XGMI 地址互转的实现细节。地址空间与页表权限标志页表项通过gmc.get_pte_flags生成 PTE 标志valid / system / snooped / uncached / fragment 等见 amdev.py且所有映射建立后都会触发 GC 与 MM 两个 hub 的 TLB 失效amdev.py保证 CPU 侧修改页表后 GPU 能看到最新映射。故障恢复机制虽然原文档未展开但 AM 驱动的恢复路径与启动流程直接相关值得补充当 IH 收集到中断并判定设备进入错误状态后recover见 amdev.py会执行gfx.reset_mec()清空 MEC 上的遗留队列随后由上层重置队列读写指针并重建计算队列见 ops_amd.py。这一机制与regSCRATCH_REG6会话标志配合构成了异常 → 复位 → 重建队列的闭环也是为什么快速启动必须校验上次会话是否正常结束的原因。总结与延伸阅读AM 驱动是 tinygrad 在 AMD 平台上用软件吃掉硬件管理的实践计算队列直绑 MEC 绕开 MES 降低延迟regSCRATCH_REG7 独立 boot 内存实现秒级快速启动单 VMID 三级页目录 全局虚拟地址空间支撑多卡统一寻址。围绕本文主题可以在仓库中继续深入docs/developer/am.md本文所依据的原始开发文档tinygrad/runtime/support/am/amdev.pyAMDev设备对象、固件加载与内存/页表管理核心实现tinygrad/runtime/support/am/ip.py各 IP 块含 GFX/SDMA/PSP/SMU的初始化细节tinygrad/runtime/ops_amd.pyAMDDevice设备接口、队列创建与作业提交路径docs/env_vars.mdDEV变量的完整取值与后端选择规则extra/amdpci/setup_vfio.shvfio-pci 绑定辅助脚本用于满足 AM 驱动的 IRQ 处理前置条件。实际使用 AM 驱动时请确保系统满足amdgpu 模块已卸载的前提并视需要配合无 IOMMU 环境与 vfio-pci 以获得 P2P/SDMA 与中断能力。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考