maybe

.github/workflows/autogen.yml

@@ -130,28 +130,6 @@ jobs:
         mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
         python3 -c "from tinygrad.runtime.autogen import mesa"
         diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
-    - name: Verify libclang autogen
-      run: |
-        mv tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
-        python3 -c "from tinygrad.runtime.autogen import libclang"
-        diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
-  autogen-mac:
-    name: In-tree Autogen (macos)
-    runs-on: macos-14
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        llvm: 'true'
-        pydeps: 'clang>=20'
-    - name: Verify macos autogen
-      run: |
-        mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
-        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
-        diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
   autogen-comgr-3:
     name: In-tree Autogen (comgr 3)
     runs-on: ubuntu-24.04

test/external/external_openpilot_image_warp.py

@@ -1,55 +0,0 @@
-import time
-from tinygrad.tensor import Tensor, Device
-
-MODEL_WIDTH = 512
-MODEL_HEIGHT = 256
-MODEL_FRAME_SIZE = MODEL_WIDTH * MODEL_HEIGHT * 3 // 2
-IMG_INPUT_SHAPE = (1, 12, 128, 256)
-
-def tensor_arange(end): return Tensor([float(i) for i in range(end)])
-def tensor_round(tensor:Tensor): return (tensor + 0.5).floor()
-
-h_src, w_src = 1208, 1928
-h_dst, w_dst = MODEL_HEIGHT, MODEL_WIDTH
-x = tensor_arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst)
-y = tensor_arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst)
-ones = Tensor.ones_like(x)
-dst_coords = x.reshape((1,-1)).cat(y.reshape((1,-1))).cat(ones.reshape((1,-1)))
-
-def warp_perspective_tinygrad(src:Tensor, M_inv:Tensor) -> Tensor:
-  src_coords = M_inv @ dst_coords
-  src_coords = src_coords / src_coords[2:3, :]
-
-  x_src = src_coords[0].reshape(h_dst, w_dst)
-  y_src = src_coords[1].reshape(h_dst, w_dst)
-
-  x_nearest = tensor_round(x_src).clip(0, w_src - 1).cast('int')
-  y_nearest = tensor_round(y_src).clip(0, h_src - 1).cast('int')
-
-  # TODO: make 2d indexing fast
-  idx = y_nearest*src.shape[1] + x_nearest
-  dst = src.flatten()[idx]
-  return dst.reshape(h_dst, w_dst)
-
-if __name__ == "__main__":
-  from tinygrad.engine.jit import TinyJit
-  update_img_jit = TinyJit(warp_perspective_tinygrad, prune=True)
-
-  step_times = []
-  for _ in range(10):
-    # regenerate inputs
-    inputs = [Tensor.randn(1928,1208), Tensor.randn(3,3)]
-    Tensor.realize(*inputs)
-    Device.default.synchronize()
-
-    # do the warp
-    st = time.perf_counter()
-    out = update_img_jit(*inputs)
-    mt = time.perf_counter()
-    val = out.contiguous().realize()
-    Device.default.synchronize()
-    et = time.perf_counter()
-
-    # measure the time
-    step_times.append((et-st)*1e3)
-    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")

test/test_fold.py

@@ -1,31 +0,0 @@
-import numpy as np
-import unittest
-from tinygrad import Tensor, UOp
-from tinygrad.uop.ops import AxisType
-
-class TestFold(unittest.TestCase):
-  def test_reduce_add(self):
-    a = Tensor.randn(10, 10).realize()
-    a_red = a.sum(axis=1)
-    np.testing.assert_allclose(a_red.numpy(), a.numpy().sum(axis=1), atol=1e-6)
-
-  def test_fold_add(self):
-    a = Tensor.randn(10, 10).realize()
-    init = Tensor.zeros(10, 1).contiguous()
-    a_red = (init+a).fold(init).reshape(10)
-    np.testing.assert_allclose(a_red.numpy(), a.numpy().sum(axis=1), atol=1e-6)
-
-  #@unittest.skip("no outer fold yet")
-  def test_fold_matmul(self):
-    vec = Tensor.randn(1, 10).realize()
-    mats = Tensor.randn(3, 10, 10).realize()
-    np_mats = mats.numpy()
-    np_ref = ((vec.numpy() @ np_mats[0]) @ np_mats[1]) @ np_mats[2]
-
-    i = UOp.range(3, -1, AxisType.OUTER)
-    out = (vec @ mats[i]).contiguous().fold(vec, i)
-
-    np.testing.assert_allclose(out.numpy(), np_ref, atol=1e-6)
-
-if __name__ == '__main__':
-  unittest.main()

tinygrad/codegen/late/devectorizer.py

@@ -308,19 +308,9 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
   if len(reduce_range) == 0: return ret
   return acc.after(acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)).index(UOp.const(dtypes.int, 0))
 
-def fold_to_store(x:UOp):
-  _, acc, ranges = x.src[0], x.src[1], x.src[2:]
-  assert acc.op is Ops.INDEX
-  buf = acc.src[0]
-  ret = x.substitute({buf: buf.rtag().after(*ranges)}).substitute({buf.rtag(): buf})
-  base, acc, ranges = ret.src[0], ret.src[1], ret.src[2:]
-  return buf.after(acc.store(base).end(*ranges)).index(acc.src[1])
-
 pm_reduce = PatternMatcher([
-  # REDUCE -> DEFINE_ACC+STORE
+  # REDUCE -> DEFINE_ACC+ASSIGN
   (UPat(Ops.REDUCE, name="red"), reduce_to_acc),
-  # FOLD -> STORE
-  (UPat(Ops.FOLD, name="x"), fold_to_store),
   # tensor core built in accumulate
   (UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
     lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)),

tinygrad/codegen/opt/postrange.py

@@ -2,8 +2,7 @@ from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp
-from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, axis_letters, axis_colors
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
@@ -13,6 +12,10 @@ from tinygrad.renderer import Renderer
 
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
+# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
+axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
+               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
+
 class Scheduler:
   def __init__(self, ast:UOp, ren:Renderer):
     self.ast, self.ren = ast, ren

tinygrad/runtime/autogen/__init__.py

@@ -4,7 +4,6 @@ from tinygrad.helpers import fetch, flatten, system
 root = (here:=pathlib.Path(__file__).parent).parents[2]
 nv_src = {"nv_570": "https://github.com/NVIDIA/open-gpu-kernel-modules/archive/81fe4fb417c8ac3b9bdcc1d56827d116743892a5.tar.gz",
           "nv_580": "https://github.com/NVIDIA/open-gpu-kernel-modules/archive/2af9f1f0f7de4988432d4ae875b5858ffdb09cc2.tar.gz"}
-macossdk = "/var/db/xcode_select_link/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk"
 
 def load(name, dll, files, **kwargs):
   if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists():
@@ -121,12 +120,4 @@ python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
 python3 src/compiler/builtin_types_h.py gen/builtin_types.h""", cwd=path, shell=True, check=True),
   tarball="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.4/mesa-25.2.4.tar.gz",
   prolog=["import gzip, base64", "from tinygrad.helpers import OSX"], epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
-    case "libclang":
-      return load("libclang", ["os.getenv('LIBCLANG_PATH', find_library('clang-20'))"],
-                  lambda: [system("llvm-config-20 --includedir")+"/clang-c/Index.h"], args=lambda: system("llvm-config-20 --cflags").split(),
-                  types={"CXString":"ci._CXString","CXType":"ci.Type","CXCursor":"ci.Cursor"}, prolog=["import clang.cindex as ci"])
-    case "metal":
-      return load("metal", ["find_library('Metal')"],[f"{macossdk}/System/Library/Frameworks/Metal.framework/Headers/MTL{s}.h" for s in
-                  ["ComputeCommandEncoder", "ComputePipeline", "CommandQueue", "Device", "IndirectCommandBuffer", "Resource", "CommandEncoder"]],
-                  args=["-xobjective-c","-isysroot",macossdk], types={"dispatch_data_t":"objc.id_"})
     case _: raise AttributeError(f"no such autogen: {nm}")

tinygrad/runtime/graph/metal.py

@@ -5,9 +5,15 @@ from tinygrad.helpers import dedup, getenv, merge_dicts, PROFILE
 from tinygrad.device import Buffer, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.engine.realize import ExecItem, CompiledRunner
 from tinygrad.engine.jit import GraphRunner, GraphException
-from tinygrad.runtime.ops_metal import wait_check, to_ns_str
-from tinygrad.runtime.autogen import metal
-from tinygrad.runtime.support import objc
+from tinygrad.runtime.ops_metal import wait_check, msg, libobjc, to_struct, objc_instance,\
+  MTLResourceOptions, cmdbuf_st_time, cmdbuf_en_time, objc_id, to_ns_str
+
+class MTLIndirectCommandType:
+  MTLIndirectCommandTypeConcurrentDispatch = (1 << 5)
+
+class MTLResourceUsage:
+  MTLResourceUsageRead = 0b01
+  MTLResourceUsageWrite = 0b10
 
 class MetalGraph(GraphRunner):
   def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[str, int]):
@@ -15,17 +21,16 @@ class MetalGraph(GraphRunner):
     if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException
 
     # create metal batch exec
-    icb_descriptor = metal.MTLIndirectCommandBufferDescriptor.new()
-    icb_descriptor.setCommandTypes(metal.MTLIndirectCommandTypeConcurrentDispatch)
-    icb_descriptor.setInheritBuffers(False)
-    icb_descriptor.setInheritPipelineState(False)
-    icb_descriptor.setMaxKernelBufferBindCount(31)
+    icb_descriptor = msg("new", objc_instance)(libobjc.objc_getClass(b"MTLIndirectCommandBufferDescriptor"))
+    msg("setCommandTypes:")(icb_descriptor, MTLIndirectCommandType.MTLIndirectCommandTypeConcurrentDispatch)
+    msg("setInheritBuffers:")(icb_descriptor, False)
+    msg("setInheritPipelineState:")(icb_descriptor, False)
+    msg("setMaxKernelBufferBindCount:")(icb_descriptor, 31)
 
-    self.icb = self.dev.sysdevice.newIndirectCommandBufferWithDescriptor_maxCommandCount_options(icb_descriptor, len(jit_cache),
-                                                                                                 metal.MTLResourceCPUCacheModeDefaultCache)
+    self.icb = msg("newIndirectCommandBufferWithDescriptor:maxCommandCount:options:", objc_instance)(self.dev.sysdevice,
+      icb_descriptor, len(jit_cache), MTLResourceOptions.MTLResourceCPUCacheModeDefaultCache)
     if self.icb.value is None: raise GraphException("create indirect command buffer failed, does your system support this?")
-    # TODO: needs categories
-    icb_label = bytes(objc.msg("UTF8String", ctypes.c_char_p)(objc.msg("description")(self.icb).retained())).decode()
+    icb_label = bytes(msg("UTF8String", ctypes.c_char_p)(msg("description", objc_instance)(self.icb))).decode()
     self.needs_icb_fix = int((m := re.search(r'AGXG(\d+)XFamily', icb_label)) is None or int(m.group(1)) < 15) # not required on M3+
 
     self.fixedvars = merge_dicts([ji.fixedvars for ji in jit_cache])
@@ -35,25 +40,26 @@ class MetalGraph(GraphRunner):
     all_pipelines, all_resources = [], [self.int_buf.buf] if len(self.varlist) else []
     for j,ji in enumerate(jit_cache):
       prg: CompiledRunner = cast(CompiledRunner, ji.prg)
-      icb_command = self.icb.indirectComputeCommandAtIndex(j).retained()
+      icb_command = msg("indirectComputeCommandAtIndex:", objc_instance)(self.icb, j)
       all_pipelines.append(prg._prg.pipeline_state)
-      icb_command.setComputePipelineState(prg._prg.pipeline_state)
+      msg("setComputePipelineState:")(icb_command, prg._prg.pipeline_state)
       for i,b in enumerate(ji.bufs):
         if b is not None and b not in input_rawbuffers:
-          icb_command.setKernelBuffer_offset_atIndex(b._buf.buf, b._buf.offset, i)
+          msg("setKernelBuffer:offset:atIndex:")(icb_command, b._buf.buf, b._buf.offset, i)
           all_resources.append(b._buf.buf)
-      for i,v in enumerate(prg.p.vars): icb_command.setKernelBuffer_offset_atIndex(self.int_buf.buf, self.varlist.index(v.expr)*4, len(ji.bufs)+i)
+      for i,v in enumerate(prg.p.vars):
+        msg("setKernelBuffer:offset:atIndex:")(icb_command, self.int_buf.buf, self.varlist.index(v.expr)*4, len(ji.bufs)+i)
 
       global_size, local_size = prg.p.launch_dims(var_vals)
-      icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size))
-      icb_command.setBarrier()
+      msg("concurrentDispatchThreadgroups:threadsPerThreadgroup:")(icb_command, to_struct(*global_size), to_struct(*local_size))
+      msg("setBarrier")(icb_command)
 
     self.all_resources = dedup(all_resources)
     self.all_pipelines = dedup(all_pipelines)
     self.command_buffer: Any = None
     if len(self.varlist): self.int_buf_view = self.dev.allocator._as_buffer(self.int_buf).cast('i')
     for var in self.fixedvars: self.int_buf_view[self.varlist.index(var)] = self.fixedvars[var]
-    self.range = metal.NSRange(0, len(jit_cache))
+    self.range = to_struct(0, len(jit_cache))
 
   def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[str, int], wait=False) -> float|None:
     if self.command_buffer is not None and self.command_buffer in self.dev.mtl_buffers_in_flight: wait_check(self.command_buffer)
@@ -62,18 +68,18 @@ class MetalGraph(GraphRunner):
 
     all_resources = dedup(self.all_resources + [input_rawbuffers[input_idx]._buf.buf for input_idx in self.input_replace.values()])
     for (j,i),input_idx in self.input_replace.items():
-      computeCommand = self.icb.indirectComputeCommandAtIndex(j)
-      computeCommand.setKernelBuffer_offset_atIndex(input_rawbuffers[input_idx]._buf.buf, input_rawbuffers[input_idx]._buf.offset, i)
+      computeCommand = msg("indirectComputeCommandAtIndex:", objc_id)(self.icb, j)
+      msg("setKernelBuffer:offset:atIndex:")(computeCommand, input_rawbuffers[input_idx]._buf.buf, input_rawbuffers[input_idx]._buf.offset, i)
 
     for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
-      computeCommand = self.icb.indirectComputeCommandAtIndex(j)
-      computeCommand.concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_dims), metal.MTLSize(*local_dims))
+      computeCommand = msg("indirectComputeCommandAtIndex:", objc_id)(self.icb, j)
+      msg("concurrentDispatchThreadgroups:threadsPerThreadgroup:")(computeCommand, to_struct(*global_dims), to_struct(*local_dims))
     for var in self.vars: self.int_buf_view[self.varlist.index(var)] = var_vals[var]
 
-    command_buffer = self.dev.mtl_queue.commandBuffer().retained()
-    encoder = command_buffer.computeCommandEncoder().retained()
-    encoder.useResources_count_usage(ctypes.cast((metal.MTLBuffer * len(all_resources))(*all_resources), ctypes.POINTER(metal.MTLResource)),
-                                     len(all_resources), metal.MTLResourceUsageRead | metal.MTLResourceUsageWrite)
+    command_buffer = msg("commandBuffer", objc_instance)(self.dev.mtl_queue)
+    encoder = msg("computeCommandEncoder", objc_instance)(command_buffer)
+    msg("useResources:count:usage:")(encoder, (objc_id * len(all_resources))(*all_resources), len(all_resources),
+        MTLResourceUsage.MTLResourceUsageRead | MTLResourceUsage.MTLResourceUsageWrite)
 
     # NOTE: the pipelines likely need to be added to the used resources to fix the crash on M1/M2, but I haven't figured out how
     # this is a O(n) hack to get them used. what should work is:
@@ -82,24 +88,24 @@ class MetalGraph(GraphRunner):
     # to repro the crash (which can also crash other running GPU apps), run with FIX_METAL_ICB=0
     if getenv("FIX_METAL_ICB", self.needs_icb_fix):
       for ps in self.all_pipelines:
-        encoder.setComputePipelineState(ps)
-        encoder.dispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(0,0,0), metal.MTLSize(0,0,0))
+        msg("setComputePipelineState:")(encoder, ps)
+        msg("dispatchThreadgroups:threadsPerThreadgroup:")(encoder, to_struct(0,0,0), to_struct(0,0,0))
 
-    encoder.executeCommandsInBuffer_withRange(self.icb, self.range)
-    encoder.endEncoding()
-    command_buffer.setLabel(to_ns_str(f"batched {len(self.jit_cache)}"))
-    command_buffer.commit()
+    msg("executeCommandsInBuffer:withRange:")(encoder, self.icb, self.range)
+    msg("endEncoding")(encoder)
+    msg("setLabel:")(command_buffer, to_ns_str(f"batched {len(self.jit_cache)}"))
+    msg("commit")(command_buffer)
     self.command_buffer = command_buffer
 
     self.dev.mtl_buffers_in_flight.append(command_buffer)
     if wait:
       wait_check(command_buffer)
-      return command_buffer.GPUEndTime() - command_buffer.GPUStartTime()
+      return cmdbuf_en_time(command_buffer) - cmdbuf_st_time(command_buffer)
     return None
 
   def collect_timestamps(self):
     # create a graph event and evenly space each program
-    st, en = decimal.Decimal(self.command_buffer.GPUStartTime()) * 1000000, decimal.Decimal(self.command_buffer.GPUEndTime()) * 1000000
+    st, en = decimal.Decimal(cmdbuf_st_time(self.command_buffer)) * 1000000, decimal.Decimal(cmdbuf_en_time(self.command_buffer)) * 1000000
     ents = [ProfileGraphEntry(self.device, cast(CompiledRunner, ji.prg)._prg.name, i, i+1, is_copy=False) for i,ji in enumerate(self.jit_cache)]
     step = (en-st)/len(ents)
     self.dev.profile_events += [ProfileGraphEvent(ents, [], [st+step*i for i in range(len(ents)+1)])]

tinygrad/runtime/ops_metal.py

@@ -1,42 +1,74 @@
-import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform, sys
-from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, init_c_struct_t, PROFILE, ProfileRangeEvent, cpu_profile, unwrap
-import tinygrad.runtime.support.objc as objc
+import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
+from typing import Any, cast
+from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, T, init_c_struct_t, PROFILE, ProfileRangeEvent, cpu_profile, unwrap
 from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent
 from tinygrad.renderer.cstyle import MetalRenderer
-from tinygrad.runtime.autogen import metal
+
+class objc_id(ctypes.c_void_p): # This prevents ctypes from converting response to plain int, and dict.fromkeys() can use it to dedup
+  def __hash__(self): return hash(self.value)
+  def __eq__(self, other): return self.value == other.value
+
+class objc_instance(objc_id): # method with name "new", "alloc" should be freed after use
+  def __del__(self):
+    # CPython doesn't make any guarantees about order in which globals (like `msg` or `libobjc`) are destroyed when the interpreter shuts down
+    # https://github.com/tinygrad/tinygrad/pull/8949 triggered the unlucky ordering which lead to a bunch of errors at exit
+    # TODO: Why isn't `sys.is_finalizing` working?
+    if msg is not None and libobjc is not None: msg("release")(self)
+
+class MTLResourceOptions:
+  MTLResourceCPUCacheModeDefaultCache = 0
+  MTLResourceStorageModeShared = 0 << 4
+
+class MTLPipelineOption:
+  MTLPipelineOptionNone = 0
 
 # 13 is requestType that metal uses to compile source code into MTLB, there aren't any docs or symbols.
 REQUEST_TYPE_COMPILE = 13
 
+libobjc = ctypes.CDLL("/usr/lib/libobjc.dylib")
+libmetal = ctypes.CDLL("/System/Library/Frameworks/Metal.framework/Metal")
 # Must be loaded for default Metal Device: https://developer.apple.com/documentation/metal/1433401-mtlcreatesystemdefaultdevice?language=objc
 ctypes.CDLL("/System/Library/Frameworks/CoreGraphics.framework/CoreGraphics")
+libdispatch = ctypes.CDLL("/usr/lib/libSystem.dylib") # libdispatch is part of libSystem on mac
+libobjc.objc_getClass.restype = objc_id
+libobjc.sel_registerName.restype = objc_id
+libmetal.MTLCreateSystemDefaultDevice.restype = objc_instance
+libdispatch.dispatch_data_create.restype = objc_instance
 
-# FIXME: these need autogen to support objc categories
-# https://developer.apple.com/library/archive/documentation/Cocoa/Conceptual/ObjectiveC/Chapters/ocCategories.html
 @functools.cache
-def to_ns_str(s: str): return ctypes.cast(objc.msg("stringWithUTF8String:")(metal.NSString._objc_class_, s.encode()), metal.NSString)
-def from_ns_str(s): return bytes(objc.msg("UTF8String", ctypes.c_char_p)(s)).decode()
+def msg(selector: str, restype: type[T] = objc_id):  # type: ignore [assignment]
+  resname = libobjc.sel_registerName(selector.encode())
+  sender = libobjc["objc_msgSend"] # Using attribute access returns a new reference so setting restype is safe
+  sender.restype = restype
+  def _msg(ptr: objc_id, *args: Any) -> T: return sender(ptr, resname, *args)
+  return _msg
+
+@functools.cache
+def to_ns_str(s: str): return msg("stringWithUTF8String:", objc_instance)(libobjc.objc_getClass(b"NSString"), s.encode())
+def from_ns_str(s): return bytes(msg("UTF8String", ctypes.c_char_p)(s)).decode()
 
 def to_struct(*t: int, _type: type[ctypes._SimpleCData] = ctypes.c_ulong):
   return init_c_struct_t(tuple([(f"field{i}", _type) for i in range(len(t))]))(*t)
 
-def wait_check(cbuf:metal.MTLCommandBuffer):
-  cbuf.waitUntilCompleted()
-  error_check(cbuf.error().retained())
+def wait_check(cbuf: Any):
+  msg("waitUntilCompleted")(cbuf)
+  error_check(msg("error", objc_instance)(cbuf))
 
-def cmdbuf_label(cbuf:metal.MTLCommandBuffer) -> str|None: return from_ns_str(label) if (label:=cbuf.label()).value is not None else None
+def cmdbuf_label(cbuf: objc_id) -> str|None: return from_ns_str(label) if (label:=msg("label", objc_id)(cbuf)).value is not None else None
+def cmdbuf_st_time(cbuf: objc_id) -> float: return cast(float, msg("GPUStartTime", ctypes.c_double)(cbuf))
+def cmdbuf_en_time(cbuf: objc_id) -> float: return cast(float, msg("GPUEndTime", ctypes.c_double)(cbuf))
 
-def error_check(error: metal.NSError, error_constructor: type[Exception] = RuntimeError):
+def error_check(error: objc_instance, error_constructor: type[Exception] = RuntimeError):
   if error.value is None: return None
-  raise error_constructor(from_ns_str(error.localizedDescription().retained()))
+  raise error_constructor(from_ns_str(msg("localizedDescription", objc_instance)(error)))
 
 class MetalDevice(Compiled):
   def __init__(self, device:str):
-    self.sysdevice = metal.MTLCreateSystemDefaultDevice()
-    self.mtl_queue = self.sysdevice.newCommandQueueWithMaxCommandBufferCount(1024)
+    self.sysdevice = libmetal.MTLCreateSystemDefaultDevice()
+    self.mtl_queue = msg("newCommandQueueWithMaxCommandBufferCount:", objc_instance)(self.sysdevice, 1024)
     if self.mtl_queue is None: raise RuntimeError("Cannot allocate a new command queue")
-    self.mtl_buffers_in_flight: list[metal.MTLCommandBuffer] = []
-    self.timeline_signal = self.sysdevice.newSharedEvent()
+    self.mtl_buffers_in_flight: list[Any] = []
+    self.timeline_signal = msg("newSharedEvent", objc_instance)(self.sysdevice)
     self.timeline_value = 0
 
     Compiled.profile_events += [ProfileDeviceEvent(device)]
@@ -44,22 +76,23 @@ class MetalDevice(Compiled):
     from tinygrad.runtime.graph.metal import MetalGraph
     # NOTE: GitHub CI macOS runners use paravirtualized metal which is broken with graph.
     # This can be reproduced locally with any virtualization software (like utm) that can create macOS VMs with apple's own virtualization framework.
-    super().__init__(device, MetalAllocator(self), [(MetalRenderer, MetalCompiler), (MetalRenderer, Compiler)], functools.partial(MetalProgram, self),
-                     MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None)
+    super().__init__(device, MetalAllocator(self), [(MetalRenderer, MetalCompiler), (MetalRenderer, Compiler)],
+                     functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(msg('name')(self.sysdevice)).lower() else None)
 
   def synchronize(self):
     for cbuf in self.mtl_buffers_in_flight:
       wait_check(cbuf)
-      st, en = decimal.Decimal(cbuf.GPUStartTime()) * 1000000, decimal.Decimal(cbuf.GPUEndTime()) * 1000000
+      st, en = decimal.Decimal(cmdbuf_st_time(cbuf)) * 1000000, decimal.Decimal(cmdbuf_en_time(cbuf)) * 1000000
       # NOTE: command buffers from MetalGraph are not profiled here
       if PROFILE and (lb:=cmdbuf_label(cbuf)) is not None and not lb.startswith("batched"):
         Compiled.profile_events += [ProfileRangeEvent(self.device, lb, st, en, is_copy=lb.startswith("COPY"))]
     self.mtl_buffers_in_flight.clear()
 
-def metal_src_to_library(device:MetalDevice, src:str) -> metal.MTLLibrary:
-  options = metal.MTLCompileOptions.new()
-  options.setFastMathEnabled(getenv("METAL_FAST_MATH"))
-  library = device.sysdevice.newLibraryWithSource_options_error(to_ns_str(src), options, ctypes.byref(compileError:=metal.NSError().retained()))
+def metal_src_to_library(device:MetalDevice, src:str) -> objc_instance:
+  options = msg("new", objc_instance)(libobjc.objc_getClass(b"MTLCompileOptions"))
+  msg("setFastMathEnabled:")(options, getenv("METAL_FAST_MATH"))
+  library = msg("newLibraryWithSource:options:error:", objc_instance)(device.sysdevice, to_ns_str(src),
+                                                                      options, ctypes.byref(compileError:=objc_instance()))
   error_check(compileError, CompileError)
   return library
 
@@ -122,71 +155,71 @@ class MetalProgram:
     self.dev, self.name, self.lib = dev, name, lib
     if lib[:4] == b"MTLB":
       # binary metal library
-      data = objc.dispatch_data_create(lib, len(lib), None, None)
-      self.library = self.dev.sysdevice.newLibraryWithData_error(data, ctypes.byref(error_lib:=metal.NSError().retained())).retained()
+      data = libdispatch.dispatch_data_create(lib, len(lib), None, None)
+      self.library = msg("newLibraryWithData:error:", objc_instance)(self.dev.sysdevice, data, ctypes.byref(error_lib:=objc_instance()))
       error_check(error_lib)
     else:
       # metal source. rely on OS caching
       try: self.library = metal_src_to_library(self.dev, lib.decode())
       except CompileError as e: raise RuntimeError from e
-    self.fxn = self.library.newFunctionWithName(to_ns_str(name)).retained()
-    descriptor = metal.MTLComputePipelineDescriptor.new()
-    descriptor.setComputeFunction(self.fxn)
-    descriptor.setSupportIndirectCommandBuffers(True)
-    self.pipeline_state = self.dev.sysdevice.newComputePipelineStateWithDescriptor_options_reflection_error(descriptor, metal.MTLPipelineOptionNone,
-      None, ctypes.byref(error_pipeline_creation:=metal.NSError().retained()))
+    self.fxn = msg("newFunctionWithName:", objc_instance)(self.library, to_ns_str(name))
+    descriptor = msg("new", objc_instance)(libobjc.objc_getClass(b"MTLComputePipelineDescriptor"))
+    msg("setComputeFunction:")(descriptor, self.fxn)
+    msg("setSupportIndirectCommandBuffers:")(descriptor, True)
+    self.pipeline_state = msg("newComputePipelineStateWithDescriptor:options:reflection:error:", objc_instance)(self.dev.sysdevice,
+      descriptor, MTLPipelineOption.MTLPipelineOptionNone, None, ctypes.byref(error_pipeline_creation:=objc_instance()))
     error_check(error_pipeline_creation)
     # cache these msg calls
-    self.max_total_threads: int = self.pipeline_state.maxTotalThreadsPerThreadgroup()
+    self.max_total_threads: int = cast(int, msg("maxTotalThreadsPerThreadgroup", ctypes.c_ulong)(self.pipeline_state))
 
   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.max_total_threads:
-      exec_width = self.pipeline_state.threadExecutionWidth()
-      memory_length = self.pipeline_state.staticThreadgroupMemoryLength()
+      exec_width = msg("threadExecutionWidth", ctypes.c_ulong)(self.pipeline_state)
+      memory_length = msg("staticThreadgroupMemoryLength", ctypes.c_ulong)(self.pipeline_state)
       raise RuntimeError(f"local size {local_size} bigger than {self.max_total_threads} with exec width {exec_width} memory length {memory_length}")
-    command_buffer = self.dev.mtl_queue.commandBuffer().retained() # FIXME: is this really ARC?
-    encoder = command_buffer.computeCommandEncoder().retained() # FIXME: is this really ARC?
-    encoder.setComputePipelineState(self.pipeline_state)
-    for i,a in enumerate(bufs): encoder.setBuffer_offset_atIndex(a.buf, a.offset, i)
-    for i,a in enumerate(vals, start=len(bufs)): encoder.setBytes_length_atIndex(bytes(ctypes.c_int(a)), 4, i)
-    encoder.dispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size))
-    encoder.endEncoding()
-    command_buffer.setLabel(to_ns_str(self.name)) # TODO: is this always needed?
-    command_buffer.commit()
+    command_buffer = msg("commandBuffer", objc_instance)(self.dev.mtl_queue)
+    encoder = msg("computeCommandEncoder", objc_instance)(command_buffer)
+    msg("setComputePipelineState:")(encoder, self.pipeline_state)
+    for i,a in enumerate(bufs): msg("setBuffer:offset:atIndex:")(encoder, a.buf, a.offset, i)
+    for i,a in enumerate(vals, start=len(bufs)): msg("setBytes:length:atIndex:")(encoder, bytes(ctypes.c_int(a)), 4, i)
+    msg("dispatchThreadgroups:threadsPerThreadgroup:")(encoder, to_struct(*global_size), to_struct(*local_size))
+    msg("endEncoding")(encoder)
+    msg("setLabel:")(command_buffer, to_ns_str(self.name)) # TODO: is this always needed?
+    msg("commit")(command_buffer)
     self.dev.mtl_buffers_in_flight.append(command_buffer)
     if wait:
       wait_check(command_buffer)
-      return command_buffer.GPUEndTime() - command_buffer.GPUStartTime()
+      return cmdbuf_en_time(command_buffer) - cmdbuf_st_time(command_buffer)
 
 class MetalBuffer:
-  def __init__(self, buf:metal.MTLBuffer, size:int, offset=0): self.buf, self.size, self.offset = buf, size, offset
+  def __init__(self, buf:Any, size:int, offset=0): self.buf, self.size, self.offset = buf, size, offset
 
 class MetalAllocator(LRUAllocator[MetalDevice]):
   def _alloc(self, size:int, options) -> MetalBuffer:
-    if options.external_ptr: return MetalBuffer(metal.MTLBuffer(options.external_ptr), size)
+    if options.external_ptr: return MetalBuffer(objc_id(options.external_ptr), size)
 
     # Buffer is explicitly released in _free() rather than garbage collected via reference count
-    ret = self.dev.sysdevice.newBufferWithLength_options(size, metal.MTLResourceStorageModeShared)
-    ret.retain = False
+    ret = msg("newBufferWithLength:options:", objc_id)(self.dev.sysdevice, ctypes.c_ulong(size), MTLResourceOptions.MTLResourceStorageModeShared)
     if ret.value is None: raise MemoryError(f"Metal OOM while allocating {size=}")
     return MetalBuffer(ret, size)
   def _free(self, opaque:MetalBuffer, options):
-    if not sys.is_finalizing(): opaque.buf.release
+    if msg is not None and libobjc is not None: msg("release")(opaque.buf)
   def _transfer(self, dest:MetalBuffer, src:MetalBuffer, sz:int, src_dev:MetalDevice, dest_dev:MetalDevice):
     dest_dev.synchronize()
-    src_command_buffer = src_dev.mtl_queue.commandBuffer().retained()
-    encoder = src_command_buffer.blitCommandEncoder().retained()
-    encoder.copyFromBuffer_sourceOffset_toBuffer_destinationOffset_size(src.buf, src.offset, dest.buf, dest.offset, sz)
-    encoder.endEncoding()
+    src_command_buffer = msg("commandBuffer", objc_instance)(src_dev.mtl_queue)
+    encoder = msg("blitCommandEncoder", objc_instance)(src_command_buffer)
+    msg("copyFromBuffer:sourceOffset:toBuffer:destinationOffset:size:")(encoder, src.buf, ctypes.c_ulong(src.offset),
+        dest.buf, ctypes.c_ulong(dest.offset), ctypes.c_ulong(sz))
+    msg("endEncoding")(encoder)
     if src_dev != dest_dev:
-      src_command_buffer.encodeSignalEvent_value(ctypes.cast(src_dev.timeline_signal, metal.MTLEvent), src_dev.timeline_value)
-      dest_command_buffer = dest_dev.mtl_queue.commandBuffer().retained()
-      dest_command_buffer.encodeWaitForEvent_value(ctypes.cast(src_dev.timeline_signal, metal.MTLEvent), src_dev.timeline_value)
-      dest_command_buffer.commit()
+      msg("encodeSignalEvent:value:")(src_command_buffer, src_dev.timeline_signal, src_dev.timeline_value)
+      dest_command_buffer = msg("commandBuffer", objc_instance)(dest_dev.mtl_queue)
+      msg("encodeWaitForEvent:value:")(dest_command_buffer, src_dev.timeline_signal, src_dev.timeline_value)
+      msg("commit")(dest_command_buffer)
       dest_dev.mtl_buffers_in_flight.append(dest_command_buffer)
       src_dev.timeline_value += 1
-    src_command_buffer.setLabel(to_ns_str(f"COPY {src_dev.device} -> {dest_dev.device}"))
-    src_command_buffer.commit()
+    msg("setLabel:")(src_command_buffer, to_ns_str(f"COPY {src_dev.device} -> {dest_dev.device}"))
+    msg("commit")(src_command_buffer)
     src_dev.mtl_buffers_in_flight.append(src_command_buffer)
     # Transfers currently synchronize the completion. Otherwise, copies can sometimes lead to incorrect values.
     # There is no real metal multidevice support for now, so transfer is used only for tests.
@@ -195,7 +228,7 @@ class MetalAllocator(LRUAllocator[MetalDevice]):
     with cpu_profile(prof_desc, self.dev.device, is_copy=True): dst[:] = src
   def _as_buffer(self, src:MetalBuffer) -> memoryview:
     self.dev.synchronize()
-    return to_mv(src.buf.contents(), src.size + src.offset)[src.offset:]
+    return to_mv(cast(int, msg("contents", objc_id)(src.buf).value), src.size + src.offset)[src.offset:]
   def _copyin(self, dest:MetalBuffer, src:memoryview): self._cp_mv(self._as_buffer(dest), src, "TINY -> METAL")
   def _copyout(self, dest:memoryview, src:MetalBuffer): self._cp_mv(dest, self._as_buffer(src), "METAL -> TINY")
   def _offset(self, buf:MetalBuffer, size:int, offset:int): return MetalBuffer(buf.buf, size, offset)

tinygrad/runtime/support/autogen.py

@@ -1,12 +1,9 @@
 import ctypes.util, importlib.metadata, itertools, re, functools, os
-from tinygrad.helpers import flatten, unwrap, fromimport
-
+from tinygrad.helpers import flatten, unwrap
 assert importlib.metadata.version('clang')[:2] == "20", 'clang version 20 required, pip install "clang==20.1.0"'
-from clang.cindex import Config, Index, Cursor, Type, CursorKind as CK, TranslationUnit as TU, LinkageKind as LK, TokenKind as ToK, TypeKind as TK
+from clang.cindex import Config, Index, CursorKind as CK, TranslationUnit as TU, LinkageKind as LK, TokenKind as ToK, TypeKind as TK
 from clang.cindex import PrintingPolicy as PP, PrintingPolicyProperty as PPP, SourceRange
 
-libclang = functools.partial(fromimport, "tinygrad.runtime.autogen.libclang") # we can't actually import this, because then we can't generate it
-
 if not Config.loaded: Config.set_library_file(os.getenv("LIBCLANG_PATH", ctypes.util.find_library("clang-20")))
 
 def fst(c): return next(c.get_children())
@@ -17,25 +14,18 @@ def readext(f, fst, snd=None):
     return f.read((fst.end.offset if isinstance(fst, SourceRange) else snd)-start)
 def attrs(c): return list(filter(lambda k: (v:=k.value) >= 400 and v < 500, map(lambda c: c.kind, c.get_children())))
 
-def protocols(t): yield from (Cursor.from_result(libclang("clang_Type_getObjCProtocolDecl")(t, i), t)
-                              for i in range(libclang("clang_Type_getNumObjCProtocolRefs")(t)))
-def basetype(t): return Type.from_result(libclang("clang_Type_getObjCObjectBaseType")(t), (t,))
-
 base_rules = [(r'\s*\\\n\s*', ' '), (r'\s*\n\s*', ' '), (r'//.*', ''), (r'/\*.*?\*/', ''), (r'\b(0[xX][0-9a-fA-F]+|\d+)[uUlL]+\b', r'\1'),
               (r'\b0+(?=\d)', ''), (r'\s*&&\s*', r' and '), (r'\s*\|\|\s*', r' or '), (r'\s*!\s*', ' not '),
               (r'(struct|union|enum)\s*([a-zA-Z_][a-zA-Z0-9_]*\b)', r'\1_\2'),
               (r'\((unsigned )?(char|uint64_t)\)', ''), (r'^.*\d+:\d+.*$', ''), (r'^.*\w##\w.*$', '')]
 
 ints = (TK.INT, TK.UINT, TK.LONG, TK.ULONG, TK.LONGLONG, TK.ULONGLONG)
-specs = (CK.OBJC_SUPER_CLASS_REF,)
-# https://clang.llvm.org/docs/AutomaticReferenceCounting.html#arc-method-families
-arc_families = ['alloc', 'copy', 'mutableCopy', 'new']
 
 def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_errno=False, anon_names={}, types={}, parse_macros=True):
-  macros, lines, anoncnt, types, objc = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}, False
+  macros, lines, anoncnt, types = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}
   def tname(t, suggested_name=None, typedef=None) -> str:
     suggested_name = anon_names.get(f"{(decl:=t.get_declaration()).location.file}:{decl.location.line}", suggested_name)
-    nonlocal lines, types, anoncnt, objc
+    nonlocal lines, types, anoncnt
     tmap = {TK.VOID:"None", TK.CHAR_U:"ctypes.c_ubyte", TK.UCHAR:"ctypes.c_ubyte", TK.CHAR_S:"ctypes.c_char", TK.SCHAR:"ctypes.c_char",
             **{getattr(TK, k):f"ctypes.c_{k.lower()}" for k in ["BOOL", "WCHAR", "FLOAT", "DOUBLE", "LONGDOUBLE"]},
             **{getattr(TK, k):f"ctypes.c_{'u' if 'U' in k else ''}int{sz}" for sz,k in
@@ -47,7 +37,6 @@ def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_e
       return f"ctypes.CFUNCTYPE({tname(f.get_result())}{(', '+', '.join(map(tname, f.argument_types()))) if f.kind==TK.FUNCTIONPROTO else ''})"
     match t.kind:
       case TK.POINTER: return "ctypes.c_void_p" if (ptr:=t.get_pointee()).kind == TK.VOID else f"ctypes.POINTER({tname(ptr)})"
-      case TK.OBJCOBJECTPOINTER: return tname(t.get_pointee()) # TODO: this seems wrong
       case TK.ELABORATED: return tname(t.get_named_type(), suggested_name)
       case TK.TYPEDEF if t.spelling == t.get_canonical().spelling: return tname(t.get_canonical())
       case TK.TYPEDEF:
@@ -87,62 +76,8 @@ def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_e
       case TK.CONSTANTARRAY:
         return f"({tname(t.get_array_element_type(), suggested_name.rstrip('s') if suggested_name else None)} * {t.get_array_size()})"
       case TK.INCOMPLETEARRAY: return f"({tname(t.get_array_element_type(), suggested_name.rstrip('s') if suggested_name else None)} * 0)"
-      case TK.OBJCINTERFACE:
-        is_defn = bool([f.kind for f in decl.get_children() if f.kind in (CK.OBJC_INSTANCE_METHOD_DECL, CK.OBJC_CLASS_METHOD_DECL)])
-        if (nm:=t.spelling) not in types: lines.append(f"class {nm}(objc.Spec): pass")
-        types[nm] = nm, is_defn
-        if is_defn:
-          ims, cms = parse_objc_spec(decl, t.spelling, CK.OBJC_INSTANCE_METHOD_DECL), parse_objc_spec(decl, t.spelling, CK.OBJC_CLASS_METHOD_DECL)
-          lines.extend([*([f"{nm}._bases_ = [{', '.join(bs)}]"] if (bs:=[tname(b.type) for b in decl.get_children() if b.kind in specs]) else []),
-                        *([f"{nm}._methods_ = [", *ims, ']'] if ims else []), *([f"{nm}._classmethods_ = [", *cms, ']'] if cms else [])])
-        return nm
-      case TK.OBJCSEL: return "objc.id_"
-      case TK.OBJCID: return (objc:=True, "objc.id_")[1]
-      case TK.OBJCOBJECT:
-        if basetype(t).kind != TK.OBJCID: raise NotImplementedError(f"generics unsupported: {t.spelling}")
-        ps = [proto(p) for p in protocols(t)]
-        if len(ps) == 0:
-          types[t.spelling] = "objc.id_", True
-          return "objc.id_"
-        if len(ps) == 1:
-          types[t.spelling] = ps[0], True
-          return ps[0]
-        types[t.spelling] = (nm:=f"_anondynamic{anoncnt()}"), True
-        lines.append(f"class {nm}({', '.join(p for p in ps)}): pass # {t.spelling}")
-        return nm
       case _: raise NotImplementedError(f"unsupported type {t.kind}")
 
-  # parses an objc @interface or @protocol, returning a list of declerations that objc.Spec can parse, for the specified kind
-  # NB: ivars are unsupported
-  def parse_objc_spec(decl:Cursor, nm:str, kind:CK) -> list[str]:
-    nonlocal lines, types
-    if decl is None: return []
-    ms = []
-    for d in filter(lambda d: d.kind == kind, decl.get_children()):
-      rollback = lines, types
-      try: ms.append(f"  ('{d.spelling}', {repr('instancetype') if (rt:=d.result_type).spelling=='instancetype' else tname(rt)}, "
-        f"[{', '.join('instancetype' if a.spelling == 'instancetype' else tname(a.type) for a in d.get_arguments())}]" +
-        (", True" if CK.NS_RETURNS_RETAINED in attrs(d) or (any(d.spelling.startswith(s) for s in arc_families) and rt.kind!=TK.VOID) else "") + "),")
-      except NotImplementedError as e:
-        print(f"skipping {nm}.{d.spelling}: {e}")
-        lines, types = rollback
-    return ms
-
-  # libclang doesn't have a "type" for @protocol, so we have to do this here...
-  def proto(decl):
-    nonlocal lines, types
-    if (nm:=decl.spelling) in types and types[nm][1]: return types[nm][0]
-    # check if this is a forward declaration
-    is_defn = bool([f.kind for f in decl.get_children() if f.kind in (CK.OBJC_INSTANCE_METHOD_DECL, CK.OBJC_CLASS_METHOD_DECL)])
-    if nm not in types: lines.append(f"class {nm}(objc.Spec): pass")
-    types[nm] = nm, is_defn
-    if is_defn:
-      bs = [proto(b) for b in decl.get_children() if b.kind==CK.OBJC_PROTOCOL_REF and b.spelling != decl.spelling]
-      ims, cms = parse_objc_spec(decl, nm, CK.OBJC_INSTANCE_METHOD_DECL), parse_objc_spec(decl, nm, CK.OBJC_CLASS_METHOD_DECL)
-      lines.extend([*([f"{nm}._bases_ = [{', '.join(bs)}]"] if bs else []),
-                    *([f"{nm}._methods_ = [", *ims, "]"] if ims else []), *([f"{nm}._classmethods_ = [", *cms, "]"] if cms else [])])
-    return nm
-
   for f in files:
     tu = Index.create().parse(f, args, options=TU.PARSE_DETAILED_PROCESSING_RECORD)
     (pp:=PP.create(tu.cursor)).set_property(PPP.TerseOutput, 1)
@@ -155,8 +90,7 @@ def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_e
             # TODO: we could support name-mangling
             lines.append(f"# {c.pretty_printed(pp)}\ntry: ({c.spelling}:=dll.{c.spelling}).restype, {c.spelling}.argtypes = "
               f"{tname(c.result_type)}, [{', '.join(tname(arg.type) for arg in c.get_arguments())}]\nexcept AttributeError: pass\n")
-            if CK.NS_RETURNS_RETAINED in attrs(c): lines.append(f"{c.spelling} = objc.returns_retained({c.spelling})")
-          case CK.STRUCT_DECL | CK.UNION_DECL | CK.TYPEDEF_DECL | CK.ENUM_DECL | CK.OBJC_INTERFACE_DECL: tname(c.type)
+          case CK.STRUCT_DECL | CK.UNION_DECL | CK.TYPEDEF_DECL | CK.ENUM_DECL: tname(c.type)
           case CK.MACRO_DEFINITION if parse_macros and len(toks:=list(c.get_tokens())) > 1:
             if toks[1].spelling == '(' and toks[0].extent.end.column == toks[1].extent.start.column:
               it = iter(toks[1:])
@@ -173,13 +107,12 @@ def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_e
             else: macros += [f"{c.spelling} = {tname(c.type)}({readext(f, last(c).extent)})"]
           case CK.VAR_DECL if c.linkage == LK.EXTERNAL and dll:
             lines.append(f"try: {c.spelling} = {tname(c.type)}.in_dll(dll, '{c.spelling}')\nexcept (ValueError,AttributeError): pass")
-          case CK.OBJC_PROTOCOL_DECL: proto(c)
       except NotImplementedError as e:
         print(f"skipping {c.spelling}: {e}")
         lines, types = rollback
   main = (f"# mypy: ignore-errors\nimport ctypes{', os' if any('os' in s for s in dll) else ''}\n"
     "from tinygrad.helpers import unwrap\nfrom tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR\n" + '\n'.join([*prolog,
-      *(["from ctypes.util import find_library"]*any('find_library' in s for s in dll)), *(["from tinygrad.runtime.support import objc"]*objc),
+      *(["from ctypes.util import find_library"]*any('find_library' in s for s in dll)),
       *(["def dll():",*flatten([[f"  try: return ctypes.CDLL(unwrap({d}){', use_errno=True' if use_errno else ''})",'  except: pass'] for d in dll]),
          "  return None", "dll = dll()\n"]*bool(dll)), *lines]) + '\n')
   macros = [r for m in macros if (r:=functools.reduce(lambda s,r:re.sub(r[0], r[1], s), rules + base_rules, m))]

tinygrad/runtime/support/objc.py

@@ -1,70 +0,0 @@
-import ctypes, ctypes.util, functools, sys
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING: id_ = ctypes.c_void_p
-else:
-  class id_(ctypes.c_void_p):
-    retain: bool = False
-    # This prevents ctypes from converting response to plain int, and dict.fromkeys() can use it to dedup
-    def __hash__(self): return hash(self.value)
-    def __eq__(self, other): return self.value == other.value
-    def __del__(self):
-      if self.retain and not sys.is_finalizing(): self.release()
-    def release(self): msg("release")(self)
-    def retained(self):
-      setattr(self, 'retain', True)
-      return self
-
-def returns_retained(f): return functools.wraps(f)(lambda *args, **kwargs: f(*args, **kwargs).retained())
-
-lib = ctypes.CDLL(ctypes.util.find_library('objc'))
-lib.sel_registerName.restype = id_
-getsel = functools.cache(lib.sel_registerName)
-lib.objc_getClass.restype = id_
-dispatch_data_create = ctypes.CDLL("/usr/lib/libSystem.dylib").dispatch_data_create
-dispatch_data_create.restype = id_
-dispatch_data_create = returns_retained(dispatch_data_create)
-
-def msg(sel:str, restype=id_, argtypes=[], retain=False, clsmeth=False):
-  # Using attribute access returns a new reference so setting restype is safe
-  (sender:=lib["objc_msgSend"]).restype, sender.argtypes = restype, [id_, id_]+argtypes if argtypes else []
-  def f(ptr, *args): return sender(ptr._objc_class_ if clsmeth else ptr, getsel(sel.encode()), *args)
-  return returns_retained(f) if retain else f
-
-if TYPE_CHECKING:
-  import _ctypes
-  class MetaSpec(_ctypes._PyCSimpleType):
-    _objc_class_: id_
-    def __getattr__(cls, nm:str) -> Any: ...
-else:
-  class MetaSpec(type(id_)):
-    def __new__(mcs, name, bases, dct):
-      cls = super().__new__(mcs, name, bases, {'_objc_class_': lib.objc_getClass(name.encode()), '_children_': set(), **dct})
-      cls._methods_, cls._classmethods_ = dct.get('_methods_', []), dct.get('_classmethods_', [])
-      return cls
-
-    def __setattr__(cls, k, v):
-      super().__setattr__(k, v)
-      if k in ("_methods_", "_classmethods_"):
-        for m in v: cls._addmeth(m, clsmeth=(v=="_classmethods_"))
-        for c in cls._children_: c._inherit(cls)
-      if k == "_bases_":
-        for b in v:
-          b._children_.add(cls)
-          cls._inherit(b)
-
-    def _inherit(cls, b):
-      for _b in getattr(b, "_bases_", []): cls._inherit(_b)
-      for m in getattr(b, "_methods_", []): cls._addmeth(m)
-      for m in getattr(b, "_classmethods_", []): cls._addmeth(m, True)
-      for c in cls._children_: c._inherit(cls)
-
-    def _addmeth(cls, m, clsmeth=False):
-      nm = m[0].strip(':').replace(':', '_')
-      if clsmeth: setattr(cls, nm, classmethod(msg(m[0], cls if m[1] == 'instancetype' else m[1],
-                                                   [cls if a == 'instancetype' else a for a in m[2]], *m[3:], clsmeth=True))) # type: ignore[misc]
-      else: setattr(cls, nm, msg(m[0], cls if m[1] == 'instancetype' else m[1], [cls if a == 'instancetype' else a for a in m[2]], *m[3:]))
-
-class Spec(id_, metaclass=MetaSpec):
-  if TYPE_CHECKING:
-    def __getattr__(self, nm:str) -> Any: ...

tinygrad/schedule/indexing.py

@@ -6,7 +6,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, g
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
 
-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
                      Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL}
 
@@ -16,10 +16,9 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
   for s in rb.src:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
-def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
+def realize_store(ctx:dict[UOp, None], a:UOp) -> None:
   if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
-  # if it's a kernel, we don't realize it
-  if a.src[1].op is not Ops.KERNEL: ctx[a] = None
+  ctx[a] = None
 
 pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
@@ -29,7 +28,7 @@ pm_generate_realize_map = PatternMatcher([
   # realize srcs of COPY, MSELECT, MSTACK
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # realize ASSIGN and input to assign (might be optimized out)
-  (UPat(Ops.ASSIGN, name="a"), realize_assign),
+  (UPat(Ops.STORE, name="a"), realize_store),
 ])
 
 @dataclass(frozen=True)
@@ -54,20 +53,23 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
   if x.op in {Ops.BUFFERIZE, Ops.INDEX}: return None
   if x.op is Ops.AFTER and x.src[1].op is Ops.KERNEL: return None
   new_srcs = []
-  for i,s in enumerate(x.src):
+  for s in x.src:
     new_src = s
-    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.AFTER and s.src[1].op is Ops.KERNEL):
+    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
       if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
     elif s in ctx.realize_map:
       realized_ranges = ctx.realize_map[s]
       assert isinstance(realized_ranges, list), "realize map must contain range list"
       closed_ranges = tuple([r for i,r in enumerate(ctx.range_map[s][1]) if i in realized_ranges])
-      # None in the device assigns it a number later
-      opts = BufferizeOpts(device=s.device) if len(ctx.range_map[s][1]) == len(realized_ranges) else BufferizeOpts(None, AddrSpace.LOCAL)
-      new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts, tag=s.tag if opts.addrspace == AddrSpace.GLOBAL else None)
-      if x in ctx.range_map:
-        # for scan we use the output ranges on the 2nd arg
-        new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][int(x.op is Ops.FOLD and i == 1)]) if i in realized_ranges])
+      if s.op is Ops.STORE:
+        # add the ends if this is a store
+        new_src = s.end(*[r for r in closed_ranges if r.op is Ops.RANGE])
+        del ctx.realize_map[s]
+      else:
+        # None in the device assigns it a number later
+        opts = BufferizeOpts(device=s.device) if len(ctx.range_map[s][1]) == len(realized_ranges) else BufferizeOpts(None, AddrSpace.LOCAL)
+        new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts, tag=s.tag if opts.addrspace == AddrSpace.GLOBAL else None)
+        if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][0]) if i in realized_ranges])
     new_srcs.append(new_src)
   # NOTE: do we need this?
   return x.replace(src=tns) if x.src != (tns:=tuple(new_srcs)) else None
@@ -86,15 +88,8 @@ def convert_reduce_axis_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
   ctx.range_map[ret] = ctx.range_map[x]
   return ret
 
-def add_ranges_to_scan(ctx:IndexingContext, x:UOp):
-  if x not in ctx.range_map: return None
-  new_ranges = [r for r,ar in zip(*ctx.range_map[x]) if r is not ar and r not in x.src]
-  ret = x.replace(src=x.src+tuple(new_ranges))
-  ctx.range_map[ret] = ctx.range_map[x]
-  return ret
-
 def remove_movement_op_after_rangeify(ctx:IndexingContext, x:UOp):
-  if x in ctx.range_map or x.src[0].op is Ops.INDEX: return x.src[0]
+  if (x in ctx.range_map or x.src[0].op is Ops.INDEX): return x.src[0]
 
 def add_third_op_to_assign_to_track_shape(ctx:IndexingContext, assign:UOp):
   if assign.src[1].op is Ops.KERNEL: return None
@@ -106,8 +101,6 @@ def add_third_op_to_assign_to_track_shape(ctx:IndexingContext, assign:UOp):
 pm_apply_rangeify = PatternMatcher([
   # REDUCE_AXIS -> REDUCE
   (UPat(Ops.REDUCE_AXIS, name="x"), convert_reduce_axis_to_reduce_with_ranges),
-  # SCAN -> SCAN (with new ranges)
-  (UPat(Ops.FOLD, name="x"), add_ranges_to_scan),
   # PAD -> WHERE
   (UPat(Ops.PAD, name="x"), convert_pad_to_where_to_keep_behavior_local),
   # add third op to assign
@@ -187,7 +180,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       # mark all ranges as ended
       assert rctx.realize_map[x] is None
       rctx.realize_map[x] = list(range(len(x.shape)))
-    elif x.op in {Ops.MSTACK, Ops.MSELECT}:
+    elif x.op in {Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
       # treat MSTACK/MSELECT like SINK
       continue
     elif len(consumer_rngs) == 0:
@@ -220,6 +213,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       out_rngs = tuple(_out_rngs)
 
       # we have to (partially) realize here if there's new ranges
+      print(_realize_axis)
       if len(_realize_axis): rctx.realize_map[x] = _realize_axis
 
     # if this element is a reduce and there's ended ranges, we might have to end some other ranges
@@ -255,9 +249,6 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # REDUCE_AXIS creates ranges for the axes it is reducing
     if x.op is Ops.REDUCE_AXIS:
       rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape)))
-    if x.op is Ops.FOLD:
-      rngs = tuple(rctx.new_range(s, axistype=AxisType.FOLD) if resolve(x.src[1].shape[i] == 1) else r \
-                   for i,(r,s) in enumerate(zip(rngs, x.src[0].shape)))
 
     if debug:
       realized_ranges = rctx.realize_map.get(x, None)

tinygrad/schedule/rangeify.py

@@ -396,7 +396,6 @@ def handle_after(ctx:LocalAddBufferContext, after:UOp):
   return buf
 
 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
-  if r.arg[-1] == AxisType.OUTER: return None
   if r.tag != (): return None
   ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None)
   ctx.range += 1
@@ -470,10 +469,7 @@ pm_add_range_tags = PatternMatcher([
 ])
 
 def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
-  if len([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
-
-  # ends of outer range don't go in kernels
-  if x.op is Ops.END and x.src[1].op is Ops.RANGE and x.src[1].arg[-1] == AxisType.OUTER: return None
+  if len(x.ranges): return None
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()
@@ -538,6 +534,12 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
 
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
 
+  # link any movementops with tags to sink, and remove the tags from other parts of the graph
+  # we add "None" to the tag so it's deduped from the movementop that exists in the graph
+  #tagged_mops = [x.replace(tag=x.tag+(None,)) for x in tsink.toposort() if x.op in GroupOp.Movement and x.tag is not None]
+  #tsink = tsink.replace(src=tsink.src+tuple(tagged_mops))
+  #tsink = tsink.substitute({})
+
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)
 

tinygrad/tensor.py

@@ -299,7 +299,7 @@ class Tensor(OpMixin):
     assert self.shape == x.shape, f"assign shape mismatch {self.shape} != {x.shape}"
     assert self.device == x.device, f"assign device mismatch {self.device} != {x.device}"
     assert self.dtype == x.dtype, f"assign dtype mismatch {self.dtype} != {x.dtype}"
-    return self.replace(self._apply_uop(UOp.assign, x))
+    return self.replace(self._apply_uop(lambda x,y: x.after(x.store(y)), x))
 
   def detach(self) -> Tensor:
     """
@@ -1509,9 +1509,6 @@ class Tensor(OpMixin):
 
   # ***** reduce ops *****
 
-  def fold(self, init:Tensor, *ranges:UOp) -> Tensor:
-    return self._apply_uop(UOp.fold, init, extra_args=ranges)
-
   def _reduce(self, op:Ops, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
     axis = tuple(self._resolve_dim(x) for x in (range(self.ndim) if axis is None else make_tuple(axis, 1)))
     if self.ndim == 0: axis = ()

tinygrad/uop/__init__.py

@@ -92,9 +92,6 @@ class Ops(FastEnum):
   # reduce
   REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto()
 
-  # scan
-  FOLD = auto()
-
   # errors/placeholders
   REWRITE_ERROR = auto(); SENTINEL = auto()
 

tinygrad/uop/ops.py

@@ -13,20 +13,14 @@ if TYPE_CHECKING:
 
 class AxisType(Enum):
   def __repr__(self): return str(self)
-  GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); FOLD = auto() # noqa: E702
-  UPCAST = auto(); UNROLL = auto() # noqa: E702
-  THREAD = auto(); OUTER = auto() # noqa: E702
+  GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
+  THREAD = auto()
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
-                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.FOLD: "F", AxisType.UNROLL: "r", AxisType.OUTER: "O"}
+                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
 axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
-               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.FOLD: "red", AxisType.UNROLL: "magenta",
-               AxisType.OUTER: "green"}
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
 
-# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
-               AxisType.GROUP_REDUCE: 2, AxisType.FOLD: 4, AxisType.REDUCE: 5, AxisType.UNROLL: 6, AxisType.OUTER: -2}
-
-range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.FOLD: 2}
+range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1}
 
 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> ConstType: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
@@ -219,10 +213,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.BUFFERIZE: return tuple([int(r.vmax+1) for r in self.src[1:]])
       case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
 
-      # shape of init
-      case Ops.FOLD:
-        return self.src[1]._shape
-
       # passthrough ops
       case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END:
         return self.src[0]._shape
@@ -447,7 +437,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     assert self.dtype.scalar() is dtypes.index, "Can only call get_valid on index dtype"
     return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid)
   def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
-  def fold(self, *src:UOp, **kwargs): return UOp(Ops.FOLD, self.dtype, (self,)+src, **kwargs)
 
   def is_contiguous(self):
     # TODO: this is is_realized

tinygrad/uop/spec.py

@@ -40,9 +40,6 @@ shared_spec = PatternMatcher([
     rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \
       all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)),
   (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:]) or None),
-
-  # FOLD, 2 ops + ranges
-  (UPat(Ops.FOLD, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[2:])),
 ])
 
 # ***** UOp spec in the Tensor graph *****
@@ -175,7 +172,7 @@ kernel_spec = PatternMatcher([
   # bufferize can be on anything
   (UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: True),
 
-  # reduce/fold must be on ranges
+  # reduce must be on ranges
   (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
 ])+movement_ops+shared_codegen_spec+shared_spec
 

tinygrad/viz/serve.py

@@ -17,7 +17,7 @@ from tinygrad.dtype import dtypes
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
                Ops.DEFINE_GLOBAL:"#cb9037", **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B",
                Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
-               Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", Ops.FOLD: "#FF7B7B",
+               Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
                Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",