Update251203 (#233)

tinygrad_repo/tinygrad/runtime/support/am/amdev.py

@@ -84,7 +84,7 @@ class AMFirmware:
     self.descs += [self.desc(blob, hdr0.header.ucode_array_offset_bytes, hdr0.header.ucode_size_bytes, am.GFX_FW_TYPE_RLC_G)]
 
   def load_fw(self, fname:str, *headers, versioned_header:str|None=None):
-    fpath = fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/45f59212aebd226c7630aff4b58598967c0c8c91/amdgpu/{fname}", subdir="fw")
+    fpath = fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/a9f26799247aa60fbaa3b64267a18f20b72b5235/amdgpu/{fname}", subdir="fw")
     blob = memoryview(bytearray(fpath.read_bytes()))
     if AM_DEBUG >= 1: print(f"am {self.adev.devfmt}: loading firmware {fname}: {hashlib.sha256(blob).hexdigest()}")
     if versioned_header:

tinygrad_repo/tinygrad/runtime/support/am/ip.py

@@ -224,7 +224,8 @@ class AM_GFX(AM_IP):
     self._grbm_select()
     self.adev.regGCVM_CONTEXT0_CNTL.write(0)
 
-  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, doorbell:int, pipe:int, queue:int):
+  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, doorbell:int, pipe:int, queue:int,
+                 aql:bool):
     mqd = self.adev.mm.valloc(0x1000, uncached=True, contiguous=True)
 
     struct_t = getattr(am, f"struct_v{self.adev.ip_ver[am.GC_HWIP][0]}_compute_mqd")
@@ -235,9 +236,10 @@ class AM_GFX(AM_IP):
       cp_hqd_pq_rptr_report_addr_lo=lo32(rptr_addr), cp_hqd_pq_rptr_report_addr_hi=hi32(rptr_addr),
       cp_hqd_pq_wptr_poll_addr_lo=lo32(wptr_addr), cp_hqd_pq_wptr_poll_addr_hi=hi32(wptr_addr),
       cp_hqd_pq_doorbell_control=self.adev.regCP_HQD_PQ_DOORBELL_CONTROL.encode(doorbell_offset=doorbell*2, doorbell_en=1),
-      cp_hqd_pq_control=self.adev.regCP_HQD_PQ_CONTROL.encode(rptr_block_size=5, unord_dispatch=0, queue_size=(ring_size//4).bit_length()-2),
+      cp_hqd_pq_control=self.adev.regCP_HQD_PQ_CONTROL.encode(rptr_block_size=5, unord_dispatch=0, queue_size=(ring_size//4).bit_length()-2,
+        **({'queue_full_en':1, 'slot_based_wptr':2, 'no_update_rptr':1} if aql else {})),
       cp_hqd_ib_control=self.adev.regCP_HQD_IB_CONTROL.encode(min_ib_avail_size=0x3), cp_hqd_hq_status0=0x20004000,
-      cp_mqd_control=self.adev.regCP_MQD_CONTROL.encode(priv_state=1), cp_hqd_vmid=0,
+      cp_mqd_control=self.adev.regCP_MQD_CONTROL.encode(priv_state=1), cp_hqd_vmid=0, cp_hqd_aql_control=int(aql),
       cp_hqd_eop_base_addr_lo=lo32(eop_addr>>8), cp_hqd_eop_base_addr_hi=hi32(eop_addr>>8),
       cp_hqd_eop_control=self.adev.regCP_HQD_EOP_CONTROL.encode(eop_size=(eop_size//4).bit_length()-2))
 

tinygrad_repo/tinygrad/runtime/support/compiler_amd.py

@@ -9,7 +9,7 @@ try:
     assert comgr.AMD_COMGR_LANGUAGE_HIP == 3
 except AttributeError: pass  # ignore if ROCm isn't installed
 from tinygrad.device import Compiler, CompileError
-from tinygrad.runtime.ops_llvm import LLVMCompiler
+from tinygrad.runtime.support.compiler_cpu import LLVMCompiler
 from tinygrad.helpers import OSX, to_char_p_p
 
 def amdgpu_disassemble(lib:bytes):

tinygrad_repo/tinygrad/runtime/support/compiler_cpu.py

@@ -0,0 +1,86 @@
+import ctypes, platform, sys, subprocess
+from tinygrad.device import Compiler
+from tinygrad.helpers import OSX, getenv, capstone_flatdump, DEBUG
+from tinygrad.runtime.support.elf import jit_loader
+try: import tinygrad.runtime.autogen.llvm as llvm
+except (ImportError, FileNotFoundError): llvm = None #type:ignore[assignment]
+
+class ClangJITCompiler(Compiler):
+  def __init__(self, cachekey="compile_clang_jit"): super().__init__(cachekey)
+
+  def compile(self, src:str) -> bytes:
+    # -fno-math-errno is required for __builtin_sqrt to become an instruction instead of a function call
+    # x18 is a reserved platform register. It is clobbered on context switch in macos and is used to store TEB pointer in windows on arm, don't use it
+    target = 'x86_64' if sys.platform == 'win32' else platform.machine()
+    # on arm march means "runs on this arch and superset" instead of "optimize for this arch". x86 march == arm mcpu
+    arch = {'x86_64': '-march=native', 'AMD64': '-march=native', 'riscv64': '-march=rv64g'}.get(platform.machine(), "-mcpu=native")
+    args = [arch, f'--target={target}-none-unknown-elf', '-O2', '-fPIC', '-ffreestanding', '-fno-math-errno', '-nostdlib', '-fno-ident']
+    arch_args = ['-ffixed-x18'] if target == 'arm64' else []
+    obj = subprocess.check_output([getenv("CC", 'clang'), '-c', '-x', 'c', *args, *arch_args, '-', '-o', '-'], input=src.encode('utf-8'))
+    return jit_loader(obj)
+
+  def disassemble(self, lib:bytes): return capstone_flatdump(lib)
+
+def cerr(): return ctypes.pointer(ctypes.pointer(ctypes.c_char()))
+
+def expect(x, err, ret=None):
+  if x: raise RuntimeError(llvm.string_cast(err.contents) if not isinstance(err, str) else err)
+  return ret
+
+class LLVMCompiler(Compiler):
+  jit = True
+  target_arch = {'arm64': 'AArch64', 'aarch64': 'AArch64', 'x86_64': 'X86', 'AMD64': 'X86', 'riscv64': 'riscv64'}[platform.machine()]
+  def __init__(self, processor:str, feats:str):
+    for component in ['Target', 'TargetInfo', 'TargetMC', 'AsmParser', 'AsmPrinter']: getattr(llvm, f'LLVMInitialize{self.target_arch}{component}')()
+
+    triple = {'AArch64': b'aarch64-none-unknown-elf', 'X86': b'x86_64-none-unknown-elf', 'AMDGPU': b'amdgcn-amd-amdhsa'}[self.target_arch]
+    target = expect(llvm.LLVMGetTargetFromTriple(triple, ctypes.pointer(tgt:=llvm.LLVMTargetRef()), err:=cerr()), err, tgt)
+    if DEBUG >= 3: print(f"LLVM init for {processor!r} with {feats!r}")
+    self.target_machine = llvm.LLVMCreateTargetMachine(target, triple, processor.encode(), feats.encode(),
+                                                       llvm.LLVMCodeGenLevelDefault, llvm.LLVMRelocPIC, llvm.LLVMCodeModelDefault)
+
+    self.pbo = llvm.LLVMCreatePassBuilderOptions()
+    if (opt:=bool(getenv("LLVMOPT", "1"))):
+      self.passes = b'default<O2>'
+      llvm.LLVMPassBuilderOptionsSetLoopUnrolling(self.pbo, True)
+      llvm.LLVMPassBuilderOptionsSetLoopVectorization(self.pbo, True)
+      llvm.LLVMPassBuilderOptionsSetSLPVectorization(self.pbo, True)
+      llvm.LLVMPassBuilderOptionsSetVerifyEach(self.pbo, True)
+    else:
+      self.passes = b'default<O0>'
+
+    self.diag_msgs: list[str] = []
+    @ctypes.CFUNCTYPE(None, llvm.LLVMDiagnosticInfoRef, ctypes.c_void_p)
+    def handle_diag(diag_ref, _arg):
+      severity = llvm.LLVMGetDiagInfoSeverity(diag_ref)
+      msg = ctypes.string_at(llvm.LLVMGetDiagInfoDescription(diag_ref)).decode()
+      if severity == llvm.LLVMDSError:
+        self.diag_msgs.append(msg)
+    self.handle_diag = handle_diag
+    llvm.LLVMContextSetDiagnosticHandler(llvm.LLVMGetGlobalContext(), handle_diag, None)
+    super().__init__(f"compile_llvm_{self.target_arch}{'_jit' if self.jit else ''}{'_opt' if opt else ''}")
+
+  def __del__(self): llvm.LLVMDisposePassBuilderOptions(self.pbo)
+
+  def compile(self, src:str) -> bytes:
+    self.diag_msgs.clear()
+    src_buf = llvm.LLVMCreateMemoryBufferWithMemoryRangeCopy(ctypes.create_string_buffer(src_bytes:=src.encode()), len(src_bytes), b'src')
+    mod = expect(llvm.LLVMParseIRInContext(llvm.LLVMGetGlobalContext(), src_buf, ctypes.pointer(m:=llvm.LLVMModuleRef()), err:=cerr()), err, m)
+    expect(llvm.LLVMVerifyModule(mod, llvm.LLVMReturnStatusAction, err:=cerr()), err)
+    expect(llvm.LLVMRunPasses(mod, self.passes, self.target_machine, self.pbo), 'failed to run passes')
+    if DEBUG >= 7: print(ctypes.string_at(llvm.LLVMPrintModuleToString(mod)).decode())
+    obj_buf = expect(llvm.LLVMTargetMachineEmitToMemoryBuffer(self.target_machine, mod, llvm.LLVMObjectFile, err:=cerr(),
+                                                              ctypes.pointer(buf:=llvm.LLVMMemoryBufferRef())), err, buf)
+    llvm.LLVMDisposeModule(mod)
+    obj = ctypes.string_at(llvm.LLVMGetBufferStart(obj_buf), llvm.LLVMGetBufferSize(obj_buf))
+    llvm.LLVMDisposeMemoryBuffer(obj_buf)
+    if self.diag_msgs: raise RuntimeError("llvm diagnostic: " + "\n".join(self.diag_msgs))
+    return jit_loader(obj) if self.jit else obj
+
+  def disassemble(self, lib:bytes): capstone_flatdump(lib)
+
+class CPULLVMCompiler(LLVMCompiler):
+  def __init__(self):
+    # +reserve-x18 here does the same thing as -ffixed-x18 in ops_cpu.py, see comments there for why it's needed on arm osx
+    cpu, feats = ctypes.string_at(llvm.LLVMGetHostCPUName()), (b'+reserve-x18,' if OSX else b'') + ctypes.string_at(llvm.LLVMGetHostCPUFeatures())
+    super().__init__(cpu.decode(), feats.decode())

tinygrad_repo/tinygrad/runtime/support/compiler_cuda.py

@@ -4,7 +4,7 @@ from tinygrad.helpers import to_char_p_p, colored, init_c_var, getenv
 import tinygrad.runtime.autogen.nvrtc as nvrtc
 from tinygrad.device import Compiler, CompileError
 
-PTX, CUDA_PATH = getenv("PTX"), getenv("CUDA_PATH", "")  # PTX shouldn't be here, in fact, it shouldn't exist
+CUDA_PATH = getenv("CUDA_PATH", "")
 
 def _get_bytes(arg, get_str, get_sz, check) -> bytes:
   sz = init_c_var(ctypes.c_size_t(), lambda x: check(get_sz(arg, ctypes.byref(x))))

tinygrad_repo/tinygrad/runtime/support/hcq.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
-from typing import cast, Callable, Type, TypeVar, Generic, Any
+from typing import cast, Callable, Type, TypeVar, Generic, Any, Sequence
 import contextlib, decimal, statistics, time, ctypes, array, os, struct, traceback, collections
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
 from tinygrad.helpers import PROFILE, getenv, to_mv, round_up, ProfileRangeEvent
-from tinygrad.renderer import Renderer
-from tinygrad.device import BufferSpec, Compiler, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent
-from tinygrad.uop.ops import sym_infer, sint, Variable, UOp
+from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, CompilerPairT
+from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
 
 class MMIOInterface:
@@ -28,10 +27,7 @@ class FileIOInterface:
   def __del__(self):
     if hasattr(self, 'fd'): os.close(self.fd)
   def ioctl(self, request, arg): return fcntl.ioctl(self.fd, request, arg)
-  def mmap(self, start, sz, prot, flags, offset):
-    x = libc.mmap(start, sz, prot, flags, self.fd, offset)
-    if x == 0xffffffffffffffff: raise OSError(f"Failed to mmap {sz} bytes at {hex(start)}: {os.strerror(ctypes.get_errno())}")
-    return x
+  def mmap(self, start, sz, prot, flags, offset): return FileIOInterface._mmap(start, sz, prot, flags, self.fd, offset)
   def read(self, size=None, binary=False, offset=None):
     if offset is not None: self.seek(offset)
     with open(self.fd, "rb" if binary else "r", closefd=False) as file: return file.read(size)
@@ -41,11 +37,13 @@ class FileIOInterface:
   def listdir(self): return os.listdir(self.path)
   def seek(self, offset): os.lseek(self.fd, offset, os.SEEK_SET)
   @staticmethod
-  def anon_mmap(start, sz, prot, flags, offset):
-    x = libc.mmap(start, sz, prot, flags, -1, offset)
+  def _mmap(start, sz, prot, flags, fd, offset):
+    x = libc.mmap(start, sz, prot, flags, fd, offset)
     if x == 0xffffffffffffffff: raise OSError(f"Failed to mmap {sz} bytes at {hex(start)}: {os.strerror(ctypes.get_errno())}")
     return x
   @staticmethod
+  def anon_mmap(start, sz, prot, flags, offset): return FileIOInterface._mmap(start, sz, prot, flags, -1, offset)
+  @staticmethod
   def munmap(buf, sz): return libc.munmap(buf, sz)
   @staticmethod
   def exists(path): return os.path.exists(path)
@@ -192,7 +190,7 @@ class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
       if isinstance(val, int): mv[i] = val if mask is None else ((mv[i] & ~mask) | val)
       else: self.mv_sints.append((mv, i, self._new_sym(val), mask))
 
-  def _apply_var_vals(self, var_vals:dict[Variable, int]):
+  def _apply_var_vals(self, var_vals:dict[str, int]):
     resolved_syms = [sym_infer(sym, var_vals) for sym in self.syms]
 
     for off, sym_idx in self.q_sints:
@@ -205,7 +203,7 @@ class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
 
     self._prev_resolved_syms = cast(list[int|None], resolved_syms)
 
-  def submit(self, dev:HCQDeviceType, var_vals:dict[Variable, int]|None=None):
+  def submit(self, dev:HCQDeviceType, var_vals:dict[str, int]|None=None):
     """
     Submits the command queue to a specific device for execution.
 
@@ -360,12 +358,12 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   signal_pool: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
   cpu_devices: list[HCQCompiled] = []
 
-  def __init__(self, device:str, allocator:HCQAllocatorBase, renderer:Renderer, compiler:Compiler, runtime, signal_t:Type[SignalType],
+  def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:Sequence[CompilerPairT], runtime, signal_t:Type[SignalType],
                comp_queue_t:Callable[[], HWQueue], copy_queue_t:Callable[[], HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
     from tinygrad.runtime.graph.hcq import HCQGraph
-    super().__init__(device, allocator, renderer, compiler, runtime, HCQGraph)
+    super().__init__(device, allocator, compilers, runtime, HCQGraph)
 
     # TODO: peer logic is determined based on device name.
     self.peer_group = device.split(":")[0]
@@ -383,15 +381,20 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True))
     self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(self.kernargs_buf.size, wrap=True)
 
+    self.error_state:Exception|None = None # Exception if error is unrecoverable and sync will always fail
+
     if self._is_cpu(): HCQCompiled.cpu_devices.append(self)
 
   def synchronize(self):
+    if self.error_state is not None: raise self.error_state
+
     # If we have any work on CPU devices, need to synchronize them. This is just an optimization to release GIL allowing to finish faster.
     if not self._is_cpu():
       for dev in HCQCompiled.cpu_devices: dev.synchronize()
 
     try: self.timeline_signal.wait(self.timeline_value - 1)
     except RuntimeError as e:
+      self.error_state = e
       if hasattr(self, 'on_device_hang'): self.on_device_hang()
       else: raise e
 
@@ -437,16 +440,21 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     except MemoryError: buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
     return buf, realloced
 
+  def _make_no_iface_error(self, errs:str, err_short:str) -> RuntimeError:
+    # Keep it in a separate function to avoid creating a traceback <-> locals ref cycle
+    e = RuntimeError(f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available")
+    if hasattr(e, "add_note"): e.add_note(errs + err_short)
+    return e
+
   def _select_iface(self, *ifaces:Type):
     errs, err_short = "", ""
     if val:=getenv(f'{type(self).__name__[:-6].upper()}_IFACE', ""): ifaces = tuple(x for x in ifaces if x.__name__.startswith(val.upper()))
     for iface_t in ifaces:
       try: return iface_t(self, self.device_id)
-      except Exception as e: errs, err_short = errs + f"\n{iface_t.__name__}: {traceback.format_exc()}", err_short + f"\n{iface_t.__name__}: {e}"
-    raise RuntimeError(f"{errs}\nNo interface for {type(self).__name__[:-6]}:{self.device_id} is available:{err_short}\n" \
-                       f"\nForce an interface with {type(self).__name__[:-6].upper()}_IFACE={('|'.join(x.__name__[:-5] for x in ifaces))}.")
+      except Exception as e: errs, err_short = errs + f"\n{iface_t.__name__}: {traceback.format_exc()}", err_short + f"\n{iface_t.__name__}: {e}."
+    raise self._make_no_iface_error(errs, err_short)
 
-  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] in ("CPU", "LLVM")
+  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
 
   def finalize(self):
     try: self.synchronize() # Try to finalize device in any case.

tinygrad_repo/tinygrad/runtime/support/memory.py

@@ -77,11 +77,10 @@ class TLSFAllocator:
       if self.lv1_entries[l1] == 0: continue
       for l2 in range(self.lv2(size) if l1 == size.bit_length() else 0, (1 << self.l2_cnt)):
         if len(self.storage[l1][l2]) > 0:
-          nsize = self.blocks[self.storage[l1][l2][0]][0]
-          assert nsize >= size, "block must be larger"
-
           # Block start address.
           start = self.storage[l1][l2][0]
+          nsize = self.blocks[start][0]
+          assert nsize >= size, "block must be larger"
 
           # If request contains alignment, split the block into two parts.
           if (new_start:=round_up(start, align)) != start:

tinygrad_repo/tinygrad/runtime/support/nv/nvdev.py

@@ -118,7 +118,7 @@ class NVDev(PCIDevImplBase):
 
     self.include("src/common/inc/swref/published/turing/tu102/dev_fb.h")
     if self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0:
-      if DEBUG >= 2: print(f"nv {self.devfmt}: WPR2 is up. Issuing a full reset.")
+      if DEBUG >= 2: print(f"nv {self.devfmt}: WPR2 is up. Issuing a full reset.", flush=True)
       System.pci_reset(self.devfmt)
       time.sleep(0.5)
 

tinygrad_repo/tinygrad/runtime/support/system.py

@@ -165,7 +165,7 @@ class PCIIfaceBase:
   def map(self, b:HCQBuffer):
     if b.owner is not None and b.owner._is_cpu():
       System.lock_memory(cast(int, b.va_addr), b.size)
-      paddrs, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, False
+      paddrs, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, True
     elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, PCIIfaceBase):
       paddrs = [(paddr if b.meta.mapping.system else (paddr + ifa.p2p_base_addr), size) for paddr,size in b.meta.mapping.paddrs]
       snooped, uncached = b.meta.mapping.snooped, b.meta.mapping.uncached