From 7b0ce86e2a77513ba8edf0fc5cdbf13bc48fb708 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 23 Dec 2025 18:15:58 -0500 Subject: [PATCH] more early compilers --- tinygrad/renderer/cstyle.py | 38 ++++++++++++++++++++++++++++++++++++ tinygrad/renderer/llvmir.py | 5 +++++ tinygrad/renderer/nir.py | 5 +++++ tinygrad/renderer/ptx.py | 12 ++++++++++++ tinygrad/runtime/ops_amd.py | 9 ++++----- tinygrad/runtime/ops_cpu.py | 5 ++--- tinygrad/runtime/ops_cuda.py | 12 ++++++------ tinygrad/runtime/ops_hip.py | 5 ++--- tinygrad/runtime/ops_nv.py | 11 +++++------ 9 files changed, 79 insertions(+), 23 deletions(-) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 0a26b02b65..1b52ab47c5 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -386,6 +386,18 @@ class CUDARenderer(CStyleLanguage): self.tensor_cores = tc.cuda_sm89 if int(arch[3:]) >= 89 else tc.cuda_sm80 if int(arch[3:]) >= 80 else tc.cuda_sm75 if int(arch[3:]) >= 75 else [] def __reduce__(self): return self.__class__, (self.arch,) +class CUDACUDARenderer(CUDARenderer): + def __init__(self, arch:str): + super().__init__(arch) + from tinygrad.runtime.support.compiler_cuda import CUDACompiler + self.compiler = CUDACompiler(arch) + +class CUDANVCCRenderer(CUDARenderer): + def __init__(self, arch:str): + super().__init__(arch) + from tinygrad.runtime.support.compiler_cuda import NVCCCompiler + self.compiler = NVCCCompiler(arch) + # language options # https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html kernel_typedef = "extern \"C\" __global__ void __launch_bounds__({launch_bounds})" @@ -538,6 +550,32 @@ class AMDRenderer(CStyleLanguage): for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d;\n}""") return super().render_kernel(function_name, kernel, bufs, uops, prefix) +class AMDHIPRenderer(AMDRenderer): + def __init__(self, arch:str): + super().__init__(arch) + from tinygrad.runtime.support.compiler_amd import HIPCompiler + self.compiler = HIPCompiler(arch) + +class AMDHIPCCRenderer(AMDRenderer): + def __init__(self, arch:str): + super().__init__(arch) + from tinygrad.runtime.support.compiler_amd import HIPCCCompiler + self.compiler = HIPCCCompiler(arch) + class NVRenderer(CUDARenderer): device = "NV" + +class NVNVRenderer(NVRenderer): + def __init__(self, arch:str): + super().__init__(arch) + from tinygrad.runtime.support.compiler_cuda import NVCompiler + self.compiler = NVCompiler(arch) + class HIPRenderer(AMDRenderer): device = "HIP" + +class HIPHIPRenderer(HIPRenderer): + def __init__(self, arch:str): + super().__init__(arch) + from tinygrad.runtime.support.compiler_amd import HIPCompiler + self.compiler = HIPCompiler(arch) + class QCOMRenderer(OpenCLRenderer): device = "QCOM" diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 3d5db9c4d1..9f3c8b9f79 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -143,6 +143,9 @@ class LLVMRenderer(Renderer): if AMX: tensor_cores = tc.amx extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast + def __init__(self): + from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler + self.compiler = CPULLVMCompiler() def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops))) def _render_footer(self, uops: list[UOp]) -> str: return 'attributes #0 = { alwaysinline nounwind "no-builtins" "no-trapping-math"="true" }' def _render_fn(self, name:str, args:list[tuple[str,DType]], kernel:list[str], prefix:list[str]|None=None) -> str: @@ -254,7 +257,9 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"'] return 'attributes #0 = { ' + ' '.join(attributes) + ' }' def __init__(self, arch:str): + from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler self.arch = arch + self.compiler = AMDLLVMCompiler(arch) self.tensor_cores = AMDRenderer.get_tensor_cores(arch) self.is_cdna = AMDRenderer.is_cdna(arch) self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))]) diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 4e0c8e859a..be4cffa5ad 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -245,6 +245,11 @@ class LVPRenderer(NIRRenderer): srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz: setattr(self, "param_idx", self.param_idx+sz))(lambda self,b,x,sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo)) + def __init__(self): + from tinygrad.runtime.support.compiler_mesa import LVPCompiler + super().__init__() + self.compiler = LVPCompiler() + def prerender(self, uops:list[UOp]): super().prerender(uops) self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)]) diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index e61fc3eda1..07223f7296 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -240,3 +240,15 @@ class PTXRenderer(Renderer): if u.op is Ops.SPECIAL: kernel = [f".reg .u32 %{u.arg};"] + kernel return self.render_kernel(kernel, name, bufs, c.items(), uops) + +class CUDAPTXRenderer(PTXRenderer): + def __init__(self, arch:str): + super().__init__(arch, "CUDA") + from tinygrad.runtime.support.compiler_cuda import PTXCompiler + self.compiler = PTXCompiler(arch) + +class NVPTXRenderer(PTXRenderer): + def __init__(self, arch:str): + super().__init__(arch, "NV") + from tinygrad.runtime.support.compiler_cuda import NVPTXCompiler + self.compiler = NVPTXCompiler(arch) diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index da824c42c9..3c6f9cd163 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -9,11 +9,10 @@ from tinygrad.uop.ops import sint from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet, CompilerPair from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, ceildiv -from tinygrad.renderer.cstyle import AMDRenderer +from tinygrad.renderer.cstyle import AMDHIPRenderer, AMDHIPCCRenderer from tinygrad.renderer.llvmir import AMDLLVMRenderer from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt from tinygrad.runtime.autogen.am import am -from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler, AMDLLVMCompiler from tinygrad.runtime.support.elf import elf_loader from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets, import_pmc @@ -931,9 +930,9 @@ class AMDDevice(HCQCompiled): max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000 self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20)) - compilers = CompilerSet([CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch)), - CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch), AMD_LLVM), - CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCCCompiler, self.arch))], ctrl_var=AMD_CC) + compilers = CompilerSet([CompilerPair(functools.partial(AMDHIPRenderer, self.arch), None), + CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), None, AMD_LLVM), + CompilerPair(functools.partial(AMDHIPCCRenderer, self.arch), None)], ctrl_var=AMD_CC) super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal, functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self), diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 2e76328e22..dc6ea8b29a 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -8,7 +8,6 @@ from tinygrad.runtime.support.hcq import CLikeArgsState from tinygrad.renderer.cstyle import ClangJITRenderer from tinygrad.renderer.llvmir import LLVMRenderer from tinygrad.renderer.nir import LVPRenderer -from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler from tinygrad.runtime.support.compiler_mesa import LVPCompiler from tinygrad.runtime.support.elf import jit_loader from tinygrad.uop.ops import sint @@ -136,6 +135,6 @@ class CPUDevice(HCQCompiled): def __init__(self, device:str=""): self.tasks:queue.Queue = queue.Queue() CPUWorker(self, self.tasks, thread_id=0).start() - compilers = CompilerSet([CompilerPair(ClangJITRenderer, None), CompilerPair(LLVMRenderer, CPULLVMCompiler, ctrl_var=CPU_LLVM), - CompilerPair(LVPRenderer, LVPCompiler, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC) + compilers = CompilerSet([CompilerPair(ClangJITRenderer, None), CompilerPair(LLVMRenderer, None, ctrl_var=CPU_LLVM), + CompilerPair(LVPRenderer, None, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC) super().__init__(device, CPUAllocator(self), compilers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue) diff --git a/tinygrad/runtime/ops_cuda.py b/tinygrad/runtime/ops_cuda.py index 2dcbad04d2..1857cd88ee 100644 --- a/tinygrad/runtime/ops_cuda.py +++ b/tinygrad/runtime/ops_cuda.py @@ -2,10 +2,10 @@ from __future__ import annotations import ctypes, functools from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing, CUDA_CC, CUDA_PTX from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerPair, CompilerSet -from tinygrad.renderer.cstyle import CUDARenderer -from tinygrad.renderer.ptx import PTXRenderer +from tinygrad.renderer.cstyle import CUDACUDARenderer, CUDANVCCRenderer +from tinygrad.renderer.ptx import CUDAPTXRenderer from tinygrad.runtime.autogen import cuda -from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler, NVCCCompiler +from tinygrad.runtime.support.compiler_cuda import pretty_ptx if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl # noqa: F401 # pylint: disable=unused-import if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported @@ -117,9 +117,9 @@ class CUDADevice(Compiled): CUDADevice.devices.append(self) from tinygrad.runtime.graph.cuda import CUDAGraph - compilers = CompilerSet([CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(CUDACompiler, self.arch)), - CompilerPair(functools.partial(PTXRenderer, self.arch), functools.partial(PTXCompiler, self.arch), CUDA_PTX), - CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(NVCCCompiler, self.arch))], ctrl_var=CUDA_CC) + compilers = CompilerSet([CompilerPair(functools.partial(CUDACUDARenderer, self.arch), None), + CompilerPair(functools.partial(CUDAPTXRenderer, self.arch), None, CUDA_PTX), + CompilerPair(functools.partial(CUDANVCCRenderer, self.arch), None)], ctrl_var=CUDA_CC) super().__init__(device, CUDAAllocator(self), compilers, functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph) def synchronize(self): diff --git a/tinygrad/runtime/ops_hip.py b/tinygrad/runtime/ops_hip.py index 02fe7d9a36..6deed062ae 100644 --- a/tinygrad/runtime/ops_hip.py +++ b/tinygrad/runtime/ops_hip.py @@ -2,8 +2,7 @@ import ctypes, functools from tinygrad.helpers import init_c_var, mv_address, init_c_struct_t, getenv from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet, CompilerPair from tinygrad.runtime.autogen import hip -from tinygrad.runtime.support.compiler_amd import HIPCompiler -from tinygrad.renderer.cstyle import HIPRenderer +from tinygrad.renderer.cstyle import HIPHIPRenderer if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl # noqa: F401 # pylint: disable=unused-import def check(status): @@ -15,7 +14,7 @@ class HIPDevice(Compiled): self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode() self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t(), lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)] - compilers = CompilerSet([CompilerPair(functools.partial(HIPRenderer, self.arch), functools.partial(HIPCompiler, self.arch))]) + compilers = CompilerSet([CompilerPair(functools.partial(HIPHIPRenderer, self.arch), None)]) super().__init__(device, HIPAllocator(self), compilers, functools.partial(HIPProgram, self)) def synchronize(self): check(hip.hipSetDevice(self.device_id)) diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index 1da4fd6217..1a7f900e65 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -8,9 +8,8 @@ from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU from tinygrad.uop.ops import sint from tinygrad.device import BufferSpec, CompilerPair, CompilerSet from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, NV_CC, NV_PTX, NV_NAK -from tinygrad.renderer.ptx import PTXRenderer -from tinygrad.renderer.cstyle import NVRenderer -from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler +from tinygrad.renderer.ptx import CUDAPTXRenderer, NVPTXRenderer +from tinygrad.renderer.cstyle import NVNVRenderer, CUDACUDARenderer from tinygrad.runtime.support.compiler_mesa import NAKCompiler from tinygrad.runtime.autogen import nv_570, nv_580, pci, mesa from tinygrad.runtime.support.elf import elf_loader @@ -583,9 +582,9 @@ class NVDevice(HCQCompiled[HCQSignal]): self.arch: str = "sm_120" if self.sm_version==0xa04 else f"sm_{(self.sm_version>>8)&0xff}{(val>>4) if (val:=self.sm_version&0xff) > 0xf else val}" self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf) - cucc, ptxcc = (CUDACompiler, PTXCompiler) if MOCKGPU else (NVCompiler, NVPTXCompiler) - compilers = CompilerSet(ctrl_var=NV_CC, cset=[CompilerPair(functools.partial(NVRenderer, self.arch),functools.partial(cucc, self.arch)), - CompilerPair(functools.partial(PTXRenderer, self.arch, device="NV"), functools.partial(ptxcc, self.arch), NV_PTX), + nvr, ptxr = (CUDACUDARenderer, CUDAPTXRenderer) if MOCKGPU else (NVNVRenderer, NVPTXRenderer) + compilers = CompilerSet(ctrl_var=NV_CC, cset=[CompilerPair(functools.partial(nvr, self.arch), None), + CompilerPair(functools.partial(ptxr, self.arch), None, NV_PTX), CompilerPair(functools.partial(NAKRenderer, dev=self), functools.partial(NAKCompiler, self.arch, self.max_warps_per_sm), NV_NAK)]) super().__init__(device, NVAllocator(self), compilers, functools.partial(NVProgram, self), HCQSignal, NVComputeQueue, NVCopyQueue)