diff --git a/tinygrad/device.py b/tinygrad/device.py index e7cae72e12..85742adf3f 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -2,7 +2,7 @@ from __future__ import annotations from dataclasses import dataclass, replace from collections import defaultdict from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING -import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal +import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up @@ -310,6 +310,14 @@ class Compiler: if self.cachekey is not None: diskcache_put(self.cachekey, src, lib) return lib def disassemble(self, lib:bytes): pass + def server(self, cmd:str, arch:str, *args) -> subprocess.Popen: + argv = f"{cmd} {pathlib.Path(__file__).parent}/runtime/support/compileserver.py {type(self).__module__}:{type(self).__name__} {arch}" + return subprocess.Popen(argv.split() + [str(a) for a in args], stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0) + def compile_server(self, src:str, proc:subprocess.Popen) -> bytes: + unwrap(proc.stdin).write(struct.pack("I", len(src.encode())) + src.encode()) + if (lib:=unwrap(proc.stdout).read(struct.unpack("I", unwrap(proc.stdout).read(4))[0])): return lib + raise CompileError("Compilation Error") + @dataclass class TinyELF: diff --git a/tinygrad/runtime/support/compiler_cuda.py b/tinygrad/runtime/support/compiler_cuda.py index 1ba245a9c0..81124b1169 100644 --- a/tinygrad/runtime/support/compiler_cuda.py +++ b/tinygrad/runtime/support/compiler_cuda.py @@ -1,10 +1,12 @@ import hashlib, tempfile, ctypes, re, pathlib -from tinygrad.helpers import to_char_p_p, colored, getenv, system +from tinygrad.helpers import to_char_p_p, colored, getenv, system, OSX from tinygrad.runtime.support.c import init_c_var from tinygrad.runtime.autogen import nvrtc, nvjitlink as jitlink from tinygrad.device import Compiler, CompileError CUDA_PATH = getenv("CUDA_PATH", "") +root = pathlib.Path(__file__).parents[3] +osx_docker_cmd = f"docker run --rm -i -v {root}:{root} -e PYTHONPATH={root} ghcr.io/tinygrad/cuda-arm64:v2.3" def _get_bytes(arg, get_str, get_sz, check) -> bytes: x = ctypes.create_string_buffer(init_c_var(ctypes.c_size_t, lambda x: check(get_sz(arg, ctypes.byref(x)))).value) @@ -44,11 +46,14 @@ def cuda_disassemble(lib:bytes, arch:str, ptx=False): class NVRTCCompiler(Compiler): def __init__(self, arch:str, ptx=True, cache_key:str="cuda"): self.ptx, self.arch, self.compile_options = ptx, arch, [f'--gpu-architecture={arch}'] - self.compile_options += [f"-I{CUDA_PATH}/include"] if CUDA_PATH else ["-I/usr/local/cuda/include", "-I/usr/include", "-I/opt/cuda/include"] - nvrtc_check(nvrtc.nvrtcVersion((nvrtcMajor := ctypes.c_int()), (nvrtcMinor := ctypes.c_int()))) - if (nvrtcMajor.value, nvrtcMinor.value) >= (12, 4): self.compile_options.append("--minimal") + if OSX: self.compiler_process = self.server(osx_docker_cmd, arch, ptx) + else: + self.compile_options += [f"-I{CUDA_PATH}/include"] if CUDA_PATH else ["-I/usr/local/cuda/include", "-I/usr/include", "-I/opt/cuda/include"] + nvrtc_check(nvrtc.nvrtcVersion((nvrtcMajor := ctypes.c_int()), (nvrtcMinor := ctypes.c_int()))) + if (nvrtcMajor.value, nvrtcMinor.value) >= (12, 4): self.compile_options.append("--minimal") super().__init__(f"compile_{cache_key}_{self.arch}") def compile(self, src:str) -> bytes: + if OSX: return self.compile_server(src, self.compiler_process) nvrtc_check(nvrtc.nvrtcCreateProgram(ctypes.byref(prog := nvrtc.nvrtcProgram()), src.encode(), "".encode(), 0, None, None)) nvrtc_check(nvrtc.nvrtcCompileProgram(prog, len(self.compile_options), to_char_p_p([o.encode() for o in self.compile_options])), prog) data = _get_bytes(prog, nvrtc.nvrtcGetPTX if self.ptx else nvrtc.nvrtcGetCUBIN, @@ -80,9 +85,11 @@ class PTXCompiler(Compiler): class NVPTXCompiler(PTXCompiler): def __init__(self, arch:str): - jitlink_check(jitlink.nvJitLinkVersion(ctypes.byref(ctypes.c_uint()), ctypes.byref(ctypes.c_uint()))) + if OSX: self.compiler_process = self.server(osx_docker_cmd, arch) + else: jitlink_check(jitlink.nvJitLinkVersion(ctypes.byref(ctypes.c_uint()), ctypes.byref(ctypes.c_uint()))) super().__init__(arch, cache_key="nv_ptx") def compile(self, src:str) -> bytes: + if OSX: return self.compile_server(src, self.compiler_process) jitlink_check(jitlink.nvJitLinkCreate(handle := jitlink.nvJitLinkHandle(), 1, to_char_p_p([f'-arch={self.arch}'.encode()])), handle) jitlink_check(jitlink.nvJitLinkAddData(handle, jitlink.NVJITLINK_INPUT_PTX, ptxsrc:=super().compile(src), len(ptxsrc), "".encode()), handle) jitlink_check(jitlink.nvJitLinkComplete(handle), handle) diff --git a/tinygrad/runtime/support/compiler_qcom.py b/tinygrad/runtime/support/compiler_qcom.py index 0c4c3bff41..64a7cba8d6 100644 --- a/tinygrad/runtime/support/compiler_qcom.py +++ b/tinygrad/runtime/support/compiler_qcom.py @@ -1,6 +1,6 @@ -import ctypes, struct, platform, pathlib, shutil, subprocess, sys, tarfile, tempfile +import ctypes, struct, platform, pathlib, shutil, tarfile, tempfile from tinygrad.device import Compiler -from tinygrad.helpers import DEBUG, system, fetch, unwrap +from tinygrad.helpers import DEBUG, system, fetch from tinygrad.runtime.support.compiler_mesa import disas_adreno # see https://github.com/sirhcm/tinydreno from tinygrad.runtime.autogen import llvm_qcom @@ -12,12 +12,11 @@ class QCOMCompiler(Compiler): assert arch.split(',')[0] == "a630", "only a630 supported" if platform.machine() == "aarch64": self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance() else: - self.arch, self.chip_id, self.fs = arch, 0x6030001, tempfile.TemporaryDirectory() + self.arch, self.chip_id, self.fs, root = arch, 0x6030001, tempfile.TemporaryDirectory(), pathlib.Path(__file__).parents[3] with tarfile.open(fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz')) as t: t.extractall(fs:=self.fs.name) - if (qemu:=shutil.which("qemu-aarch64-static")): argv = f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3 {__file__} {arch}" - else: argv = (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {pathlib.Path(__file__).parents[2]}:/tinygrad " - f"-e PYTHONPATH=/ -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3 /tinygrad/runtime/support/compiler_qcom.py {arch}") - self.compiler_process = subprocess.Popen(argv.split(), stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0) + self.compiler_process = self.server(f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3" if (qemu:=shutil.which("qemu-aarch64-static")) + else (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {root}:{root} " + f"-e PYTHONPATH={root} -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3"), arch) super().__init__(f"compile_qcomcl_{arch}") def __del__(self): llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst) if platform.machine() == "aarch64" else self.compiler_process.kill() @@ -32,10 +31,7 @@ class QCOMCompiler(Compiler): return handle def compile(self, src) -> bytes: - if platform.machine() != "aarch64": - unwrap(self.compiler_process.stdin).write(struct.pack("I", len(src.encode())) + src.encode()) - if (lib:=unwrap(self.compiler_process.stdout).read(struct.unpack("I", unwrap(self.compiler_process.stdout).read(4))[0])): return lib - raise RuntimeError("QCOM Compilation Error") + if platform.machine() != "aarch64": return self.compile_server(src, self.compiler_process) ch = self.checked(llvm_qcom.cl_compiler_compile_source(self.llvm_inst, self.chip_id, llvm_qcom.CL_MODE_64BIT, b"", 0, 0, 0, src.encode(), 0, llvm_qcom.CL_SRC_STR, None)) if DEBUG >= 8: print(system("llvm-dis", input=ctypes.string_at((comp:=ch.contents.compiled.contents).llvm_bitcode, comp.llvm_bitcode_size))) @@ -48,12 +44,3 @@ class QCOMCompiler(Compiler): def disassemble(self, lib: bytes): disas_adreno(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)], self.chip_id) -if __name__ == "__main__": - compiler = QCOMCompiler(sys.argv[1]) - while (amt:=sys.stdin.buffer.read(4)): - try: lib = compiler.compile(sys.stdin.buffer.read(struct.unpack("I", amt)[0]).decode()) - except Exception as e: - lib = b"" - print(e, file=sys.stderr, flush=True) - sys.stdout.buffer.write(struct.pack("I", len(lib)) + lib) - sys.stdout.buffer.flush() diff --git a/tinygrad/runtime/support/compileserver.py b/tinygrad/runtime/support/compileserver.py new file mode 100644 index 0000000000..b41bd7ab6a --- /dev/null +++ b/tinygrad/runtime/support/compileserver.py @@ -0,0 +1,13 @@ +import ast, struct, sys +from tinygrad.helpers import fromimport + +if __name__ == "__main__": + assert len(sys.argv) >= 3, f"usage: {sys.argv[0]} []" + compiler = fromimport(*sys.argv[1].split(':'))(sys.argv[2], *(ast.literal_eval(arg) for arg in sys.argv[3:])) + while (amt:=sys.stdin.buffer.read(4)): + try: lib = compiler.compile(sys.stdin.buffer.read(struct.unpack("I", amt)[0]).decode()) + except Exception as e: + lib = b"" + print(e, file=sys.stderr, flush=True) + sys.stdout.buffer.write(struct.pack("I", len(lib)) + lib) + sys.stdout.buffer.flush()