tinygrad 0.8.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

tinygrad/runtime/ops_llvm.py

@@ -1,66 +1,46 @@
-import ctypes
-from typing import ClassVar, Tuple
-from tinygrad.device import Compiled, MallocAllocator
-from tinygrad.helpers import getenv, DEBUG, cpu_time_execution
-from ctypes import CFUNCTYPE
-from tinygrad.codegen.kernel import LinearizerOptions
-from tinygrad.renderer.llvmir import uops_to_llvm_ir
-
+from __future__ import annotations
+import ctypes, functools
+from typing import Tuple
+from tinygrad.device import Compiled, Compiler, MallocAllocator
+from tinygrad.helpers import DEBUG, cpu_time_execution, cpu_objdump
+from tinygrad.renderer.llvmir import LLVMRenderer
 import llvmlite.binding as llvm
 
-LLVMOPT = bool(getenv("LLVMOPT"))
+class LLVMCompiler(Compiler):
+  def __init__(self, device:LLVMDevice):
+    self.device = device
+    super().__init__("compile_llvm")
+  def compile(self, src:str) -> bytes:
+    mod = llvm.parse_assembly(src)
+    mod.verify()
+    self.device.optimizer.run(mod)
+    if DEBUG >= 5: print(self.device.target_machine.emit_assembly(mod))
+    return self.device.target_machine.emit_object(mod)
 
-class LLVM:
-  target_machine: ClassVar[llvm.targets.TargetMachine] = None
-  engine: ClassVar[llvm.executionengine.ExecutionEngine] = None
-  optimizer: ClassVar[llvm.passmanagers.ModulePassManager] = None
+class LLVMProgram:
+  def __init__(self, device:LLVMDevice, name:str, lib:bytes):
+    if DEBUG >= 6: cpu_objdump(lib)
+    self.name, self.lib = name, lib
+    device.engine.add_object_file(llvm.object_file.ObjectFileRef.from_data(lib))
+    self.fxn = device.engine.get_function_address(name)
 
-  def __init__(self):
-    if LLVM.engine is not None: return
+  def __call__(self, *bufs, vals:Tuple[int, ...]=(), wait=False):
+    if not hasattr(self, 'cfunc'):
+      self.cfunc = ctypes.CFUNCTYPE(ctypes.c_int, *([ctypes.c_void_p]*len(bufs)), *([ctypes.c_int32]*len(vals)))(self.fxn)
+    return cpu_time_execution(lambda: self.cfunc(*bufs, *vals), enable=wait)
+
+class LLVMDevice(Compiled):
+  def __init__(self, device:str):
     llvm.initialize()
     llvm.initialize_native_target()
     llvm.initialize_native_asmprinter()
     llvm.initialize_native_asmparser()
-    target = llvm.Target.from_triple(llvm.get_process_triple())
-    LLVM.optimizer = llvm.create_module_pass_manager()
-    LLVM.target_machine = target.create_target_machine(opt=2)  # this opt actually can change things. ex: opt=3 means no FMA, opt=2 means FMA
-    LLVM.target_machine.add_analysis_passes(LLVM.optimizer)
-
-    # TODO: this makes compile times so much faster
-    if LLVMOPT:
-      llvm.set_option(str(), '-force-vector-interleave=4')  # this makes sum the same speed as torch, it also doubles the (slow) conv speed
-      if DEBUG >= 4: llvm.set_option(str(), '--debug-only=loop-vectorize')
-      #llvm.set_option(str(), '--debug')
-
-      # does this do anything?
-      builder = llvm.create_pass_manager_builder()
-      builder.opt_level = 3
-      builder.size_level = 0
-      builder.loop_vectorize = True
-      builder.slp_vectorize = True
-      builder.populate(LLVM.optimizer)
-
-    LLVM.target_machine.set_asm_verbosity(True)
+    self.optimizer: llvm.passmanagers.ModulePassManager = llvm.create_module_pass_manager()
+    # this opt actually can change things. ex: opt=3 means no FMA, opt=2 means FMA
+    self.target_machine: llvm.targets.TargetMachine = llvm.Target.from_triple(llvm.get_process_triple()).create_target_machine(opt=2)
+    self.target_machine.add_analysis_passes(self.optimizer)
+    self.target_machine.set_asm_verbosity(True)
     backing_mod = llvm.parse_assembly(str())
     backing_mod.triple = llvm.get_process_triple()
-    LLVM.engine = llvm.create_mcjit_compiler(backing_mod, LLVM.target_machine)
-
-def compile_llvm(prg) -> bytes:
-  mod = llvm.parse_assembly(prg)
-  mod.verify()
-  LLVM().optimizer.run(mod)
-  if DEBUG >= 5: print(LLVM.target_machine.emit_assembly(mod))
-  return LLVM.target_machine.emit_object(mod)
-
-class LLVMProgram:
-  def __init__(self, name:str, lib:bytes):
-    self.name, self.lib = name, lib
-    LLVM().engine.add_object_file(llvm.object_file.ObjectFileRef.from_data(lib))
-    self.fxn = LLVM.engine.get_function_address(name)
-
-  def __call__(self, *bufs, vals:Tuple[int, ...]=(), wait=False):
-    self.cfunc = CFUNCTYPE(ctypes.c_int, *([ctypes.c_void_p]*len(bufs)), *([ctypes.c_int32]*len(vals)))(self.fxn)
-    return cpu_time_execution(lambda: self.cfunc(*bufs, *vals), enable=wait)
-
-LLVMDevice = Compiled(MallocAllocator, LinearizerOptions(supports_float4=False, has_local=False, has_shared=False),
-                      uops_to_llvm_ir, compile_llvm, LLVMProgram)
+    self.engine: llvm.executionengine.ExecutionEngine = llvm.create_mcjit_compiler(backing_mod, self.target_machine)
+    super().__init__(device, MallocAllocator, LLVMRenderer(), LLVMCompiler(self), functools.partial(LLVMProgram, self))

tinygrad/runtime/ops_metal.py

@@ -1,21 +1,30 @@
 from __future__ import annotations
 import os, subprocess, pathlib, ctypes, tempfile, functools
 import Metal, libdispatch
-from typing import List, Any, Tuple, Optional
-from tinygrad.codegen.kernel import LinearizerOptions
+from typing import List, Set, Any, Tuple, Optional
 from tinygrad.helpers import prod, getenv, DEBUG, unwrap2
-from tinygrad.device import Compiled, LRUAllocator
+from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator
 from tinygrad.renderer.cstyle import MetalRenderer
 
-def compile_metal(prg, use_xcode=bool(getenv("METAL_XCODE"))) -> bytes:
-  assert MetalDevice.compiler_device, "metal device creation is required for metal compile"
-  if use_xcode:
-    # NOTE: if you run llvm-dis on "air" you can see the llvm bytecode
-    air = subprocess.check_output(['xcrun', '-sdk', 'macosx', 'metal', '-x', 'metal', '-c', '-', '-o', '-'], input=prg.encode('utf-8'))
-    return subprocess.check_output(['xcrun', '-sdk', 'macosx', 'metallib', '-', '-o', '-'], input=air)
-  options = Metal.MTLCompileOptions.new()
-  library = unwrap2(MetalDevice.compiler_device.newLibraryWithSource_options_error_(prg, options, None))
-  return library.libraryDataContents().bytes().tobytes()
+def wait_check(cbuf: Any):
+  cbuf.waitUntilCompleted()
+  if (error := cbuf.error()) is not None:
+    raise RuntimeError(error)
+
+class MetalCompiler(Compiler):
+  def __init__(self, device:Optional[MetalDevice]):
+    self.device = device
+    super().__init__("compile_metal")
+  def compile(self, src:str) -> bytes:
+    if self.device is None:
+      # NOTE: if you run llvm-dis on "air" you can see the llvm bytecode
+      air = subprocess.check_output(['xcrun', '-sdk', 'macosx', 'metal', '-x', 'metal', '-c', '-', '-o', '-'], input=src.encode('utf-8'))
+      return subprocess.check_output(['xcrun', '-sdk', 'macosx', 'metallib', '-', '-o', '-'], input=air)
+    options = Metal.MTLCompileOptions.new()
+    options.setFastMathEnabled_(getenv("METAL_FAST_MATH"))
+    try: library = unwrap2(self.device.device.newLibraryWithSource_options_error_(src, options, None))
+    except AssertionError as e: raise CompileError(e) from e
+    return library.libraryDataContents().bytes().tobytes()
 
 class MetalProgram:
   def __init__(self, device:MetalDevice, name:str, lib:bytes):
@@ -24,14 +33,15 @@ class MetalProgram:
       with tempfile.NamedTemporaryFile(delete=True) as shader:
         shader.write(lib)
         shader.flush()
-        os.system(f"cd {pathlib.Path(__file__).parents[2]}/disassemblers/applegpu && python3 compiler_explorer.py {shader.name}")
+        os.system(f"cd {pathlib.Path(__file__).parents[2]}/extra/disassemblers/applegpu && python3 compiler_explorer.py {shader.name}")
+    assert lib[:4] == b"MTLB", "Invalid Metal library. Could be due to using conda. Try system python or METAL_XCODE=1 DISABLE_COMPILER_CACHE=1."
     data = libdispatch.dispatch_data_create(lib, len(lib), None, None)
     self.library = unwrap2(self.device.device.newLibraryWithData_error_(data, None))
     self.fxn = self.library.newFunctionWithName_(name)
     self.pipeline_state = unwrap2(self.device.device.newComputePipelineStateWithFunction_error_(self.fxn, None))
 
-  def __call__(self, *bufs, global_size:Tuple[int,int,int], local_size:Tuple[int,int,int], vals:Tuple[int, ...]=(), wait=False):
-    assert prod(local_size) <= self.pipeline_state.maxTotalThreadsPerThreadgroup(),f"local size {local_size} bigger than {self.pipeline_state.maxTotalThreadsPerThreadgroup()} with exec width {self.pipeline_state.threadExecutionWidth()} memory length {self.pipeline_state.staticThreadgroupMemoryLength()}"  # noqa: E501
+  def __call__(self, *bufs, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1), vals:Tuple[int, ...]=(), wait=False):
+    if prod(local_size) > self.pipeline_state.maxTotalThreadsPerThreadgroup(): raise RuntimeError(f"local size {local_size} bigger than {self.pipeline_state.maxTotalThreadsPerThreadgroup()} with exec width {self.pipeline_state.threadExecutionWidth()} memory length {self.pipeline_state.staticThreadgroupMemoryLength()}")  # noqa: E501
     command_buffer = self.device.mtl_queue.commandBuffer()
     encoder = command_buffer.computeCommandEncoder()
     encoder.setComputePipelineState_(self.pipeline_state)
@@ -41,19 +51,26 @@ class MetalProgram:
     encoder.endEncoding()
     command_buffer.commit()
     if wait:
-      command_buffer.waitUntilCompleted()
+      wait_check(command_buffer)
       return command_buffer.GPUEndTime() - command_buffer.GPUStartTime()
     self.device.mtl_buffers_in_flight.append(command_buffer)
 
 class MetalAllocator(LRUAllocator):
   def __init__(self, device:MetalDevice):
     self.device:MetalDevice = device
+    self.track_cross_device: Set[MetalDevice] = set()
     super().__init__()
-  def _alloc(self, size:int) -> Any:
+  def free_cache(self):
+    self.device.synchronize()
+    for x in self.track_cross_device: x.synchronize()
+    self.track_cross_device.clear()
+    return super().free_cache()
+  def _alloc(self, size:int, options) -> Any:
     ret = self.device.device.newBufferWithLength_options_(size, Metal.MTLResourceStorageModeShared)
     if ret is None: raise MemoryError(f"Metal OOM while allocating {size=}")
     return ret
-  def transfer(self, dest:Any, src:Any, sz:int):
+  def transfer(self, dest:Any, src:Any, sz:int, src_dev: MetalDevice, **kwargs):
+    src_dev.synchronize()
     command_buffer = self.device.mtl_queue.commandBuffer()
     encoder = command_buffer.blitCommandEncoder()
     encoder.copyFromBuffer_sourceOffset_toBuffer_destinationOffset_size_(src, 0, dest, 0, sz)
@@ -64,7 +81,7 @@ class MetalAllocator(LRUAllocator):
     ret = self.device.device.newBufferWithBytesNoCopy_length_options_deallocator_(src, len(src), Metal.MTLResourceStorageModeShared, None)
     if ret: self.device.mv_in_metal.append(src)
     return ret
-  def _free(self, opaque:Any): opaque.release()
+  def _free(self, opaque:Any, options): opaque.release()
   def as_buffer(self, src:Any) -> memoryview:
     self.device.synchronize()
     return src.contents().as_buffer(src.length())
@@ -72,17 +89,17 @@ class MetalAllocator(LRUAllocator):
   def copyout(self, dest:memoryview, src:Any): dest[:] = self.as_buffer(src)
 
 class MetalDevice(Compiled):
-  compiler_device = None
   def __init__(self, device:str):
     self.device = Metal.MTLCreateSystemDefaultDevice()
-    if MetalDevice.compiler_device is None: MetalDevice.compiler_device = self.device
     self.mtl_queue = self.device.newCommandQueueWithMaxCommandBufferCount_(1024)
     self.mtl_buffers_in_flight: List[Any] = []
     self.mv_in_metal: List[memoryview] = []
+    self.track_cross_buffer: List[Any] = []
     from tinygrad.runtime.graph.metal import MetalGraph
-    super().__init__(MetalAllocator(self), LinearizerOptions(device="METAL"), MetalRenderer,
-                     compile_metal, functools.partial(MetalProgram, self), functools.partial(MetalGraph, self))
+    super().__init__(device, MetalAllocator(self), MetalRenderer(), MetalCompiler(None if getenv("METAL_XCODE") else self),
+                     functools.partial(MetalProgram, self), MetalGraph)
   def synchronize(self):
-    for cbuf in self.mtl_buffers_in_flight: cbuf.waitUntilCompleted()
+    for cbuf in self.mtl_buffers_in_flight: wait_check(cbuf)
     self.mv_in_metal.clear()
     self.mtl_buffers_in_flight.clear()
+    self.track_cross_buffer.clear()

tinygrad/runtime/ops_npy.py

@@ -0,0 +1,9 @@
+import numpy as np
+from tinygrad.helpers import flat_mv
+from tinygrad.device import Compiled, Allocator
+
+class NpyAllocator(Allocator):  # pylint: disable=abstract-method
+  def copyout(self, dest:memoryview, src:np.ndarray): dest[:] = flat_mv(np.require(src, requirements='C').data)
+
+class NpyDevice(Compiled):
+  def __init__(self, device:str): super().__init__(device, NpyAllocator(), None, None, None)