compileonce

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, subprocess, struct, multiprocessing
+import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
 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,16 +310,6 @@ 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}"
-    # the lock is created with the server and travels with the process in a spawn/fork, so compiles to the server stay serialized
-    self.compile_server_lock = multiprocessing.Lock()
-    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:
-    with self.compile_server_lock:
-      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

tinygrad/runtime/support/compileonce.py

@@ -0,0 +1,15 @@
+import ast, subprocess, sys
+from tinygrad.device import CompileError
+from tinygrad.helpers import fromimport
+
+# run argv as a one shot compiler: src goes to stdin, the compiled binary comes back on stdout (see __main__ below)
+def compile_once(argv:list[str], src:str, env:dict[str,str]|None=None) -> bytes:
+  with subprocess.Popen(argv, stdin=subprocess.PIPE, stdout=subprocess.PIPE, env=env) as p:
+    ret, _ = p.communicate(src.encode())
+    if p.returncode != 0: raise CompileError(f"Compilation Error: {' '.join(argv)}")
+  return ret
+
+if __name__ == "__main__":
+  assert len(sys.argv) >= 3, f"usage: {sys.argv[0]} <compiler> <arch> [<args>]"
+  compiler = fromimport(*sys.argv[1].split(':'))(sys.argv[2], *(ast.literal_eval(arg) for arg in sys.argv[3:]))
+  sys.stdout.buffer.write(compiler.compile(sys.stdin.buffer.read().decode()))

tinygrad/runtime/support/compiler_cuda.py

@@ -1,12 +1,16 @@
 import hashlib, tempfile, ctypes, re, pathlib
 from tinygrad.helpers import to_char_p_p, colored, getenv, system, OSX
 from tinygrad.runtime.support.c import init_c_var
+from tinygrad.runtime.support.compileonce import compile_once
 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"
+# one shot compile in a container, no server to lifecycle manage (see compileonce.py)
+def osx_compiler_cmd(compiler:str, *args) -> list[str]:
+  root = pathlib.Path(__file__).parents[3]
+  return (f"docker run --rm -i -v {root}:{root} -e PYTHONPATH={root} ghcr.io/tinygrad/cuda-arm64:v2.3 "
+          f"python3 {pathlib.Path(__file__).parent}/compileonce.py {compiler}").split() + [str(a) for a in args]
 
 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)
@@ -46,14 +50,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}']
-    if OSX: self.compiler_process = self.server(osx_docker_cmd, arch, ptx)
+    if OSX: self.compiler_cmd = osx_compiler_cmd(f"{NVRTCCompiler.__module__}:NVRTCCompiler", 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)
+    if OSX: return compile_once(self.compiler_cmd, src)
     nvrtc_check(nvrtc.nvrtcCreateProgram(ctypes.byref(prog := nvrtc.nvrtcProgram()), src.encode(), "<null>".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,
@@ -85,11 +89,11 @@ class PTXCompiler(Compiler):
 
 class NVPTXCompiler(PTXCompiler):
   def __init__(self, arch:str):
-    if OSX: self.compiler_process = self.server(osx_docker_cmd, arch)
+    if OSX: self.compiler_cmd = osx_compiler_cmd(f"{NVPTXCompiler.__module__}:NVPTXCompiler", 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)
+    if OSX: return compile_once(self.compiler_cmd, src)
     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), "<null>".encode()), handle)
     jitlink_check(jitlink.nvJitLinkComplete(handle), handle)

tinygrad/runtime/support/compiler_qcom.py

@@ -1,6 +1,7 @@
-import ctypes, struct, platform, pathlib, shutil, tarfile, tempfile
+import ctypes, os, struct, platform, pathlib, shutil, tarfile, tempfile
 from tinygrad.device import Compiler
 from tinygrad.helpers import DEBUG, system, fetch
+from tinygrad.runtime.support.compileonce import compile_once
 from tinygrad.runtime.support.compiler_mesa import disas_adreno
 # see https://github.com/sirhcm/tinydreno
 from tinygrad.runtime.autogen import llvm_qcom
@@ -10,16 +11,24 @@ def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
 class QCOMCompiler(Compiler):
   def __init__(self, arch:str):
     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()
+    self.arch, self.chip_id = arch, 0x6030001
+    if platform.machine() == "aarch64": self.llvm_inst = llvm_qcom.cl_compiler_create_llvm_instance()
     else:
-      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)
-      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)
+      self.fs, root = 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(self.fs.name)
+      once = f"{pathlib.Path(__file__).parent}/compileonce.py {QCOMCompiler.__module__}:QCOMCompiler {arch}"
+      self.compiler_env = None
+      # the compiler is aarch64 only, run it emulated. for qemu user mode PYTHONPATH must point at the tinygrad on the host
+      if (qemu:=shutil.which("qemu-aarch64-static")):
+        self.compiler_cmd = f"{qemu} -cpu max,pauth=off -L {self.fs.name} {self.fs.name}/usr/bin/python3 {once}".split()
+        self.compiler_env = {**os.environ, "PYTHONPATH": str(root)}
+      else:
+        self.compiler_cmd = (f"docker run --rm -i --platform linux/aarch64 -v {self.fs.name}/usr:/usr -v {root}:{root} "
+                             f"-e PYTHONPATH={root} -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3 {once}").split()
     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()
+  def __del__(self):
+    if platform.machine() == "aarch64": llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst)
 
   def __reduce__(self): return QCOMCompiler, (self.arch,)
 
@@ -31,7 +40,7 @@ class QCOMCompiler(Compiler):
     return handle
 
   def compile(self, src) -> bytes:
-    if platform.machine() != "aarch64": return self.compile_server(src, self.compiler_process)
+    if platform.machine() != "aarch64": return compile_once(self.compiler_cmd, src, self.compiler_env)
     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)))
@@ -43,4 +52,3 @@ class QCOMCompiler(Compiler):
     return ret
 
   def disassemble(self, lib: bytes): disas_adreno(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)], self.chip_id)
-

tinygrad/runtime/support/compileserver.py

@@ -1,13 +0,0 @@
-import ast, struct, sys
-from tinygrad.helpers import fromimport
-
-if __name__ == "__main__":
-  assert len(sys.argv) >= 3, f"usage: {sys.argv[0]} <compiler> <arch> [<args>]"
-  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()