nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

remu fixes

Browse Source
This commit is contained in:
geohot
2025-12-29 15:54:14 +00:00
parent 8a8e7d6103
commit e0d9c8ef2b
8 changed files with 109 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/test.yml
+10 -18
View File
@@ -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:
extra/remu/src/lib.rs
+7 -1
View File
@@ -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;
}
extra/remu/src/work_group.rs
+40 -9
View File
@@ -13,6 +13,7 @@ pub struct WorkGroup<'a> {
kernel_args: *const u64,
launch_bounds: [u32; 3],
wave_state: HashMap<usize, WaveState>,
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<u32>, 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<u32>, 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);
}
test/mockgpu/amd/amdgpu.py
+10 -5
View File
@@ -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):
test/mockgpu/helpers.py
+3
View File
@@ -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")
tinygrad/helpers.py
+1 -1
View File
@@ -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
tinygrad/renderer/rdna.py
+33 -13
View File
@@ -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] = []
tinygrad/runtime/ops_null.py
+5 -2
View File
@@ -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)