nx_kfd.nx source

1// nx_kfd.nx -- the native-Linux AMD executor shim (/dev/kfd AM/KFD path), sibling to nx_nv (NVIDIA) + nx_dxg 2// (WSL2). Modeled on tinygrad's ops_amd.py AM driver: PM4 ring buffer + wptr + doorbell, bypassing the MES 3// scheduler firmware = the MOST-sovereign GPU path per deep-research. Composes PM4 packets (nx_pm4_asm) into a 4// compute-queue ring, updates wptr, rings the doorbell. Runs on a native-Linux / rented AMD cloud instance 5// (brick-safe, rule #26). 6// 7// LOCALLY-VERIFIABLE (gated): the PM4 compute-dispatch stream composition (via nx_pm4_asm, already gated). 8// CLOUD-PENDING (structured): the KFD queue creation (AMDKFD_IOC_CREATE_QUEUE on /dev/kfd) + the ring/doorbell 9// BAR map -- completed against the instance's ROCt/KFD ABI on first cloud run. Pure funcs, no main. license_tier: ORIGINAL 10import "nx_syscalls.nx" 11import "nx_pm4_asm.nx" 12 13// AMD KFD ioctls (AMDKFD_IOCTL_BASE='K'=0x4B): the ioctl = _IOWR(0x4B, nr, size). (struct sizes = TODO_CLOUD) 14const KFD_IOC_CREATE_QUEUE: i64 = 0x02 15const KFD_IOC_CREATE_EVENT: i64 = 0x05 16const KFD_IOC_ALLOC_MEMORY_OF_GPU: i64 = 0x16 17 18func kfd_open() -> i64 { return sys_openat_rd("/dev/kfd" as *u8) } 19func kfd_open_render() -> i64 { return sys_openat_rd("/dev/dri/renderD128" as *u8) } 20 21// write one 32-bit dword into the PM4 ring at byte offset `off`. 22func pm4_ring_put(ring: *u8, off: i64, dw: i64) -> i64 { let w: *u8 = (ring as i64 + off) as *u8; w[0]=(dw&0xff) as u8; w[1]=((dw>>8)&0xff) as u8; w[2]=((dw>>16)&0xff) as u8; w[3]=((dw>>24)&0xff) as u8; return 0 } 23 24// compose a minimal PM4 COMPUTE DISPATCH into `ring`: SET_SH_REG(kernel params) + DISPATCH_DIRECT(dims). 25// reg values are 0-framed here (real kernel address/rsrc are cloud-supplied); returns bytes written. 26func pm4_compose_dispatch(ring: *u8, reg_base: i64, n_regs: i64, dim_x: i64, dim_y: i64, dim_z: i64) -> i64 { 27 var o: i64 = 0 28 pm4_ring_put(ring, o, pm4_packet3(PM4_IT_SET_SH_REG, n_regs + 1)); o = o + 4 29 pm4_ring_put(ring, o, reg_base); o = o + 4 30 var i: i64 = 0 31 while i < n_regs { pm4_ring_put(ring, o, 0); o = o + 4; i = i + 1 } 32 pm4_ring_put(ring, o, pm4_packet3(PM4_IT_DISPATCH_DIRECT, 3)); o = o + 4 33 pm4_ring_put(ring, o, dim_x); o = o + 4 34 pm4_ring_put(ring, o, dim_y); o = o + 4 35 pm4_ring_put(ring, o, dim_z); o = o + 4 36 pm4_ring_put(ring, o, 1); o = o + 4 // dispatch_initiator (COMPUTE_SHADER_EN) 37 return o 38} 39 40// ring the doorbell: 32-bit write of the wptr (in dwords) to the queue's mapped doorbell address. 41func kfd_doorbell_ring(doorbell: i64, wptr_dwords: i64) -> i64 { 42 let b: *u8 = doorbell as *u8 43 b[0]=(wptr_dwords&0xff) as u8; b[1]=((wptr_dwords>>8)&0xff) as u8; b[2]=((wptr_dwords>>16)&0xff) as u8; b[3]=((wptr_dwords>>24)&0xff) as u8 44 return 0 45}