From e0d9c8ef2bda78c0c8fe43b61ef4b75b4d7f98f8 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 19 Dec 2025 19:05:29 +0000 Subject: [PATCH] remu fixes --- .github/workflows/test.yml | 28 ++++++++------------- extra/remu/src/lib.rs | 8 +++++- extra/remu/src/work_group.rs | 49 +++++++++++++++++++++++++++++------- test/mockgpu/amd/amdgpu.py | 15 +++++++---- test/mockgpu/helpers.py | 3 +++ tinygrad/helpers.py | 2 +- tinygrad/renderer/rdna.py | 46 +++++++++++++++++++++++---------- tinygrad/runtime/ops_null.py | 7 ++++-- 8 files changed, 109 insertions(+), 49 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c2cb389794..eea6ec39f9 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -656,35 +656,27 @@ jobs: uses: ./.github/actions/process-replay testrdna: - name: Linux (RDNA) - runs-on: ubuntu-22.04 - timeout-minutes: 20 + name: Compile-only (RDNA) + runs-on: ubuntu-24.04 + timeout-minutes: 15 env: - AMD: 1 - AMD_RDNA: 1 - MOCKGPU: 1 - FORWARD_ONLY: 1 + NULL: 1 + NULL_RDNA: 1 steps: - name: Checkout Code uses: actions/checkout@v4 - name: Setup Environment uses: ./.github/actions/setup-tinygrad with: - key: rdna-minimal + key: compile-rdna deps: testing_minimal amd: 'true' - name: Check Device.DEFAULT and print some source run: | - python3 -c "from tinygrad import Device; assert Device.DEFAULT == 'AMD', Device.DEFAULT" - DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add - - name: Run pytest (rdna) - run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/device/test_hcq.py --durations=20 - - name: Run TRANSCENDENTAL math - run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20 - - name: Run RDNA renderer tests - run: python -m pytest test/test_rdna_renderer.py --durations=20 - - name: Run process replay tests - uses: ./.github/actions/process-replay + python3 -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL', Device.DEFAULT" + DEBUG=4 python3 test/test_ops.py TestOps.test_add + - name: Run test_ops (compile-only) + run: python -m pytest -n=auto test/test_ops.py --durations=20 testnvidia: strategy: diff --git a/extra/remu/src/lib.rs b/extra/remu/src/lib.rs index 3f1610fd7d..7166b63a0d 100644 --- a/extra/remu/src/lib.rs +++ b/extra/remu/src/lib.rs @@ -10,6 +10,12 @@ mod work_group; #[no_mangle] pub extern "C" fn run_asm(lib: *const c_char, lib_sz: u32, gx: u32, gy: u32, gz: u32, lx: u32, ly: u32, lz: u32, args_ptr: *const u64) -> i32 { + // Legacy entry point - uses hardcoded SGPR layout (s13/14/15 for workgroup IDs) + run_asm_with_rsrc2(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, 0) +} + +#[no_mangle] +pub extern "C" fn run_asm_with_rsrc2(lib: *const c_char, lib_sz: u32, gx: u32, gy: u32, gz: u32, lx: u32, ly: u32, lz: u32, args_ptr: *const u64, rsrc2: u32) -> i32 { if lib.is_null() || (lib_sz % 4) != 0 { panic!("Pointer is null or length is not properly aligned to 4 bytes"); } @@ -22,7 +28,7 @@ pub extern "C" fn run_asm(lib: *const c_char, lib_sz: u32, gx: u32, gy: u32, gz: for gx in 0..gx { for gy in 0..gy { for gz in 0..gz { - let mut wg = WorkGroup::new(dispatch_dim, [gx, gy, gz], [lx, ly, lz], &kernel, args_ptr); + let mut wg = WorkGroup::new(dispatch_dim, [gx, gy, gz], [lx, ly, lz], &kernel, args_ptr, rsrc2); if let Err(err) = wg.exec_waves() { return err; } diff --git a/extra/remu/src/work_group.rs b/extra/remu/src/work_group.rs index 5f8df5600a..7f33332613 100644 --- a/extra/remu/src/work_group.rs +++ b/extra/remu/src/work_group.rs @@ -13,6 +13,7 @@ pub struct WorkGroup<'a> { kernel_args: *const u64, launch_bounds: [u32; 3], wave_state: HashMap, + rsrc2: u32, // compute_pgm_rsrc2 from kernel descriptor } #[derive(Debug, Clone)] @@ -119,8 +120,8 @@ impl WaveContext { } impl<'a> WorkGroup<'a> { - pub fn new(dispatch_dim: u32, id: [u32; 3], launch_bounds: [u32; 3], kernel: &'a Vec, kernel_args: *const u64) -> Self { - Self { dispatch_dim, id, kernel, launch_bounds, kernel_args, lds: VecDataStore::new(), wave_state: HashMap::new() } + pub fn new(dispatch_dim: u32, id: [u32; 3], launch_bounds: [u32; 3], kernel: &'a Vec, kernel_args: *const u64, rsrc2: u32) -> Self { + Self { dispatch_dim, id, kernel, launch_bounds, kernel_args, lds: VecDataStore::new(), wave_state: HashMap::new(), rsrc2 } } pub fn exec_waves(&mut self) -> Result<(), i32> { @@ -157,10 +158,40 @@ impl<'a> WorkGroup<'a> { scalar_reg.write64(0, self.kernel_args as u64); let [gx, gy, gz] = self.id; - match self.dispatch_dim { - 3 => (scalar_reg[13], scalar_reg[14], scalar_reg[15]) = (gx, gy, gz), - 2 => (scalar_reg[14], scalar_reg[15]) = (gx, gy), - _ => scalar_reg[15] = gx, + + // If rsrc2 is provided, use it to determine workgroup ID placement + // Otherwise fall back to legacy behavior (s13/14/15) + if self.rsrc2 != 0 { + // Parse compute_pgm_rsrc2 to determine SGPR layout + // Bits 1-5: USER_SGPR count + // Bit 7: ENABLE_SGPR_WORKGROUP_ID_X + // Bit 8: ENABLE_SGPR_WORKGROUP_ID_Y + // Bit 9: ENABLE_SGPR_WORKGROUP_ID_Z + let user_sgpr_count = ((self.rsrc2 >> 1) & 0x1f) as usize; + let enable_wg_id_x = (self.rsrc2 >> 7) & 1 != 0; + let enable_wg_id_y = (self.rsrc2 >> 8) & 1 != 0; + let enable_wg_id_z = (self.rsrc2 >> 9) & 1 != 0; + + // Workgroup IDs are placed after user SGPRs + let mut sgpr_idx = user_sgpr_count; + if enable_wg_id_x { + scalar_reg[sgpr_idx] = gx; + sgpr_idx += 1; + } + if enable_wg_id_y { + scalar_reg[sgpr_idx] = gy; + sgpr_idx += 1; + } + if enable_wg_id_z { + scalar_reg[sgpr_idx] = gz; + } + } else { + // Legacy behavior: place workgroup IDs at s13/14/15 based on dispatch_dim + match self.dispatch_dim { + 3 => (scalar_reg[13], scalar_reg[14], scalar_reg[15]) = (gx, gy, gz), + 2 => (scalar_reg[14], scalar_reg[15]) = (gx, gy), + _ => scalar_reg[15] = gx, + } } let mut vec_reg = VGPR::new(); @@ -289,7 +320,7 @@ mod test_workgroup { ]; let addr = (&mut ret as *mut u32) as u64; let kernel = global_store_sgpr(addr, kernel, 106); - let mut wg = WorkGroup::new(1, [0, 0, 0], [3, 1, 1], &kernel, [addr].as_ptr()); + let mut wg = WorkGroup::new(1, [0, 0, 0], [3, 1, 1], &kernel, [addr].as_ptr(), 0); wg.exec_waves().unwrap(); assert_eq!(ret, 0b100); } @@ -305,7 +336,7 @@ mod test_workgroup { ]; let addr = (&mut ret as *mut u32) as u64; let kernel = global_store_sgpr(addr, kernel, 126); - let mut wg = WorkGroup::new(1, [0, 0, 0], [4, 1, 1], &kernel, [addr].as_ptr()); + let mut wg = WorkGroup::new(1, [0, 0, 0], [4, 1, 1], &kernel, [addr].as_ptr(), 0); wg.exec_waves().unwrap(); assert_eq!(ret, 0b0111); } @@ -316,7 +347,7 @@ mod test_workgroup { let kernel = vec![0xBE8D00FF, 0x7FFFFFFF, 0x7E1402FF, u32::MAX, 0xD700000A, 0x0002010A]; let addr = (&mut ret as *mut u32) as u64; let kernel = global_store_sgpr(addr, kernel, 0); - let mut wg = WorkGroup::new(1, [0, 0, 0], [5, 1, 1], &kernel, [addr].as_ptr()); + let mut wg = WorkGroup::new(1, [0, 0, 0], [5, 1, 1], &kernel, [addr].as_ptr(), 0); wg.exec_waves().unwrap(); assert_eq!(ret, 0b11110); } diff --git a/test/mockgpu/amd/amdgpu.py b/test/mockgpu/amd/amdgpu.py index c3fc27d705..a051a2c7fc 100644 --- a/test/mockgpu/amd/amdgpu.py +++ b/test/mockgpu/amd/amdgpu.py @@ -7,7 +7,7 @@ import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am. SDMA_MAX_COPY_SIZE = 0x400000 regCOMPUTE_PGM_LO = 0x1bac + amd_gpu.GC_BASE__INST0_SEG0 -regCOMPUTE_PGM_RSRC2 = 0x1bb3 + amd_gpu.GC_BASE__INST0_SEG0 +regCOMPUTE_PGM_RSRC1 = 0x1bb2 + amd_gpu.GC_BASE__INST0_SEG0 # 0x2e12 - address used by ops_amd.py regCOMPUTE_USER_DATA_0 = 0x1be0 + amd_gpu.GC_BASE__INST0_SEG0 regCOMPUTE_NUM_THREAD_X = 0x1ba7 + amd_gpu.GC_BASE__INST0_SEG0 regGRBM_GFX_INDEX = 0x2200 + amd_gpu.GC_BASE__INST0_SEG1 @@ -180,17 +180,22 @@ class PM4Executor(AMDQueue): prg_addr = (self.gpu.regs[regCOMPUTE_PGM_LO] + (self.gpu.regs[regCOMPUTE_PGM_LO + 1] << 32)) << 8 args_addr = self.gpu.regs[regCOMPUTE_USER_DATA_0] + (self.gpu.regs[regCOMPUTE_USER_DATA_0 + 1] << 32) lc = [self.gpu.regs[i] for i in range(regCOMPUTE_NUM_THREAD_X, regCOMPUTE_NUM_THREAD_X+3)] - rsrc2 = self.gpu.regs[regCOMPUTE_PGM_RSRC2] + # rsrc2 is at COMPUTE_PGM_RSRC1+1 (rsrc1 and rsrc2 are written together) + # Try all SE indexes since broadcast mode might be active + rsrc2 = 0 + for se in range(6): + if (v := self.gpu.regs.regs.get((regCOMPUTE_PGM_RSRC1 + 1, se), 0)) != 0: + rsrc2 = v + break prg_sz = 0 for st,sz in self.gpu.mapped_ranges: if st <= prg_addr < st+sz: prg_sz = sz - (prg_addr - st) assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)" - # Pass valid memory ranges and rsrc2 to Python emulator for bounds checking and SGPR layout + # Pass valid memory ranges to Python emulator for bounds checking if hasattr(remu, 'valid_mem_ranges'): remu.valid_mem_ranges = self.gpu.mapped_ranges - if hasattr(remu, 'rsrc2'): remu.rsrc2 = rsrc2 - err = remu.run_asm(prg_addr, prg_sz, *gl, *lc, args_addr) + err = remu.run_asm_with_rsrc2(prg_addr, prg_sz, *gl, *lc, args_addr, rsrc2) if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel") def _exec_indirect_buffer(self, n): diff --git a/test/mockgpu/helpers.py b/test/mockgpu/helpers.py index 39b8daa336..82e5200b9c 100644 --- a/test/mockgpu/helpers.py +++ b/test/mockgpu/helpers.py @@ -38,6 +38,9 @@ def _try_dlopen_remu(): remu.run_asm.restype = ctypes.c_int32 remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p] + remu.run_asm_with_rsrc2.restype = ctypes.c_int32 + remu.run_asm_with_rsrc2.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, + ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p, ctypes.c_uint32] except OSError: pass else: return remu print("Could not find libremu.so") diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 7514bd34ae..51bf0edb37 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -192,7 +192,7 @@ CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasat # Compilers CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0) NV_PTX, CUDA_PTX, NV_NAK, QCOM_IR3 = ContextVar("NV_PTX", 0), ContextVar("CUDA_PTX", 0), ContextVar("NV_NAK", 0), ContextVar("QCOM_IR3", 0) -NULL_IR3, NULL_NAK = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0) +NULL_IR3, NULL_NAK, NULL_RDNA = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0), ContextVar("NULL_RDNA", 0) AMD_CC, CPU_CC, NV_CC, CUDA_CC = ContextVar("AMD_CC", ""), ContextVar("CPU_CC", ""), ContextVar("NV_CC", ""), ContextVar("CUDA_CC", "") QCOM_CC = ContextVar("QCOM_CC", "") # VIZ implies PROFILE, but you can run PROFILE without VIZ diff --git a/tinygrad/renderer/rdna.py b/tinygrad/renderer/rdna.py index 89bf2b666a..66df77db65 100644 --- a/tinygrad/renderer/rdna.py +++ b/tinygrad/renderer/rdna.py @@ -1,5 +1,5 @@ from typing import Callable, cast -import struct, yaml +import struct from collections import defaultdict from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp from tinygrad.dtype import dtypes, DType, PtrDType, AddrSpace @@ -727,17 +727,37 @@ class RDNARenderer(Renderer): kernarg_size = (args[-1][".offset"] + args[-1][".size"]) if args else 0 - metadata = { - 'amdhsa.kernels': [{ - '.args': args, - '.group_segment_fixed_size': self.lds_size, '.kernarg_segment_align': 8, '.kernarg_segment_size': kernarg_size, - '.language': 'OpenCL C', '.language_version': [1, 2], '.max_flat_workgroup_size': 256, - '.name': function_name, '.private_segment_fixed_size': 0, '.sgpr_count': s_cnt, '.sgpr_spill_count': 0, - '.symbol': f'{function_name}.kd', '.uses_dynamic_stack': False, '.vgpr_count': v_cnt, '.vgpr_spill_count': 0, - '.wavefront_size': 32 - }], - 'amdhsa.target': f'amdgcn-amd-amdhsa--{self.arch}', 'amdhsa.version': [1, 2] - } + # Build metadata YAML manually to avoid yaml dependency + args_yaml = [] + for arg in args: + arg_lines = [' - ' + '\n '.join(f'{k}: {repr(v) if isinstance(v, str) else str(v).lower() if isinstance(v, bool) else v}' + for k, v in arg.items())] + args_yaml.extend(arg_lines) + metadata_yaml = f"""amdhsa.kernels: +- .args: +{chr(10).join(args_yaml)} + .group_segment_fixed_size: {self.lds_size} + .kernarg_segment_align: 8 + .kernarg_segment_size: {kernarg_size} + .language: OpenCL C + .language_version: + - 1 + - 2 + .max_flat_workgroup_size: 256 + .name: {function_name} + .private_segment_fixed_size: 0 + .sgpr_count: {s_cnt} + .sgpr_spill_count: 0 + .symbol: {function_name}.kd + .uses_dynamic_stack: false + .vgpr_count: {v_cnt} + .vgpr_spill_count: 0 + .wavefront_size: 32 +amdhsa.target: amdgcn-amd-amdhsa--{self.arch} +amdhsa.version: +- 1 +- 2 +""" kernel_str = '\n'.join(kernel) # NOTE: .text must be first line for HIPCompiler to detect as assembly @@ -764,7 +784,7 @@ class RDNARenderer(Renderer): " .amdhsa_system_sgpr_workgroup_id_z 1\n" + \ " .amdhsa_system_vgpr_workitem_id 2\n" + \ ".end_amdhsa_kernel\n\n" + \ - ".amdgpu_metadata\n" + yaml.dump(metadata) + ".end_amdgpu_metadata" + ".amdgpu_metadata\n" + metadata_yaml + ".end_amdgpu_metadata" def render(self, uops:list[UOp]) -> str: kernel:list[str] = [] diff --git a/tinygrad/runtime/ops_null.py b/tinygrad/runtime/ops_null.py index 14f5dc63c0..03b4d17267 100644 --- a/tinygrad/runtime/ops_null.py +++ b/tinygrad/runtime/ops_null.py @@ -4,9 +4,11 @@ from tinygrad.engine.jit import MultiGraphRunner from tinygrad.renderer.cstyle import Renderer, CStyleLanguage from tinygrad.renderer.llvmir import AMDLLVMRenderer from tinygrad.uop.ops import Ops -from tinygrad.helpers import cpu_profile, EMULATE, NULL_IR3, NULL_NAK +from tinygrad.helpers import cpu_profile, EMULATE, NULL_IR3, NULL_NAK, NULL_RDNA from tinygrad.renderer.nir import IR3Renderer, NAKRenderer +from tinygrad.renderer.rdna import RDNARenderer from tinygrad.runtime.support.compiler_mesa import IR3Compiler, NAKCompiler +from tinygrad.runtime.support.compiler_amd import HIPCompiler class NullRenderer(CStyleLanguage): device = "NULL" @@ -41,5 +43,6 @@ class NullDevice(Compiled): case _: raise RuntimeError(f"can't EMULATE device: {EMULATE.value}") compilers = CompilerSet([CompilerPair(renderer, Compiler), CompilerPair(functools.partial(IR3Renderer, self), functools.partial(IR3Compiler, 0x6030001), NULL_IR3), # adreno 630 - CompilerPair(functools.partial(NAKRenderer, self), functools.partial(NAKCompiler, "sm_120", 48), NULL_NAK)]) # 5090 + CompilerPair(functools.partial(NAKRenderer, self), functools.partial(NAKCompiler, "sm_120", 48), NULL_NAK), # 5090 + CompilerPair(functools.partial(RDNARenderer, "gfx1100"), functools.partial(HIPCompiler, "gfx1100"), NULL_RDNA)]) super().__init__(device, NullAllocator(self), compilers, functools.partial(NullProgram, device), NullGraph)