sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions Activity

more early compilers

Browse Source
This commit is contained in:
geohot
2025-12-23 18:15:58 -05:00
parent 43c6e973d8
commit 7b0ce86e2a
9 changed files with 79 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/renderer/cstyle.py
+38
View File
@@ -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"
tinygrad/renderer/llvmir.py
+5
View File
@@ -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))])
tinygrad/renderer/nir.py
+5
View File
@@ -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)])
tinygrad/renderer/ptx.py
+12
View File
@@ -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)
tinygrad/runtime/ops_amd.py
+4 -5
View File
@@ -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),
tinygrad/runtime/ops_cpu.py
+2 -3
View File
@@ -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)
tinygrad/runtime/ops_cuda.py
+6 -6
View File
@@ -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):
tinygrad/runtime/ops_hip.py
+2 -3
View File
@@ -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))
tinygrad/runtime/ops_nv.py
+5 -6
View File
@@ -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)