Merge branch 'master' into call_inline

cl_compiler instead of compiler because overriding Compiled.compiler seems more confusing

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '15'
+  CACHE_VERSION: '16'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -581,7 +581,7 @@ jobs:
         load: true
         tags: qemu-hexagon:latest
         cache-from: type=gha
-        cache-to: type=gha,mode=min
+        cache-to: ${{ github.event_name != 'pull_request' && 'type=gha,mode=min' || '' }}
     - name: Set MOCKDSP env
       run: printf "MOCKDSP=1" >> $GITHUB_ENV
     - name: Run test_tiny on DSP
@@ -739,7 +739,7 @@ jobs:
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run pytest (cuda)
         # skip multitensor because it's slow
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_multitensor.py --durations=20
       - name: Run TestOps.test_add with PMA
         run: VIZ=-1 PMA=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
       - name: Run process replay tests

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -2,6 +2,7 @@
 
 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 
@@ -12,7 +13,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=8 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=1
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -20,13 +21,14 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
 
-export SEED=5760
+export SEED=${SEED:-5760}
 
-export JITBEAM=3
+export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 
 export FAKEDATA=1 BENCHMARK=10 LLAMA_LAYERS=2

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -13,7 +13,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -21,9 +21,10 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
 
 export SEED=${SEED:-5760}
 

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -1,4 +1,6 @@
 #!/bin/bash
 export BENCHMARK=5
+export EVAL_BS=0
 export VIZ=${VIZ:--1}
 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20

extra/gemm/asm/cdna/asm.py

@@ -14,7 +14,12 @@ GEMM_ARGS = {
   (8192, 8192, 8192): (256, 128, 131072),
   (4096, 4096, 4096): (256, 64, 16384),
   (4096, 14336, 4096): (256, 64, 57344),
+  (4096, 14336, 8192): (256, 128, 114688),
   (4096, 4096, 14336): (256, 224, 57344),
+  (14336, 4096, 8192): (256, 128, 114688),
+  (4096, 8192, 14336): (256, 224, 114688),
+  (4096, 4096, 8192): (256, 128, 32768),
+  (4096, 8192, 4096): (256, 64, 32768),
 }
 ITERS_ARGS = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
 

extra/gemm/asm/cdna/gemm.py

@@ -30,13 +30,13 @@ atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters
 def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
   if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
   if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
-  # only sharding on the batch is tested, others might work too
-  if isinstance(a.device, tuple) and not (a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None):
-    return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
   batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
   N = b.shape[1]
+  # only sharding on the batch or K is tested, others might work too
   if isinstance(a.device, tuple):
-    batch //= len(a.device)
+    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
+    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
     dname = a.device[0]
   else: dname = a.device
   arch = getattr(Device[dname].renderer, "arch", "")
@@ -65,6 +65,8 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
   out, a, b = kernel.src[1:]
   assert all_same([gradient.device, a.device, b.device, out.device])
   a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
+  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
+  g_t = g_t[:a.shape[0]]
   grad_a = (g_t @ b_t.T).uop
   grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
   return (None, grad_a, grad_b)
@@ -80,9 +82,10 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
   batch, M, K = a.shape
   N = b.shape[1]
   is_multi = isinstance(a.device, tuple)
+  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
 
   if is_multi:
-    out = Tensor(Tensor.empty(batch//len(a.device), M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+    out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
   else:
     out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
 
@@ -93,4 +96,5 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
     out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), grad_fxn=custom_gemm_bw)[0]
   else:
     out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
+  if k_sharded: out = out.sum(0)
   return out.squeeze(0) if squeeze else out

extra/models/llama.py

@@ -89,13 +89,13 @@ class Attention:
       assert start_pos == 0
       keys, values = xk, xv
 
-    if Tensor.training:
-      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
-      attn = xq.scaled_dot_product_attention(keys, values, is_causal=True, enable_gqa=True).transpose(1, 2)
-    else:
+    if self.max_context:
       keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
       xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
       attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    else:
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, is_causal=True, enable_gqa=True).transpose(1, 2)
     if getenv("STUB_ATTENTION"):
       from tinygrad.uop.ops import UOp, KernelInfo
       def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
@@ -203,7 +203,7 @@ class Transformer:
     h = self.tok_embeddings(tokens)
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]
 
-    if not Tensor.training and seqlen > 1:
+    if self.max_context != 0 and seqlen > 1:
       mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1)
     else: mask = None
     for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)

extra/viz/cli.py

@@ -31,6 +31,7 @@ if __name__ == "__main__":
   g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
   g_profile = parser.add_argument_group("profile options")
   g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
+  g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
   g_rewrites = parser.add_argument_group("rewrites options")
   g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
                           help="Select an item within the chosen kernel (optional name, default: only list names)")
@@ -72,9 +73,9 @@ if __name__ == "__main__":
           total += et
     if agg and total > 0:
       items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
-      sel = items[:10]
+      sel = items if args.top == -1 else items[:args.top]
       table = [[name, time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in sel]
-      if (other:=items[len(sel):]):
+      if args.top != -1 and (other:=items[len(sel):]):
         other_t = total-sum(t for _, (t, _) in sel)
         table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
       print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))

test/device/test_metal.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad.device import CompileError, Device, Compiler
+from tinygrad.device import CompileError, Device
 if Device.DEFAULT=="METAL":
   from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
@@ -48,28 +48,4 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
 """)
     with self.assertRaises(RuntimeError):
       compiled = compiled[:40] # corrupt the compiled program
-      MetalProgram(device, "r_5", compiled)
-
-  def test_program_w_empty_compiler(self):
-    device = MetalDevice("metal")
-    compiler = Compiler(device)
-    compiled = compiler.compile("""
-#include <metal_stdlib>
-kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]){
-  data0[0] = 0;
-}
-""")
-    MetalProgram(device, "r_5", compiled)
-
-  def test_bad_program_w_empty_compiler(self):
-    device = MetalDevice("metal")
-    compiler = Compiler(device)
-    # this does not raise
-    compiled = compiler.compile("""
-#include <metal_stdlib>
-kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]){
-  invalid codes;
-}
-""")
-    with self.assertRaises(RuntimeError):
       MetalProgram(device, "r_5", compiled)

test/device/test_ocl.py

@@ -1,9 +1,19 @@
 import unittest
+from unittest.mock import patch
 from tinygrad import Device
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes
 from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler, CLProgram
 
+@unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
+class TestCLCompileCache(unittest.TestCase):
+  def test_compile_cached(self):
+    device = Device[Device.DEFAULT]
+    src = "__kernel void cached_test(__global int* a) { a[0] = 1; }"
+    CLProgram(device, name="cached_test", lib=src.encode())
+    with patch.object(CLCompiler, 'compile', side_effect=RuntimeError("compile should not be called on cache hit")):
+      CLProgram(device, name="cached_test", lib=src.encode())
+
 @unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
 class TestCLError(unittest.TestCase):
   @unittest.skip("allocates tons of memory")
@@ -17,7 +27,7 @@ class TestCLError(unittest.TestCase):
   def test_invalid_kernel_name(self):
     device = Device[Device.DEFAULT]
     with self.assertRaises(RuntimeError) as err:
-      CLProgram(device, name="", lib=CLCompiler(device, "test").compile("__kernel void test(__global int* a) { a[0] = 1; }"))
+      CLProgram(device, name="", lib="__kernel void test(__global int* a) { a[0] = 1; }".encode())
     assert str(err.exception) == "OpenCL Error -46: CL_INVALID_KERNEL_NAME"
 
   def test_unaligned_copy(self):

test/null/test_gc.py

@@ -72,7 +72,6 @@ class TestGC(unittest.TestCase):
     ys = y.schedule()
     del x
     run_schedule(ys)
-    np.testing.assert_equal(y.numpy(), np.full((256,), 2))
     self.assertEqual(bufs_allocated()-init, 1)
     del y
     self.assertEqual(bufs_allocated()-init, 0)

test/test_edgecases.py

@@ -175,7 +175,6 @@ class TestZeroFolding(unittest.TestCase):
 class TestAssignIssues(unittest.TestCase):
   # these are good failures. i'm not sure we need more, but we need to fix these.
 
-  @unittest.expectedFailure
   def test_assign_permuted_view_constant(self):
     # assigning to a permuted view should modify the underlying tensor
     arr = np.arange(6).reshape(2, 3).astype(np.float32)
@@ -185,7 +184,6 @@ class TestAssignIssues(unittest.TestCase):
     t.permute(1, 0).assign(Tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]))
     np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
 
-  @unittest.expectedFailure
   def test_assign_shrink_view_constant(self):
     # assigning to a shrunk view should update the base tensor
     arr = np.arange(9).reshape(3, 3).astype(np.float32)

test/test_fusion_op.py

@@ -1,68 +0,0 @@
-import unittest
-import time
-import numpy as np
-from tinygrad import Tensor, dtypes
-from tinygrad.engine.realize import run_schedule
-
-class TestFusionOp(unittest.TestCase):
-  def test_contiguous_add(self):
-    def test(contig=False):
-      bt = Tensor(np.arange(16), dtype=dtypes.float32).reshape(4,4)
-      x = bt.permute(1,0)
-      if contig: x = x.contiguous()
-      return (x.permute(1,0) + bt).data()
-    assert test() == test(True)
-
-  def test_expand_fuse(self):
-    bt = Tensor(np.ones((10, 1)), dtype=dtypes.float32)
-    out = (bt*2).expand(10,10).sum(1)
-    sched = out.schedule()
-    run_schedule(sched)
-    outd = out.tolist()
-    assert all(x == 20.0 for x in outd)
-
-  def test_recursive_add(self):
-    st = time.perf_counter()
-    a = Tensor([1,2,3,4])
-    for _ in range(24): a = a + a
-    sched = a.schedule()
-    sched[-1].lower()
-    self.assertLess(time.perf_counter()-st, 2.0)
-    assert len(sched[-1].prg.p.src.splitlines()) < 250
-
-  def test_recursive_add_cmp(self):
-    st = time.perf_counter()
-    a = Tensor([1,2,3,4])
-    for _ in range(24): a = a + a
-    sched1 = a.schedule()
-    b = Tensor([1,2,3,4])
-    for _ in range(24): b = b + b
-    sched2 = b.schedule()
-    c = Tensor([1,2,3,4])
-    for _ in range(23): c = c + c
-    sched3 = c.schedule()
-    self.assertEqual(sched1[-1].ast, sched2[-1].ast)
-    with self.assertRaises(AssertionError): self.assertEqual(sched1[-1].ast, sched3[-1].ast)
-    self.assertLess(time.perf_counter()-st, 2.0)
-
-  def test_recursive_pad(self):
-    st = time.perf_counter()
-    val = 1.0
-    a = Tensor(val)
-    for _ in range(24): a = Tensor.stack(a, a)[0]
-    sched = a.schedule()
-    self.assertEqual(len(sched), 0)
-    self.assertLess(time.perf_counter()-st, 2.0)
-
-  def test_recursive_reshape(self):
-    st = time.perf_counter()
-    a = Tensor.empty(32, 32).realize()
-    b = Tensor.empty(16, 2).realize()
-    r = a.sum(1)
-    for _ in range(24): r = r.reshape(16, 2) + b
-    sched = r.schedule()
-    self.assertEqual(len(sched), 1)
-    self.assertLess(time.perf_counter()-st, 2.0)
-
-if __name__ == '__main__':
-  unittest.main(verbosity=2)

test/test_jit_footguns.py

@@ -111,30 +111,18 @@ class TestJitFootguns(unittest.TestCase):
       self.assertEqual(first.numpy().item(), expected_first)
       buf = new_buf
 
-  def test_slice_assign_requires_realize(self):
-    """Slice assign then read from same buffer - assign isn't connected to read without explicit realize()."""
+  def test_slice_assign_works_without_realize(self):
+    """Slice assign then read from same buffer - pending assigns are side-realized."""
     from tinygrad import Variable
     v_pos = Variable("pos", 0, 3)
-
-    # without .realize() after assign, the read doesn't see the assigned values
     cache = Tensor.zeros(4, 4).contiguous().realize()
     @TinyJit
-    def f_broken(pos):
+    def f(pos):
       cache[pos:pos+1, :].assign(Tensor.ones(1, 4))
       return cache.sum().realize()
     for i in range(4):
       cache.assign(Tensor.zeros(4, 4)).realize()
-      self.assertEqual(f_broken(v_pos.bind(i)).item(), 0.0)  # should be 4.0!
-
-    # workaround: add .realize() after assign
-    cache2 = Tensor.zeros(4, 4).contiguous().realize()
-    @TinyJit
-    def f_fixed(pos):
-      cache2[pos:pos+1, :].assign(Tensor.ones(1, 4)).realize()
-      return cache2.sum().realize()
-    for i in range(4):
-      cache2.assign(Tensor.zeros(4, 4)).realize()
-      self.assertEqual(f_fixed(v_pos.bind(i)).item(), 4.0)
+      self.assertEqual(f(v_pos.bind(i)).item(), 4.0)
 
   def test_symbolic_pad_view_frozen(self):
     """Symbolic pad view has BIND values baked in at capture time. TODO: pad should be captured in jit."""

test/test_outerworld.py

@@ -57,7 +57,7 @@ class TestOuterRange(unittest.TestCase):
     out.realize()
 
     # TODO: testing allclose
-    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+    assert Tensor.allclose(ref, out, atol=1e-6), f"max diff {(ref-out).abs().max().item()}"
 
 class TestOuterScan(unittest.TestCase):
   def _test_scan(self):
@@ -85,7 +85,7 @@ class TestOuterScan(unittest.TestCase):
     out.realize()
 
     # TODO: testing allclose
-    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+    assert Tensor.allclose(ref, out, atol=1e-5), f"max diff {(ref-out).abs().max().item()}"
 
 class TestOuterworld(unittest.TestCase):
   def test_range_plus_1(self):
@@ -166,7 +166,7 @@ class TestVmap(unittest.TestCase):
     out.realize()
 
     # TODO: testing allclose
-    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+    assert Tensor.allclose(ref, out, atol=1e-6), f"max diff {(ref-out).abs().max().item()}"
   def test_vmap_inner_fuse(self): self.test_vmap_inner(fuse=True)
   def test_vmap_outer(self): self.test_vmap_inner(AxisType.OUTER)
   def test_vmap_outer_fuse(self): self.test_vmap_inner(AxisType.OUTER, fuse=True)

test/test_setitem.py

@@ -1,8 +1,6 @@
-import unittest
-import random
-from os import getenv
+import unittest, random
 from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
-from tinygrad.helpers import Context
+from tinygrad.helpers import Context, getenv
 import numpy as np
 
 class TestSetitem(unittest.TestCase):
@@ -13,7 +11,7 @@ class TestSetitem(unittest.TestCase):
       ((6,6), (slice(2,4), slice(3,5)), 1.0),
       ((6,6), (3, 4), 1.0),
       ((6,6), (3, None, 4, None), 1.0),
-      ((4,4,4,4), (Ellipsis, slice(1,3), slice(None)), Tensor(4)),
+      ((4,4,4,4), (Ellipsis, slice(1,3), slice(None)), Tensor(4.0)),
       ((4,4,4,4), (Ellipsis, slice(1,3)), 4),
       ((4,4,4,4), (2, slice(1,3), None, 1), 4),
       ((4,4,4,4), (slice(1,3), slice(None), slice(0,4,2)), 4),
@@ -50,6 +48,10 @@ class TestSetitem(unittest.TestCase):
         t[1] = v
         self.assertEqual(t.dtype, dt)
 
+  def test_setitem_dtype_mismatch(self):
+    t = Tensor.zeros(6, dtype=dtypes.float).contiguous().realize()
+    with self.assertRaises(RuntimeError): t[2:4] = Tensor([1, 2], dtype=dtypes.int)
+
   def test_setitem_into_noncontiguous(self):
     t = Tensor.ones(4)
     with self.assertRaises(RuntimeError): t[1] = 5
@@ -109,8 +111,6 @@ class TestSetitem(unittest.TestCase):
   def test_setitem_consecutive_inplace_operator(self):
     t = Tensor.arange(4).reshape(2, 2).contiguous()
     t[1] += 2
-    t = t.contiguous()
-    # TODO: RuntimeError: can't double realize in one schedule
     t[1] -= 1
     np.testing.assert_allclose(t.numpy(), [[0, 1], [3, 4]])
 
@@ -182,7 +182,7 @@ class TestSetitem(unittest.TestCase):
     t[:-1] = t[1:]
     self.assertEqual(t.tolist(), [[2.0], [1.0], [1.0]])
 
-  # TODO: WEBGPU pipeline validation error
+  # TODO: WEBGPU pipeline validation error. this generates (1==gidx0)|(2==gidx0)|(3==gidx0)|(4==gidx0)|(5==gidx0) ...
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU pipeline validation error")
   def test_setitem_big(self):
     idx_size, val = 256, 4
@@ -193,23 +193,23 @@ class TestSetitem(unittest.TestCase):
 
   def test_setitem_advanced_indexing(self):
     # Example from https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing
-    t = Tensor.zeros(10,20,30,40,50).contiguous()
+    t = Tensor.zeros(10,20,30,40,50, dtype=dtypes.int).contiguous()
     ind_1 = Tensor([5,3,7,8])
     ind_2 = Tensor([[[0],[1],[2]],[[3],[4],[5]]])
     v = Tensor.arange(2*3*4*10*30*50).reshape(2,3,4,10,30,50)
     t[:, ind_1, :, ind_2, :] = v
-    n = np.zeros((10,20,30,40,50))
+    n = np.zeros((10,20,30,40,50), dtype=np.int32)
     n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
-    np.testing.assert_allclose(t.numpy(), n)
+    np.testing.assert_equal(t.numpy(), n)
 
   def test_setitem_2d_tensor_indexing(self):
-    t = Tensor.zeros(2).contiguous()
+    t = Tensor.zeros(2, dtype=dtypes.int).contiguous()
     index = Tensor([[0, 1], [1,0]])
     v = Tensor.arange(2*2).reshape(2, 2).contiguous()
     t[index] = v
-    n = np.zeros((2,))
+    n = np.zeros((2,), dtype=np.int32)
     n[index.numpy()] = v.numpy()
-    np.testing.assert_allclose(t.numpy(), n)
+    np.testing.assert_equal(t.numpy(), n)
 
   @unittest.skip("slow")
   def test_setitem_tensor_indexing_fuzz(self):

test/testextra/test_asm_gemm.py

@@ -9,24 +9,24 @@ from test.helpers import needs_second_gpu
 # Use NULL=1 EMULATE=AMD_CDNA4 to also test the assembly
 def is_cdna4(): return getattr(Device[Device.DEFAULT].renderer, "arch", "").startswith("gfx950")
 
-def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
+def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None:
   Tensor.manual_seed(0)
-  a_rand = Tensor.randn((batch, M, K), dtype=dtypes.float).sub(0.5).cast(dtype)
-  b_rand = Tensor.randn((K, N), dtype=dtypes.float).sub(0.5).cast(dtype)
+  a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(dtype)
+  b_rand = Tensor.randn(b_shape, dtype=dtypes.float).sub(0.5).cast(dtype)
   with Context(DEBUG=0):
     Tensor.realize(a_rand, b_rand)
 
   devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus)) if (multi:=gpus>1) else None
 
   a, b = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
-  if multi: a, b = a.shard(devs, axis=0), b.shard(devs, axis=None)
+  if multi: a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard)
   with Context(ASM_GEMM=1):
     tst = asm_gemm(a, b)
     tst.sum().backward()
   Tensor.realize(tst, a.grad, b.grad)
 
   a_ref, b_ref = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
-  if multi: a_ref, b_ref = a_ref.shard(devs, axis=0), b_ref.shard(devs, axis=None)
+  if multi: a_ref, b_ref = a_ref.shard(devs, axis=a_shard), b_ref.shard(devs, axis=b_shard)
   with Context(ASM_GEMM=0):
     ref = asm_gemm(a_ref, b_ref)
     ref.sum().backward()
@@ -34,10 +34,18 @@ def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:i
 
   # no validation on the NULL device
   if a_rand.device.startswith("NULL"): return None
+  atol, rtol = (1e-2, 1e-3)
   with Context(DEBUG=0):
-    assert (tst - ref).square().max().float().item() < 1e-6, "forward mismatch"
-    assert (a.grad - a_ref.grad).square().max().float().item() < 1e-3, "grad_a mismatch"
-    assert (b.grad - b_ref.grad).square().max().float().item() < 1e-3, "grad_b mismatch"
+    assert tst.allclose(ref, atol=atol, rtol=rtol), "forward mismatch"
+    assert a.grad.allclose(a_ref.grad, atol=atol, rtol=rtol), "grad_a mismatch"
+    assert b.grad.allclose(b_ref.grad, atol=atol, rtol=rtol), "grad_b mismatch"
+
+
+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
+  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
+
+def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=8) -> None:
+  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)
 
 # 128x smaller than usual
 # uses the UOp GEMM, runs on non CDNA4 and CI
@@ -50,6 +58,8 @@ class TestGemm(unittest.TestCase):
   def test_gemm_batched(self): verify_asm_gemm(2, 64, 32, 32)
   @needs_second_gpu
   def test_gemm_multi(self): verify_asm_gemm(2, 64, 32, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_k_sharded(self): verify_asm_gemm_k_sharded(64, 64, 2*64, gpus=2)
 
 # uses the Asm GEMM on CDNA4 only for speed reasons
 class TestGemmLarge(unittest.TestCase):
@@ -70,6 +80,12 @@ class TestGemmLarge(unittest.TestCase):
   def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
   @unittest.skip("disabled, asm in this shape is slower than tinygrad")
   def test_gemm7(self): verify_asm_gemm(1, 8192, 128256, 4096)
+  def test_gemm8(self): verify_asm_gemm(1, 4096, 14336, 8192)
+  def test_gemm9(self): verify_asm_gemm(8, 4096, 14336, 8192, dtype=dtypes.bfloat16, gpus=8)
+  def test_gemm10(self): verify_asm_gemm(1, 4096, 8192, 4096)
+  def test_k_sharded_1(self): verify_asm_gemm_k_sharded(14336, 4096, 8*8192, gpus=8)
+  def test_k_sharded_2(self): verify_asm_gemm_k_sharded(4096, 14336, 8*8192, gpus=8)
+  def test_k_sharded_3(self): verify_asm_gemm_k_sharded(4096, 4096, 8*8192, gpus=8)
   def test_gemm_unsupported(self):
     with self.assertRaisesRegex(AssertionError, "shape not supported"):
       verify_asm_gemm(8, 1024, 1024, 4096, gpus=8)

test/unit/test_assign.py

@@ -560,30 +560,16 @@ class TestAssignOrdering(unittest.TestCase):
 
   def test_overlapping_slice_assigns(self):
     """Overlapping slice assigns - later write should win for overlapping elements."""
-    # without .realize(): assigns not executed, buffer stays zeros
     buf = Tensor.zeros(8).contiguous().realize()
     buf[0:4].assign(Tensor.ones(4))
     buf[2:6].assign(Tensor.ones(4) * 2)
-    np.testing.assert_equal(buf.numpy(), [0,0,0,0,0,0,0,0])  # TODO: wrong! should be [1,1,2,2,2,2,0,0]
-
-    # with .realize(): assigns execute in order
-    buf = Tensor.zeros(8).contiguous().realize()
-    buf[0:4].assign(Tensor.ones(4)).realize()
-    buf[2:6].assign(Tensor.ones(4) * 2).realize()
     np.testing.assert_equal(buf.numpy(), [1,1,2,2,2,2,0,0])
 
   def test_overlapping_slice_assigns_reverse(self):
     """Overlapping slice assigns in reverse order."""
-    # without .realize(): assigns not executed
     buf = Tensor.zeros(8).contiguous().realize()
     buf[2:6].assign(Tensor.ones(4) * 2)
     buf[0:4].assign(Tensor.ones(4))
-    np.testing.assert_equal(buf.numpy(), [0,0,0,0,0,0,0,0])  # TODO: wrong! should be [1,1,1,1,2,2,0,0]
-
-    # with .realize(): assigns execute in order
-    buf = Tensor.zeros(8).contiguous().realize()
-    buf[2:6].assign(Tensor.ones(4) * 2).realize()
-    buf[0:4].assign(Tensor.ones(4)).realize()
     np.testing.assert_equal(buf.numpy(), [1,1,1,1,2,2,0,0])
 
   def test_read_between_writes(self):
@@ -619,26 +605,14 @@ class TestAssignOrdering(unittest.TestCase):
 
   def test_slice_write_then_full_read(self):
     """Write to slice, then read full buffer."""
-    # without .realize(): orphan slice assign not triggered by .numpy()
     buf = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
     buf[1:3].assign(Tensor([5, 6]))
-    np.testing.assert_equal(buf.numpy(), [0, 0, 0, 0])  # TODO: wrong! should be [0, 5, 6, 0]
-
-    # with .realize(): assign executes
-    buf = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
-    buf[1:3].assign(Tensor([5, 6])).realize()
     np.testing.assert_equal(buf.numpy(), [0, 5, 6, 0])
 
   def test_chained_slice_copies(self):
     """Copy from one slice to another within same buffer."""
-    # without .realize(): orphan slice assign not triggered
     buf = Tensor([1, 2, 3, 4, 5, 6, 7, 8]).contiguous().realize()
     buf[4:8].assign(buf[0:4].contiguous())
-    np.testing.assert_equal(buf.numpy(), [1, 2, 3, 4, 5, 6, 7, 8])  # TODO: wrong! should be [1,2,3,4,1,2,3,4]
-
-    # with .realize(): assign executes
-    buf = Tensor([1, 2, 3, 4, 5, 6, 7, 8]).contiguous().realize()
-    buf[4:8].assign(buf[0:4].contiguous()).realize()
     np.testing.assert_equal(buf.numpy(), [1, 2, 3, 4, 1, 2, 3, 4])
 
   def test_swap_slices(self):
@@ -661,16 +635,9 @@ class TestAssignOrdering(unittest.TestCase):
 
   def test_reduction_after_partial_assign(self):
     """Reduction over buffer after partial assign - must see the assigned values."""
-    # without .realize(): orphan slice assign not triggered by reduction
     buf = Tensor.zeros(4, 4).contiguous().realize()
     buf[0:2, :].assign(Tensor.ones(2, 4))  # top half = 1
     total = buf.sum()
-    self.assertEqual(total.item(), 0)  # TODO: wrong! should be 8 (2*4 ones)
-
-    # with .realize(): assign executes before reduction
-    buf = Tensor.zeros(4, 4).contiguous().realize()
-    buf[0:2, :].assign(Tensor.ones(2, 4)).realize()
-    total = buf.sum()
     self.assertEqual(total.item(), 8)
 
   def test_multiple_reductions_different_views(self):
@@ -734,34 +701,18 @@ class TestAssignOrdering(unittest.TestCase):
   def test_variable_slice_ordering(self):
     """Variable-indexed slices - tests symbolic dependency tracking."""
     v_i = Variable("i", 0, 3)
-
-    # without .realize(): orphan slice assigns not triggered
     buf = Tensor.zeros(4, 4).contiguous().realize()
     buf[v_i.bind(0):v_i.bind(0)+1, :].assign(Tensor.ones(1, 4))
-    row0_sum = buf[0:1, :].sum()
-    self.assertEqual(row0_sum.item(), 0)  # TODO: wrong! should be 4
-
-    # with .realize(): assigns execute
-    buf = Tensor.zeros(4, 4).contiguous().realize()
-    buf[v_i.bind(0):v_i.bind(0)+1, :].assign(Tensor.ones(1, 4)).realize()
-    row0_sum = buf[0:1, :].sum()
-    buf[v_i.bind(1):v_i.bind(1)+1, :].assign(Tensor.ones(1, 4) * 2).realize()
-    row1_sum = buf[1:2, :].sum()
-    self.assertEqual(row0_sum.item(), 4)
-    self.assertEqual(row1_sum.item(), 8)
+    buf[v_i.bind(1):v_i.bind(1)+1, :].assign(Tensor.ones(1, 4) * 2)
+    self.assertEqual(buf[0:1, :].sum().item(), 4)
+    self.assertEqual(buf[1:2, :].sum().item(), 8)
 
   def test_multiple_slice_assigns_then_read(self):
-    """Multiple non-overlapping slice assigns then read - RAW dependencies must ensure all writes complete before read."""
+    """Multiple non-overlapping slice assigns then read."""
     buf = Tensor.zeros(4).contiguous().realize()
     buf[0:1].assign(Tensor.ones(1))
     buf[1:2].assign(Tensor.full((1,), 2.0))
     buf[2:3].assign(Tensor.full((1,), 3.0))
-    self.assertEqual(buf.sum().realize().item(), 0.0)  # TODO: wrong! should be 1 + 2 + 3 + 0 = 6
-
-    buf = Tensor.zeros(4).contiguous().realize()
-    buf[0:1].assign(Tensor.ones(1)).realize()
-    buf[1:2].assign(Tensor.full((1,), 2.0)).realize()
-    buf[2:3].assign(Tensor.full((1,), 3.0)).realize()
     self.assertEqual(buf.sum().realize().item(), 6.0)
 
 if __name__ == "__main__":

test/unit/test_call.py

@@ -13,7 +13,7 @@ class TestCall(unittest.TestCase):
     # we define a plus function
     plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
 
-    c = Tensor.call(a, b, fxn=plus_fxn)
+    c = Tensor.call(a, b, fxn=plus_fxn, inline=True)
     np.testing.assert_equal(c.numpy(), (a+b).numpy())
 
   def test_call_plus_backward(self):
@@ -30,7 +30,7 @@ class TestCall(unittest.TestCase):
 
     # we define a plus function
     plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
-    c = Tensor.call(a, b, fxn=plus_fxn, grad_fxn=grad_fxn)
+    c = Tensor.call(a, b, fxn=plus_fxn, grad_fxn=grad_fxn, inline=True)
     c.mean().backward()
 
     np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
@@ -46,7 +46,7 @@ class TestCall(unittest.TestCase):
     a.grad, b.grad = None, None
 
     plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
-    c = Tensor.call(a, b, fxn=plus_fxn)
+    c = Tensor.call(a, b, fxn=plus_fxn, inline=True)
     c.mean().backward()
 
     np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
@@ -57,7 +57,7 @@ class TestCall(unittest.TestCase):
     a = Tensor.randn(M, K)
     b = Tensor.randn(K, N)
     Tensor.realize(a, b)
-    c = Tensor.call(a, b, fxn=a.as_param(0) @ b.as_param(1))
+    c = Tensor.call(a, b, fxn=a.as_param(0) @ b.as_param(1), inline=True)
     np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
 
   @unittest.skip("needs GEMM on mixins")
@@ -70,7 +70,7 @@ class TestCall(unittest.TestCase):
     # we define a gemm function
     x = UOp.param(0, dtypes.float, shape=(M, K))
     y = UOp.param(1, dtypes.float, shape=(K, N))
-    c = Tensor.call(a, b, fxn=x@y)
+    c = Tensor.call(a, b, fxn=x@y, inline=True)
 
     np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
 
@@ -86,7 +86,7 @@ class TestCall(unittest.TestCase):
 
     p0, p1 = UOp.param(0, dtypes.float, (10,10)), UOp.param(1, dtypes.float, (10,10))
     complex_fxn = (p0*p1 + p0).exp2() * p1.reciprocal()
-    c = Tensor.call(a, b, fxn=complex_fxn)
+    c = Tensor.call(a, b, fxn=complex_fxn, inline=True)
     c.mean().backward()
 
     np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)

test/unit/test_disk_tensor.py

@@ -301,6 +301,11 @@ class TestDiskTensor(TempDirTestCase):
     # self.assertEqual(dt.tolist(), [10, 2, 20, 4, 30, 6])
     self.assertEqual(dt.tolist(), [10, 20, 30, 4, 5, 6])  # wrong!
 
+  def test_advanced_setitem_not_supported(self):
+    dt = Tensor.arange(12).reshape(3, 4).to(f"disk:{self.tmp('dt_advanced_setitem')}")
+    with self.assertRaises(RuntimeError, msg="advanced setitem is not supported for DISK tensors"):
+      dt[Tensor([0, 2]), Tensor([1, 3])] = 99
+
   def test_assign_const_to_disk(self):
     # assign from CONST (Tensor.full) to disk - source has no buffer, needs contiguous first
     dt = Tensor.empty(4, device=f"disk:{self.tmp('dt_assign_const')}", dtype=dtypes.int32)

tinygrad/codegen/__init__.py

@@ -139,7 +139,6 @@ def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
   return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),), arg=ctx.aux(list(lin.src)) if ctx.has_aux else prg.arg)
 
 def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
-  if ctx.compiler is None: return None
   lib = ctx.compiler.compile_cached(source.arg)
   return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),))
 

tinygrad/device.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Generic, TypeVar, Iterator, Generator
+from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar
 from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, NV_PTX, CUDA_PTX, NV_NAK
 from tinygrad.helpers import EMULATED_DTYPES
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
-from tinygrad.renderer import Renderer
+if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
 # **************** Device ****************
 
@@ -278,37 +278,34 @@ class Compiler:
   def disassemble(self, lib:bytes): pass
 
 @dataclass(frozen=True)
-class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial|None; ctrl_var:ContextVar|None = None # noqa: E702
-
-@dataclass(frozen=True)
-class CompilerSet: cset:list[CompilerPair]; ctrl_var:ContextVar|None = None # noqa: E702
+class CompilerSet: cset:list[tuple[type[Renderer]|functools.partial, ContextVar|None]]; ctrl_var:ContextVar|None = None # noqa: E702
 
 class Compiled:
   profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.
 
   def __init__(self, device:str, allocator:Allocator, compilers:CompilerSet|None, runtime, graph=None, group_id=None):
+    from tinygrad.renderer import Renderer
+
     self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id
 
     self.comps_ctrl_var = compilers.ctrl_var if compilers is not None else None
-    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial|None]]] = {}
-    self.cached_pair:dict[Any, tuple[Renderer, Compiler|None]] = {}
-    for cpair in (compilers.cset if compilers is not None else [CompilerPair(Renderer, Compiler)]):
-      self.comp_sets[self._compiler_name(cpair.renderer, cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
+    self.comp_sets:dict[str, tuple[ContextVar|None, type[Renderer]|functools.partial]] = {}
+    self.cached_pair:dict[Any, Renderer] = {}
+    for ren, var in (compilers.cset if compilers is not None else [(Renderer, None)]):
+      self.comp_sets[var.key.split('_', 1)[-1] if var is not None else self._compiler_name(ren)] = (var, ren)
 
   @property
-  def renderer(self) -> Renderer: return self._select_compiler_pair()[0]
+  def renderer(self) -> Renderer: return self._select_compiler_pair()
 
   @property
   def compiler(self) -> Compiler:
-    if (ret:=self.renderer.compiler or self._select_compiler_pair()[1]) is None: raise RuntimeError(f"no compiler for {self.device}")
+    if (ret:=self.renderer.compiler) is None: raise RuntimeError(f"no compiler for {self.device}")
     return ret
 
-  def _compiler_name(self, r:type[Renderer]|functools.partial, c:type[Compiler]|functools.partial|None) -> str:
-    devname = self.device.split(':')[0].upper()
-    if c is None: return unwrap_class_type(r).__name__.upper().removesuffix("RENDERER").removeprefix(devname) or devname
-    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname) or devname
+  def _compiler_name(self, r:type[Renderer]|functools.partial) -> str:
+    return unwrap_class_type(r).__name__.upper().removesuffix("RENDERER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname
 
-  def _select_compiler_pair(self) -> tuple[Renderer, Compiler|None]:
+  def _select_compiler_pair(self) -> Renderer:
     # select forced compiler from global env var.
     forced_comps = set([self.comp_sets[val][1]] if self.comps_ctrl_var is not None and (val:=self.comps_ctrl_var.value) else [])
 
@@ -397,18 +394,18 @@ def enumerate_devices_str() -> Generator[str, None, None]:
       d = Device[device]
       default_comp_pairs, default_compiler, cc_ctrl_var = d.comp_sets, d.compiler, d.comps_ctrl_var
       try:
-        for k,(en,(r,c)) in default_comp_pairs.items():
-          d.comp_sets = {k:(None,(r,c))} # env var set to None, so it doesn't interfere
+        for k,(en,r) in default_comp_pairs.items():
+          d.comp_sets = {k:(None,r)} # env var set to None, so it doesn't interfere
           d.comps_ctrl_var = None
           try:
             # d.renderer, d.compiler = r(), c()
             with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
             if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
-            set_text = f'({cc_ctrl_var.key}={d._compiler_name(r, c)} to make default)' if cc_ctrl_var is not None else ''
+            set_text = f'({cc_ctrl_var.key}={d._compiler_name(r)} to make default)' if cc_ctrl_var is not None else ''
             default_text = '(default)' if type(default_compiler) is type(d.compiler) else set_text
-            compilers_results.append(f"{colored('+', 'green')} {d._compiler_name(r, c)} {default_text}")
+            compilers_results.append(f"{colored('+', 'green')} {d._compiler_name(r)} {default_text}")
             any_works = True
-          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {d._compiler_name(r, c)}: {e}")
+          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {d._compiler_name(r)}: {e}")
       finally:
         # put the defaults back!
         d.comp_sets, d.comps_ctrl_var = default_comp_pairs, cc_ctrl_var

tinygrad/helpers.py

@@ -185,11 +185,12 @@ ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), Conte
 EMULATE, EMULATED_DTYPES = ContextVar("EMULATE", ""), ContextVar("EMULATED_DTYPES", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 # Compilers
-CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
-NV_PTX, CUDA_PTX, NV_NAK, QCOM_IR3 = ContextVar("NV_PTX", 0), ContextVar("CUDA_PTX", 0), ContextVar("NV_NAK", 0), ContextVar("QCOM_IR3", 0)
+CPU_CC, CPU_LLVM, CPU_LVP = ContextVar("CPU_CC", ""), ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0)
+NV_CC, NV_PTX, NV_NAK = ContextVar("NV_CC", ""), ContextVar("NV_PTX", 0), ContextVar("NV_NAK", 0)
+CUDA_CC, CUDA_PTX, CUDA_NVCC = ContextVar("CUDA_CC", ""), ContextVar("CUDA_PTX", 0), ContextVar("CUDA_NVCC", 0)
 NULL_IR3, NULL_NAK, NULL_ALLOW_COPYOUT = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0), ContextVar("NULL_ALLOW_COPYOUT", 0)
-AMD_CC, CPU_CC, NV_CC, CUDA_CC = ContextVar("AMD_CC", ""), ContextVar("CPU_CC", ""), ContextVar("NV_CC", ""), ContextVar("CUDA_CC", "")
-QCOM_CC = ContextVar("QCOM_CC", "")
+AMD_CC, AMD_LLVM, AMD_HIPCC  = ContextVar("AMD_CC", ""), ContextVar("AMD_LLVM", 0), ContextVar("AMD_HIPCC", 0)
+QCOM_CC, QCOM_IR3 = ContextVar("QCOM_CC", ""), ContextVar("QCOM_IR3", 0)
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
 VIZ = ContextVar("VIZ", 0)
 PROFILE = ContextVar("PROFILE", abs(VIZ.value))

tinygrad/nn/__init__.py

@@ -360,7 +360,8 @@ class Embedding:
 
   def __call__(self, idx:Tensor) -> Tensor:
     if not dtypes.is_int(idx.dtype): raise TypeError(f"Expected integer dtype for index in embedding, got {idx.dtype}")
-    if USE_ATOMICS: return Tensor.call(self.weight, idx, fxn=_embedding_fwd(self.weight.as_param(0), idx.as_param(1)), grad_fxn=_embedding_bwd)
+    if USE_ATOMICS:
+      return Tensor.call(self.weight, idx, fxn=_embedding_fwd(self.weight.as_param(0), idx.as_param(1)), grad_fxn=_embedding_bwd, inline=True)
     return _embedding_fwd(self.weight, idx)
 
 class LSTMCell:

tinygrad/renderer/__init__.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Callable, cast, TYPE_CHECKING
+from typing import Callable, cast
 import functools
 from dataclasses import dataclass, field
 from tinygrad.helpers import to_function_name, dedup, prod, DEBUG
@@ -7,7 +7,7 @@ from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, Gro
 from tinygrad.dtype import AddrSpace, PtrDType
 from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.codegen.opt import Opt
-if TYPE_CHECKING: from tinygrad.device import Compiler
+from tinygrad.device import Compiler
 
 @dataclass(frozen=True)
 class Estimates:
@@ -150,7 +150,8 @@ class Renderer:
   pre_matcher: PatternMatcher|None = None
   extra_matcher: PatternMatcher|None = None
   code_for_op: dict[Ops, Callable] = {}
-  compiler: Compiler|None = None
+
+  compiler: Compiler = Compiler()
 
   def __reduce__(self): return self.__class__, ()
   def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")

tinygrad/renderer/cstyle.py

@@ -340,7 +340,9 @@ class IntelRenderer(OpenCLRenderer):
 class MetalRenderer(CStyleLanguage):
   device = "METAL"
   shared_max = 32768
-  def __init__(self): self.tensor_cores = tc.metal if hasattr(os, 'uname') and os.uname().machine == "arm64" else []
+  def __init__(self):
+    from tinygrad.runtime.ops_metal import MetalCompiler
+    self.compiler, self.tensor_cores = MetalCompiler(), tc.metal if hasattr(os, 'uname') and os.uname().machine == "arm64" else []
 
   # language options
   kernel_typedef = "kernel void"
@@ -382,15 +384,17 @@ class MetalRenderer(CStyleLanguage):
 _nms = list("xyzwabcdefghijkl") + [f'v{i}' for i in range(16, 32)]
 
 class CUDARenderer(CStyleLanguage):
-  device = "CUDA"
   global_max = (2147483647, 65535, 65535)
   local_max = (1024, 1024, 64)
   shared_max = 49152
 
-  def __init__(self, arch:str):
-    self.arch, arch_ver = arch, int(arch[3:])
-    self.tensor_cores = tc.cuda_sm89 if arch_ver >= 89 else tc.cuda_sm80 if arch_ver >= 80 else tc.cuda_sm75 if arch_ver >= 75 else []
-  def __reduce__(self): return self.__class__, (self.arch,)
+  def __init__(self, arch:str, device:str="NV", use_nvcc=False):
+    from tinygrad.runtime.support.compiler_cuda import NVRTCCompiler, NVCCCompiler
+    from tinygrad.runtime.support.hcq import MOCKGPU
+    self.device, self.arch, self.use_nvcc = device, arch, use_nvcc
+    self.compiler = (NVCCCompiler if use_nvcc else NVRTCCompiler)(arch, ptx=bool(MOCKGPU) or device == "CUDA", cache_key=device.lower())
+    self.tensor_cores = tc.cuda_sm89 if (ver:=int(arch[3:])) >= 89 else tc.cuda_sm80 if ver >= 80 else tc.cuda_sm75 if ver >= 75 else []
+  def __reduce__(self): return self.__class__, (self.arch, self.device, self.use_nvcc)
 
   # language options
   # https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html
@@ -547,7 +551,6 @@ class AMDHIPRenderer(CStyleLanguage):
   for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d;\n}""")
     return super().render_kernel(function_name, kernel, bufs, uops, prefix)
 
-class NVRenderer(CUDARenderer): device = "NV"
 class HIPRenderer(AMDHIPRenderer): device = "HIP"
 class AMDHIPCCRenderer(AMDHIPRenderer):
   def __init__(self, arch:str):

tinygrad/renderer/llvmir.py

@@ -205,6 +205,9 @@ class CPULLVMRenderer(LLVMRenderer):
   string_rewrite = base_rewrite + PatternMatcher([(UPat(Ops.WMMA, name="wmma"), render_wmma_amx)])
   def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops)))
   def _render_footer(self, uops: list[UOp]) -> str: return 'attributes #0 = { alwaysinline nounwind "no-builtins" "no-trapping-math"="true" }'
+  def __init__(self):
+    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
+    self.compiler = CPULLVMCompiler()
 
 barrier = 'fence syncscope("workgroup") release\ntail call void @llvm.amdgcn.s.barrier()\nfence syncscope("workgroup") acquire\n'
 code_for_workitem = {"g": lambda x: f"tail call i32 @llvm.amdgcn.workgroup.id.{chr(120+int(x))}()",

tinygrad/renderer/ptx.py

@@ -144,9 +144,11 @@ class PTXRenderer(Renderer):
   tc_sm80 = [x for x in tc.cuda_sm80 if x.dtype_in in [dtypes.half, dtypes.float]]
   code_for_op = asm_for_op
   extra_matcher = ptx_matcher
-  def __init__(self, arch:str, device="CUDA"):
-    self.device, self.arch, arch_ver = device, arch, int(arch[3:])
-    self.tensor_cores = PTXRenderer.tc_sm80 if arch_ver >= 80 else tc.cuda_sm75 if arch_ver >= 75 else []
+  def __init__(self, arch:str, device="NV"):
+    from tinygrad.runtime.support.compiler_cuda import NVPTXCompiler, PTXCompiler
+    from tinygrad.runtime.support.hcq import MOCKGPU
+    self.compiler, self.device, self.arch = (PTXCompiler if bool(MOCKGPU) or device == "CUDA" else NVPTXCompiler)(arch), device, arch
+    self.tensor_cores = PTXRenderer.tc_sm80 if (ver:=int(arch[3:])) >= 80 else tc.cuda_sm75 if ver >= 75 else []
   def __reduce__(self): return self.__class__, (self.arch, self.device)
 
   # language options

tinygrad/runtime/autogen/__init__.py

@@ -115,7 +115,7 @@ def __getattr__(nm):
       return load("rocprof", "['rocprof-trace-decoder', p:='/usr/local/lib/rocprof-trace-decoder.so', p.replace('so','dylib')]",
                   [f"{{}}/include/{s}.h" for s in ["rocprof_trace_decoder", "trace_decoder_instrument", "trace_decoder_types"]],
                   tarball="https://github.com/ROCm/rocprof-trace-decoder/archive/dd0485100971522cc4cd8ae136bdda431061a04d.tar.gz")
-    case "mesa": return load("mesa", "['tinymesa_cpu', 'tinymesa']", [
+    case "mesa": return load("mesa", "([] if CPU_CC.value == 'LVP' or bool(CPU_LVP) else ['tinymesa']) + ['tinymesa_cpu']", [
         *[f"{{}}/src/compiler/nir/{s}.h" for s in ["nir", "nir_builder", "nir_shader_compiler_options", "nir_serialize"]], "{}/gen/nir_intrinsics.h",
         *[f"{{}}/src/nouveau/{s}.h" for s in ["headers/nv_device_info", "compiler/nak"]],
         *[f"{{}}/src/gallium/auxiliary/gallivm/lp_bld{s}.h" for s in ["", "_passmgr", "_misc", "_type", "_init", "_nir", "_struct", "_jit_types",
@@ -134,7 +134,8 @@ def __getattr__(nm):
           *[f"python3 src/compiler/{s}_h.py > gen/{s.split('/')[-1]}.h" for s in ["nir/nir_opcodes", "nir/nir_builder_opcodes"]],
           *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
   tarball="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
-  prolog=["import gzip, base64"], epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
+  prolog=["from tinygrad.helpers import CPU_CC, CPU_LVP", "import gzip, base64"],
+  epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
     case "libclang":
       return load("libclang", clang_lib,
                   lambda: [f"{system('llvm-config-20 --includedir')}/clang-c/{s}.h" for s in ["Index", "CXString", "CXSourceLocation", "CXFile"]],

tinygrad/runtime/autogen/mesa.py

@@ -4,8 +4,9 @@ import ctypes
 from typing import Annotated, Literal, TypeAlias
 from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
 from tinygrad.runtime.support import c
+from tinygrad.helpers import CPU_CC, CPU_LVP
 import gzip, base64
-dll = c.DLL('mesa', ['tinymesa_cpu', 'tinymesa'])
+dll = c.DLL('mesa', ([] if CPU_CC.value == 'LVP' or bool(CPU_LVP) else ['tinymesa']) + ['tinymesa_cpu'])
 class struct_u_printf_info(ctypes.Structure): pass
 u_printf_info: TypeAlias = struct_u_printf_info
 uint32_t: TypeAlias = Annotated[int, ctypes.c_uint32]

tinygrad/runtime/ops_amd.py

@@ -6,9 +6,9 @@ from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface
 from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices
 from tinygrad.uop.ops import sint
-from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar
-from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, ceildiv, unwrap
+from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, AMD_HIPCC, ceildiv, unwrap
 from tinygrad.renderer.cstyle import AMDHIPRenderer, AMDHIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt, amdgpu_kd, amdgpu_drm
@@ -962,9 +962,9 @@ class AMDDevice(HCQCompiled):
     self.sdma_queues:dict = {}
     self.has_sdma_queue = self.sdma_queue(0) is not None
 
-    compilers = CompilerSet([CompilerPair(functools.partial(AMDHIPRenderer, self.arch), None),
-                             CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), None, AMD_LLVM),
-                             CompilerPair(functools.partial(AMDHIPCCRenderer, self.arch), None)], ctrl_var=AMD_CC)
+    compilers = CompilerSet([(functools.partial(AMDHIPRenderer, self.arch), None),
+                             (functools.partial(AMDLLVMRenderer, self.arch), AMD_LLVM),
+                             (functools.partial(AMDHIPCCRenderer, self.arch), AMD_HIPCC)], ctrl_var=AMD_CC)
 
     super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                      functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),

tinygrad/runtime/ops_cl.py

@@ -5,7 +5,7 @@ from tinygrad.runtime.autogen import opencl as cl
 from tinygrad.runtime.support import c
 from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, mv_address, suppress_finalizing
 from tinygrad.renderer.cstyle import OpenCLRenderer, IntelRenderer
-from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError, CompilerPair, CompilerSet
+from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError, CompilerSet
 from tinygrad.dtype import ImageDType
 
 CC_CB = c.CFUNCTYPE[None, [c.POINTER[ctypes.c_char], c.POINTER[None], cl.size_t, c.POINTER[None]]]
@@ -39,9 +39,9 @@ class CLCompiler(Compiler):
 
 class CLProgram:
   def __init__(self, device:CLDevice, name:str, lib:bytes, buf_dtypes=[], **kwargs):
-    self.dev, self.name, self.lib, self.buf_dtypes = device, name, lib, buf_dtypes
-    self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.device_id, (ctypes.c_size_t * 1)(len(lib)),
-                                                        to_char_p_p([lib], ctypes.c_ubyte), binary_status := ctypes.c_int32(),
+    self.dev, self.name, self.lib, self.buf_dtypes = device, name, device.cl_compiler.compile_cached(lib.decode()), buf_dtypes
+    self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.device_id, (ctypes.c_size_t * 1)(len(self.lib)),
+                                                        to_char_p_p([self.lib], ctypes.c_ubyte), binary_status := ctypes.c_int32(),
                                                         errcode_ret := ctypes.c_int32()), errcode_ret)
     check(binary_status.value)
     check(cl.clBuildProgram(self.program, 1, device.device_id, None, BP_CB(), None)) # NOTE: OSX requires this
@@ -125,8 +125,9 @@ class CLDevice(Compiled):
                                            ctypes.string_at(buf, size=total.value).decode())[1]
 
     renderer = IntelRenderer if "cl_intel_subgroup_matrix_multiply_accumulate" in self.device_exts else OpenCLRenderer
-    compiler = functools.partial(CLCompiler, self, f"{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}")
-    super().__init__(device, CLAllocator(self), CompilerSet([CompilerPair(renderer, compiler)]), functools.partial(CLProgram, self))
+    self.cl_compiler = CLCompiler(self, f"{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}")
+    super().__init__(device, CLAllocator(self), CompilerSet([(renderer, None)]), functools.partial(CLProgram, self))
+
   def synchronize(self):
     check(cl.clFinish(self.queue))
     self.pending_copyin.clear()

tinygrad/runtime/ops_cpu.py

@@ -2,13 +2,12 @@ from __future__ import annotations
 import platform, sys, ctypes, functools, time, mmap, threading, queue
 from tinygrad.helpers import to_mv, OSX, WIN, mv_address, wait_cond, suppress_finalizing, unwrap, data64_le
 from tinygrad.helpers import CPU_CC, CPU_LVP, CPU_LLVM
-from tinygrad.device import BufferSpec, DMACPURef, CompilerSet, CompilerPair
+from tinygrad.device import BufferSpec, DMACPURef, CompilerSet
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.hcq import CLikeArgsState
 from tinygrad.renderer.cstyle import ClangJITRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
-from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.uop.ops import sint
 
@@ -134,6 +133,5 @@ class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
     self.tasks:queue.Queue = queue.Queue()
     CPUWorker(self, self.tasks, thread_id=0).start()
-    compilers = CompilerSet([CompilerPair(ClangJITRenderer, None), CompilerPair(CPULLVMRenderer, CPULLVMCompiler, ctrl_var=CPU_LLVM),
-                             CompilerPair(LVPRenderer, None, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC)
+    compilers = CompilerSet([(ClangJITRenderer, None), (CPULLVMRenderer, CPU_LLVM), (LVPRenderer, CPU_LVP)], ctrl_var=CPU_CC)
     super().__init__(device, CPUAllocator(self), compilers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue)

tinygrad/runtime/ops_cuda.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 import ctypes, functools
-from tinygrad.helpers import DEBUG, getenv, mv_address, suppress_finalizing, CUDA_CC, CUDA_PTX
-from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerPair, CompilerSet
+from tinygrad.helpers import DEBUG, getenv, mv_address, suppress_finalizing, CUDA_CC, CUDA_PTX, CUDA_NVCC
+from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerSet
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.runtime.autogen import cuda
-from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler, NVCCCompiler
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx
 from tinygrad.runtime.support.c import init_c_struct_t, init_c_var
 if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl  # noqa: F401  # pylint: disable=unused-import
 if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported
@@ -118,9 +118,9 @@ class CUDADevice(Compiled):
     CUDADevice.devices.append(self)
 
     from tinygrad.runtime.graph.cuda import CUDAGraph
-    compilers = CompilerSet([CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(CUDACompiler, self.arch)),
-                             CompilerPair(functools.partial(PTXRenderer, self.arch), functools.partial(PTXCompiler, self.arch), CUDA_PTX),
-                             CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(NVCCCompiler, self.arch))], ctrl_var=CUDA_CC)
+    compilers = CompilerSet([(functools.partial(CUDARenderer, self.arch, device="CUDA"), None),
+                             (functools.partial(PTXRenderer, self.arch, device="CUDA"), CUDA_PTX),
+                             (functools.partial(CUDARenderer, self.arch, device="CUDA", use_nvcc=True), CUDA_NVCC)], ctrl_var=CUDA_CC)
     super().__init__(device, CUDAAllocator(self), compilers, functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph)
 
   def synchronize(self):

tinygrad/runtime/ops_dsp.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys, subprocess, struct
 assert sys.platform != 'win32'
-from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, CompilerSet, CompilerPair
+from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, CompilerSet
 from tinygrad.dtype import dtypes, DType, PtrDType
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing
@@ -45,6 +45,8 @@ class DSPRenderer(ClangRenderer):
   type_map = { **ClangRenderer.type_map, dtypes.uint64: "unsigned long long", dtypes.int64: "long long" }
   code_for_op = {k:v for k,v in ClangRenderer.code_for_op.items() if k != Ops.SQRT}
 
+  def __init__(self): self.compiler = DSPCompiler()
+
   def _render_defines(self, uops) -> list[str]:
     return ['''/* DSP boilerplate */ struct dcvs_v2_req { int type; int _pad; _Bool dcvs_enable; char dcvs_option; _Bool set_latency; int latency;
       _Bool set_dcvs_params; short _pad2; char target_corner; char min_corner; char max_corner; int _pad3[3];};''','int HAP_power_set(void*, void*);',
@@ -116,28 +118,11 @@ class DSPAllocator(Allocator['DSPDevice']):
   def _copyout(self, dest:memoryview, src:DSPBuffer): ctypes.memmove(mv_address(dest), src.va_addr, dest.nbytes)
   def _offset(self, buf, size:int, offset:int): return DSPBuffer(buf.va_addr+offset, size, buf.share_info, buf.offset+offset)
 
-class ClangCompiler(Compiler):
-  def __init__(self, cachekey="compile_clang", args:list[str]|None=None, objdump_tool='objdump'):
-    self.args = ['-shared', '-march=native'] if args is None else args
-    self.objdump_tool = objdump_tool
-    super().__init__(cachekey)
-
-  def compile(self, src:str) -> bytes:
-    # TODO: remove file write. sadly clang doesn't like the use of /dev/stdout here
-    with tempfile.NamedTemporaryFile(delete=True) as f:
-      system(f"{getenv('CC','clang')} {' '.join(self.args)} -O2 -Wall -Werror -x c -fPIC -ffreestanding -nostdlib - -o {f.name}", input=src.encode())
-      return pathlib.Path(f.name).read_bytes()
-
-  def disassemble(self, lib:bytes): return cpu_objdump(lib, self.objdump_tool)
-
-class DSPDevice(Compiled):
-  def __init__(self, device:str=""):
-    compiler_args = ["--target=hexagon", "-mcpu=hexagonv65", "-fuse-ld=lld", "-nostdlib",  "-mhvx=v65", "-mhvx-length=128b"]
-    if getenv("MOCKDSP"):
-      mock_compilers = CompilerSet([CompilerPair(MockDSPRenderer, functools.partial(ClangCompiler, None, ["-static"]+compiler_args, 'llvm-objdump'))])
-      super().__init__(device, DSPAllocator(self), mock_compilers, MockDSPProgram)
+class DSPCompiler(Compiler):
+  def __init__(self, mock:bool=False):
+    compiler_args = "--target=hexagon -mcpu=hexagonv65 -fuse-ld=lld -nostdlib -mhvx=v65 -mhvx-length=128b"
+    if mock: self.args = f"-static {compiler_args}"
     else:
-      self.ion_fd = os.open('/dev/ion', os.O_RDONLY)
       # Generate link script to pass into clang. Aligning all used sections to 4k fixes invoke problem.
       sections = ['text', 'rela.plt', 'rela.dyn', 'plt', 'data', 'bss', 'hash', 'dynamic',
                   'got', 'got.plt', 'dynsym', 'dynstr', 'symtab', 'shstrtab', 'strtab']
@@ -146,8 +131,25 @@ class DSPDevice(Compiled):
         self.link_ld.write(f"SECTIONS {{ . = 0x0; {sections_link}\n /DISCARD/ : {{ *(.note .note.* .gnu.hash .comment) }} }}".encode())
         self.link_ld.flush()
 
-      compiler = functools.partial(ClangCompiler, "compile_dsp", ["-shared"] + compiler_args + [f"-T{self.link_ld.name}"], 'llvm-objdump')
-      super().__init__(device, DSPAllocator(self), CompilerSet([CompilerPair(DSPRenderer, compiler)]), functools.partial(DSPProgram, self))
+      self.args = f"-shared {compiler_args} -T{self.link_ld.name}"
+
+    super().__init__(None if mock else "compile_dsp")
+
+  def compile(self, src:str) -> bytes:
+    # TODO: remove file write. sadly clang doesn't like the use of /dev/stdout here
+    with tempfile.NamedTemporaryFile(delete=True) as f:
+      system(f"{getenv('CC','clang')} {self.args} -O2 -Wall -Werror -x c -fPIC -ffreestanding -nostdlib - -o {f.name}", input=src.encode())
+      return pathlib.Path(f.name).read_bytes()
+
+  def disassemble(self, lib:bytes): return cpu_objdump(lib, "llvm-objdump")
+
+
+class DSPDevice(Compiled):
+  def __init__(self, device:str=""):
+    if getenv("MOCKDSP"): super().__init__(device, DSPAllocator(self), CompilerSet([(MockDSPRenderer, None)]), MockDSPProgram)
+    else:
+      self.ion_fd = os.open('/dev/ion', os.O_RDONLY)
+      super().__init__(device, DSPAllocator(self), CompilerSet([(DSPRenderer, None)]), functools.partial(DSPProgram, self))
       fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes()))
       self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True))
       ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)
@@ -268,6 +270,7 @@ static void *mmap2(void *addr, unsigned int length, int prot, int flags, int fd,
 return (void*)syscall((long)addr, length, prot, flags, fd, offset, 222); }}'''
 
 class MockDSPRenderer(DSPRenderer):
+  def __init__(self): self.compiler = DSPCompiler(mock=True)
   def _render_defines(self, uops) -> list[str]: return ClangRenderer._render_defines(self, uops)
   def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[DType,bool]]]) -> str:
     # https://gpages.juszkiewicz.com.pl/syscalls-table/syscalls.html

tinygrad/runtime/ops_hip.py

@@ -1,6 +1,6 @@
 import ctypes, functools
 from tinygrad.helpers import mv_address, getenv, suppress_finalizing
-from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet
 from tinygrad.runtime.autogen import hip
 from tinygrad.renderer.cstyle import HIPRenderer
 from tinygrad.runtime.support.c import init_c_var, init_c_struct_t
@@ -15,7 +15,7 @@ class HIPDevice(Compiled):
     self.arch = init_c_var(hip.hipDeviceProp_t, lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode()
     self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t, lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
 
-    compilers = CompilerSet([CompilerPair(functools.partial(HIPRenderer, self.arch), None)])
+    compilers = CompilerSet([(functools.partial(HIPRenderer, self.arch), None)])
     super().__init__(device, HIPAllocator(self), compilers, functools.partial(HIPProgram, self))
   def synchronize(self):
     check(hip.hipSetDevice(self.device_id))

tinygrad/runtime/ops_metal.py

@@ -1,7 +1,7 @@
-import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
-from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, PROFILE, ProfileRangeEvent, cpu_profile, unwrap, suppress_finalizing
+import subprocess, pathlib, struct, ctypes, tempfile, functools, decimal, platform
+from tinygrad.helpers import prod, to_mv, round_up, cache_dir, PROFILE, ProfileRangeEvent, cpu_profile, unwrap, suppress_finalizing
 import tinygrad.runtime.support.objc as objc
-from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent, CompilerSet
 from tinygrad.renderer.cstyle import MetalRenderer
 from tinygrad.runtime.autogen import metal
 from tinygrad.runtime.support.c import DLL
@@ -42,7 +42,7 @@ 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), CompilerSet([CompilerPair(MetalRenderer, MetalCompiler), CompilerPair(MetalRenderer, Compiler)]),
+    super().__init__(device, MetalAllocator(self), CompilerSet([(MetalRenderer, None)]),
       functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None)
 
   def synchronize(self):
@@ -54,20 +54,12 @@ class MetalDevice(Compiled):
         Compiled.profile_events += [ProfileRangeEvent(self.device, lb, st, en)]
     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()))
-  error_check(compileError, CompileError)
-  return library
-
 class MetalCompiler(Compiler):
   # Opening METAL after LLVM doesn't fail because ctypes.CDLL opens with RTLD_LOCAL but MTLCompiler opens it's own llvm with RTLD_GLOBAL
   # This means that MTLCompiler's llvm will create it's own instances of global state because RTLD_LOCAL doesn't export symbols, but if RTLD_GLOBAL
   # library is loaded first then RTLD_LOCAL library will just use it's symbols. On linux there is RTLD_DEEPBIND to prevent that, but on macos there
   # doesn't seem to be anything we can do.
-  with contextlib.suppress(FileNotFoundError, ModuleNotFoundError):
-    import tinygrad.runtime.autogen.llvm # noqa: F401
+  import tinygrad.runtime.autogen.llvm as _
   support = DLL("MTLCompiler", "MTLCompiler")
   support.MTLCodeGenServiceCreate.restype = ctypes.c_void_p
 
@@ -118,15 +110,9 @@ class MetalCompiler(Compiler):
 class MetalProgram:
   def __init__(self, dev:MetalDevice, name:str, lib:bytes, **kwargs):
     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()
-      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
+    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()
+    error_check(error_lib)
     self.fxn = self.library.newFunctionWithName(to_ns_str(name)).retained()
     descriptor = metal.MTLComputePipelineDescriptor.new()
     descriptor.setComputeFunction(self.fxn)

tinygrad/runtime/ops_null.py

@@ -1,5 +1,5 @@
 import functools
-from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, Allocator, CompilerSet
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.renderer.cstyle import Renderer, CStyleLanguage, AMDHIPRenderer
 from tinygrad.uop.ops import Ops
@@ -39,6 +39,6 @@ class NullDevice(Compiled):
       case "AMD_CDNA4": renderer = functools.partial(AMDHIPRenderer, "gfx950")
       case "": renderer = NullRenderer
       case _: raise RuntimeError(f"can't EMULATE device: {EMULATE.value}")
-    compilers = CompilerSet([CompilerPair(renderer, Compiler), CompilerPair(functools.partial(IR3Renderer, 0x6030001), None, NULL_IR3), # adreno 630
-                             CompilerPair(functools.partial(NAKRenderer, "sm_120", 48), None, NULL_NAK)]) # 5090
+    compilers = CompilerSet([(renderer, None), (functools.partial(IR3Renderer, 0x6030001), NULL_IR3), # adreno 630
+                             (functools.partial(NAKRenderer, "sm_120", 48), NULL_NAK)]) # 5090
     super().__init__(device, NullAllocator(self), compilers, functools.partial(NullProgram, device), NullGraph)

tinygrad/runtime/ops_nv.py

@@ -6,12 +6,11 @@ from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQProgram, HCQSignal, BumpAllocator
 from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU, hcq_filter_visible_devices, hcq_profile
 from tinygrad.uop.ops import sint
-from tinygrad.device import Compiled, BufferSpec, CompilerPair, CompilerSet
+from tinygrad.device import Compiled, BufferSpec, CompilerSet
 from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, NV_CC, NV_PTX, NV_NAK, PROFILE
 from tinygrad.helpers import ContextVar, VIZ, ProfileEvent
 from tinygrad.renderer.ptx import PTXRenderer
-from tinygrad.renderer.cstyle import NVRenderer
-from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler
+from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.runtime.autogen import nv_570, nv_580, pci, mesa
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.nv.nvdev import NVDev, NVMemoryManager
@@ -620,10 +619,9 @@ class NVDevice(HCQCompiled[NVSignal]):
     self.arch: str = "sm_120" if self.sm_version==0xa04 else f"sm_{(self.sm_version>>8)&0xff}{(val>>4) if (val:=self.sm_version&0xff) > 0xf else val}"
     self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf)
 
-    cucc, ptxcc = (CUDACompiler, PTXCompiler) if MOCKGPU else (NVCompiler, NVPTXCompiler)
-    compilers = CompilerSet(ctrl_var=NV_CC, cset=[CompilerPair(functools.partial(NVRenderer, self.arch),functools.partial(cucc, self.arch)),
-       CompilerPair(functools.partial(PTXRenderer, self.arch, device="NV"), functools.partial(ptxcc, self.arch), NV_PTX),
-       CompilerPair(functools.partial(NAKRenderer, self.arch, self.max_warps_per_sm), None, NV_NAK)])
+    compilers = CompilerSet(ctrl_var=NV_CC, cset=[(functools.partial(CUDARenderer, self.arch), None),
+       (functools.partial(PTXRenderer, self.arch, device="NV"), NV_PTX),
+       (functools.partial(NAKRenderer, self.arch, self.max_warps_per_sm), NV_NAK)])
     super().__init__(device, NVAllocator(self), compilers, functools.partial(NVProgram, self), NVSignal, NVComputeQueue, NVCopyQueue)
 
     self.pma_enabled = PMA.value > 0 and PROFILE >= 1

tinygrad/runtime/ops_python.py

@@ -6,7 +6,7 @@ from typing import Any, TYPE_CHECKING
 import pickle, base64, itertools, time, struct, sys, functools
 from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate, float_to_fp16, float_to_bf16, float_to_fp8, fp8_to_float
 from tinygrad.helpers import all_same, getenv, flatten, get_single_element, EMULATE
-from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer
@@ -214,9 +214,14 @@ class PythonProgram:
         i += 1
     return time.perf_counter() - st
 
+class PythonCompiler(Compiler):
+  def compile(self, src:str) -> bytes: return base64.b64decode(src)
+
 class PythonRenderer(Renderer):
   device = "PYTHON"
   code_for_op = python_alu
+  compiler = PythonCompiler()
+
   def __init__(self):
     match EMULATE.value:
       case "METAL": self.device, self.tensor_cores = "METAL", tc.metal
@@ -236,9 +241,6 @@ class PythonRenderer(Renderer):
     lops = [(u.op, u.dtype, [uops.index(v) for v in u.src if u.op is not Ops.SPECIAL], u.arg) for u in uops]
     return base64.b64encode(pickle.dumps(lops)).decode()
 
-class PythonCompiler(Compiler):
-  def compile(self, src:str) -> bytes: return base64.b64decode(src)
-
 class PythonAllocator(Allocator['PythonDevice']):
   def _alloc(self, size, options): return memoryview(bytearray(size))
   def _copyin(self, dest, src:memoryview): dest[:] = src
@@ -246,4 +248,4 @@ class PythonAllocator(Allocator['PythonDevice']):
 
 class PythonDevice(Compiled):
   def __init__(self, device:str):
-    super().__init__(device, PythonAllocator(self), CompilerSet([CompilerPair(PythonRenderer, PythonCompiler)]), PythonProgram)
+    super().__init__(device, PythonAllocator(self), CompilerSet([(PythonRenderer, None)]), PythonProgram)

tinygrad/runtime/ops_qcom.py

@@ -2,15 +2,15 @@ from __future__ import annotations
 import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib
 assert sys.platform != 'win32'
 from typing import Any
-from tinygrad.device import BufferSpec, CompilerSet, CompilerPair, Device
+from tinygrad.device import BufferSpec, CompilerSet, Device
 from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, mesa
-from tinygrad.runtime.ops_cl import CLCompiler, CLDevice
+from tinygrad.runtime.ops_cl import CLDevice
 from tinygrad.renderer.cstyle import QCOMRenderer
 from tinygrad.renderer.nir import IR3Renderer
 from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, fromimport, cpu_profile, lo32, suppress_finalizing
-from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC, PROFILE
+from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC, PROFILE, DEBUG
 from tinygrad.dtype import ImageDType, dtypes
 from tinygrad.runtime.support.system import System
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import
@@ -49,10 +49,6 @@ def pkt7_hdr(opcode: int, cnt: int): return mesa.CP_TYPE7_PKT | cnt & 0x3FFF | p
 def pkt4_hdr(reg: int, cnt: int): return mesa.CP_TYPE4_PKT | cnt & 0x7F | parity(cnt) << 7 | (reg & 0x3FFFF) << 8 | parity(reg) << 27
 
 def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
-class QCOMCompiler(CLCompiler):
-  def __init__(self, device:str=""): super().__init__(CLDevice(device), 'compile_qcom')
-  def disassemble(self, lib:bytes):
-    fromimport('tinygrad.runtime.support.compiler_mesa', 'disas_adreno')(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)])
 
 class QCOMSignal(HCQSignal):
   def __init__(self, *args, **kwargs): super().__init__(*args, **{**kwargs, 'timestamp_divider': 19.2})
@@ -231,7 +227,7 @@ class QCOMArgsState(HCQArgsState):
 class QCOMProgram(HCQProgram):
   def __init__(self, dev: QCOMDevice, name: str, lib: bytes, buf_dtypes=[], **kwargs):
     self.dev: QCOMDevice = dev
-    self.buf_dtypes, self.name, self.lib, self.NIR = buf_dtypes, name, lib, isinstance(dev.renderer, IR3Renderer)
+    self.buf_dtypes, self.name, self.NIR = buf_dtypes, name, isinstance(dev.renderer, IR3Renderer)
 
     if self.NIR:
       from tinygrad.runtime.support.compiler_mesa import IR3Compiler
@@ -252,7 +248,9 @@ class QCOMProgram(HCQProgram):
       self.tex_off, self.ibo_off, self.samp_off = 2048, 2048 + 0x40 * self.tex_cnt, 2048 + 0x40 * (self.tex_cnt + self.ibo_cnt)
       self.fregs, self.hregs = v.info.max_reg + 1, v.info.max_half_reg + 1
       self.consts_info:list[tuple] = []
-    else: self._parse_lib()
+    else:
+      self._parse_lib(lib:=self.dev.cl_dev.cl_compiler.compile_cached(lib.decode()))
+      if DEBUG >= 7: fromimport('tinygrad.runtime.support.compiler_mesa', 'disas_adreno')(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)])
 
     self.lib_gpu: HCQBuffer = self.dev.allocator.alloc(self.image_size, buf_spec:=BufferSpec(cpu_access=True, nolru=True))
     to_mv(self.lib_gpu.va_addr, self.image_size)[:] = self.image
@@ -274,21 +272,21 @@ class QCOMProgram(HCQProgram):
       raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
     return super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait)
 
-  def _parse_lib(self):
+  def _parse_lib(self, lib):
     # Extract image binary
-    self.image_size = _read_lib(self.lib, 0x100)
-    self.image = bytearray(self.lib[(image_offset:=_read_lib(self.lib, 0xc0)):image_offset+self.image_size])
+    self.image_size = _read_lib(lib, 0x100)
+    self.image = bytearray(lib[(image_offset:=_read_lib(lib, 0xc0)):image_offset+self.image_size])
 
     # Parse image descriptors
-    image_desc_off = _read_lib(self.lib, 0x110)
-    self.prg_offset, self.brnchstck = _read_lib(self.lib, image_desc_off+0xc4), _read_lib(self.lib, image_desc_off+0x108) // 2
-    self.pvtmem, self.shmem = _read_lib(self.lib, image_desc_off+0xc8), _read_lib(self.lib, image_desc_off+0xd8)
+    image_desc_off = _read_lib(lib, 0x110)
+    self.prg_offset, self.brnchstck = _read_lib(lib, image_desc_off+0xc4), _read_lib(lib, image_desc_off+0x108) // 2
+    self.pvtmem, self.shmem = _read_lib(lib, image_desc_off+0xc8), _read_lib(lib, image_desc_off+0xd8)
 
     # Fill up constants and buffers info
     self.consts_info = []
 
     # Collect sampler info.
-    self.samp_cnt = samp_cnt_in_file = _read_lib(self.lib, image_desc_off + 0xdc)
+    self.samp_cnt = samp_cnt_in_file = _read_lib(lib, image_desc_off + 0xdc)
     assert self.samp_cnt <= 1, "Up to one sampler supported"
     if self.samp_cnt:
       self.samp_cnt += 1
@@ -298,8 +296,8 @@ class QCOMProgram(HCQProgram):
 
     # Collect kernel arguments (buffers) info.
     bdoff, binfos = round_up(image_desc_off + 0x158 + len(self.name), 4) + 8 * samp_cnt_in_file, []
-    while bdoff + 32 <= len(self.lib):
-      length, _, _, offset_words, _, _, _, typ = struct.unpack("8I", self.lib[bdoff:bdoff+32])
+    while bdoff + 32 <= len(lib):
+      length, _, _, offset_words, _, _, _, typ = struct.unpack("8I", lib[bdoff:bdoff+32])
       if length == 0: break
       binfos.append((offset_words * 4, typ))
       bdoff += length
@@ -309,16 +307,16 @@ class QCOMProgram(HCQProgram):
     self.tex_cnt, self.ibo_cnt = sum(typ is BUFTYPE_TEX for _,typ in binfos), sum(typ is BUFTYPE_IBO for _,typ in binfos)
     self.ibo_off, self.tex_off, self.samp_off = 2048, 2048 + 0x40 * self.ibo_cnt, 2048 + 0x40 * self.tex_cnt + 0x40 * self.ibo_cnt
 
-    if _read_lib(self.lib, 0xb0) != 0: # check if we have constants.
-      cdoff = _read_lib(self.lib, 0xac)
+    if _read_lib(lib, 0xb0) != 0: # check if we have constants.
+      cdoff = _read_lib(lib, 0xac)
       while cdoff + 40 <= image_offset:
-        cnst, offset_words, _, is32 = struct.unpack("I", self.lib[cdoff:cdoff+4])[0], *struct.unpack("III", self.lib[cdoff+16:cdoff+28])
+        cnst, offset_words, _, is32 = struct.unpack("I", lib[cdoff:cdoff+4])[0], *struct.unpack("III", lib[cdoff+16:cdoff+28])
         self.consts_info.append((cnst, offset_words * (sz_bytes:=(2 << is32)), sz_bytes))
         cdoff += 40
 
     # Registers info
-    reg_desc_off = _read_lib(self.lib, 0x34)
-    self.fregs, self.hregs = _read_lib(self.lib, reg_desc_off + 0x14), _read_lib(self.lib, reg_desc_off + 0x18)
+    reg_desc_off = _read_lib(lib, 0x34)
+    self.fregs, self.hregs = _read_lib(lib, reg_desc_off + 0x14), _read_lib(lib, reg_desc_off + 0x18)
 
 class QCOMTextureInfo:
   def __init__(self, pitch:int, real_stride:int, desc:list[int], ibo:list[int]):
@@ -385,8 +383,8 @@ class QCOMDevice(HCQCompiled):
     if PROFILE and self.gpu_id[:2] < (7, 3):
       System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", value="4000000000", msg="Failed to disable suspend mode", expected="4294967276")
 
-    compilers = CompilerSet(ctrl_var=QCOM_CC, cset=[CompilerPair(QCOMRenderer, functools.partial(QCOMCompiler, device)),
-                                                    CompilerPair(functools.partial(IR3Renderer, info.chip_id), None, QCOM_IR3)])
+    self.cl_dev = CLDevice(device)
+    compilers = CompilerSet(ctrl_var=QCOM_CC, cset=[(QCOMRenderer, None), (functools.partial(IR3Renderer, info.chip_id), QCOM_IR3)])
     super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
                      functools.partial(QCOMComputeQueue, self), None)
 

tinygrad/runtime/ops_webgpu.py

@@ -1,5 +1,5 @@
 import functools, struct
-from tinygrad.device import  Compiled, Allocator, Compiler, BufferSpec, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, Allocator, BufferSpec, CompilerSet
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.helpers import round_up, suppress_finalizing
 from tinygrad.runtime.autogen import webgpu
@@ -217,7 +217,7 @@ class WebGpuDevice(Compiled):
     self.device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback,
       webgpu.WGPURequestDeviceStatus, 1, 2, adapter_res, dev_desc)
 
-    super().__init__(device, WebGpuAllocator(self), CompilerSet([CompilerPair(WGSLRenderer, Compiler)]),
+    super().__init__(device, WebGpuAllocator(self), CompilerSet([(WGSLRenderer, None)]),
       functools.partial(WebGPUProgram, (self.device_res, webgpu.WGPUFeatureName_TimestampQuery in supported)))
 
   def synchronize(self):

tinygrad/runtime/support/compiler_cuda.py

@@ -1,5 +1,4 @@
-import subprocess, hashlib, tempfile, ctypes, re, pathlib
-from typing import Callable
+import hashlib, tempfile, ctypes, re, pathlib
 from tinygrad.helpers import to_char_p_p, colored, getenv, system
 from tinygrad.runtime.support.c import init_c_var
 from tinygrad.runtime.autogen import nvrtc, nvjitlink as jitlink
@@ -42,36 +41,32 @@ def cuda_disassemble(lib:bytes, arch:str, ptx=False):
     print(system(f'nvdisasm {fn}'))
   except Exception as e: print("Failed to generate SASS", str(e), "Make sure your PATH contains ptxas/nvdisasm binary of compatible version.")
 
-class CUDACompiler(Compiler):
-  def __init__(self, arch:str, cache_key:str="cuda"):
-    self.arch, self.compile_options = arch, [f'--gpu-architecture={arch}']
+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")
     super().__init__(f"compile_{cache_key}_{self.arch}")
-  def _compile_program(self, src:str, nvrtc_get_content:Callable, nvrtc_get_size:Callable) -> bytes:
+  def compile(self, src:str) -> bytes:
     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_get_content, nvrtc_get_size, nvrtc_check)
+    data = _get_bytes(prog, nvrtc.nvrtcGetPTX if self.ptx else nvrtc.nvrtcGetCUBIN,
+                      nvrtc.nvrtcGetPTXSize if self.ptx else nvrtc.nvrtcGetCUBINSize, nvrtc_check)
     nvrtc_check(nvrtc.nvrtcDestroyProgram(ctypes.byref(prog)))
     return data
-  def compile(self, src:str) -> bytes: return self._compile_program(src, nvrtc.nvrtcGetPTX, nvrtc.nvrtcGetPTXSize)
-  def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch, ptx=True)
-
-class NVCompiler(CUDACompiler):
-  def __init__(self, arch:str): super().__init__(arch, cache_key="nv")
-  def compile(self, src:str) -> bytes: return self._compile_program(src, nvrtc.nvrtcGetCUBIN, nvrtc.nvrtcGetCUBINSize)
-  def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch)
+  def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch, ptx=self.ptx)
 
 class NVCCCompiler(Compiler):
-  def __init__(self, arch:str, extra_options:list[str]=[]):
+  def __init__(self, arch:str, ptx:bool=True, cache_key:str="cuda", extra_options:list[str]=[]):
+    assert ptx, "NVCCCompiler cubin support unimplemented"
     self.arch, self.extra_options = arch, extra_options
-    super().__init__(f"compile_nvcc_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}")
+    super().__init__(f"compile_nvcc_{cache_key}_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}")
   def compile(self, src:str) -> bytes:
     with tempfile.NamedTemporaryFile(suffix=".cu") as srcf, tempfile.NamedTemporaryFile(suffix=".ptx") as libf:
       srcf.write(src.encode())
       srcf.flush()
-      subprocess.run(["nvcc", f"-arch={self.arch}", "-ptx", "-o", libf.name, srcf.name] + self.extra_options, check=True)
+      system(f"nvcc -arch={self.arch} -ptx -o {libf.name} {srcf.name}" + ' '.join(self.extra_options))
       return libf.read()
   def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch, ptx=True)
 

tinygrad/schedule/rangeify.py

@@ -33,6 +33,10 @@ pm_mops = PatternMatcher([
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
+def collapse_nested_assign(assign:UOp, target:UOp, src:UOp):
+  """nested ASSIGN to the same buffer (e.g. __iadd__ in __setitem__): collapse the redundant outer ASSIGN"""
+  if src.src[0].base is target.base: return src if src.src[0] is target else assign.replace(src=(target, src.src[1]))
+
 def assign_to_contiguous(assign:UOp, target:UOp, src:UOp):
   if (t := target.base).op is Ops.BUFFER or (t.op is Ops.MSTACK and all(s.op is Ops.BUFFER for s in t.src)): return None
   return src.f(Ops.CONTIGUOUS, tag=assign.tag)
@@ -75,9 +79,8 @@ mop_cleanup = PatternMatcher([
 ])
 
 def resolve_call(c:UOp) -> UOp|None:
-  # don't resolve real kernel calls, sink or program
-  if c.src[0].op is Ops.SINK and isinstance(c.src[0].arg, KernelInfo): return None
-  if c.src[0].op is Ops.PROGRAM: return None
+  # we only resolve here if the call is inlined
+  if not c.arg.inline: return None
   params = sorted([x for x in c.src[0].toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
   args = c.src[1:]
   # TODO: this check belongs in spec, not here
@@ -95,8 +98,8 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # resolve calls
   (UPat(Ops.CALL, name="c"), resolve_call),
 
-  # remove CONTIGUOUS if the BUFFER is already contiguous
-  (UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER), UPat()), name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
+  # remove CONTIGUOUS if the source is already contiguous
+  (UPat(Ops.RESHAPE, src=(UPat((Ops.BUFFER, Ops.CONTIGUOUS)), UPat()), name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
 
   # split_reduceop
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
@@ -132,6 +135,9 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
 
   # ** assign rules **
 
+  # collapse nested ASSIGN to the same buffer (e.g. __iadd__ in __setitem__)
+  (UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat(Ops.ASSIGN, name="src")), name="assign"), collapse_nested_assign),
+
   # move bitcast from assign target to source: a.bitcast(X).assign(src) -> a.assign(src.bitcast(a.dtype))
   (UPat(Ops.ASSIGN, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src")), name="assign"),
    lambda assign, target, src: target.assign(src.bitcast(target.dtype)).replace(tag=assign.tag)),

tinygrad/tensor.py

@@ -24,6 +24,7 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
 all_tensors: dict[weakref.ref[Tensor], None] = {}
+_pending_assigns: dict[UOp, list[UOp]] = {}  # buffer_uop -> [assign_uops in insertion order]
 def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
   with cpu_profile(TracingKey(name), "TINY"):
     # get tensors in scope
@@ -238,8 +239,8 @@ class Tensor(OpMixin):
     else:
       param = UOp.param(slot, self.dtype, self.shape, self.device)
     return Tensor(param, device=self.device)
-  def call(self, *lst:Tensor, fxn:Tensor|UOp, grad_fxn:Callable|None=None) -> Tensor:
-    return Tensor((fxn.uop if isinstance(fxn, Tensor) else fxn).call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn), device=self.device)
+  def call(self, *lst:Tensor, fxn:Tensor|UOp, **kwargs) -> Tensor:
+    return Tensor((fxn.uop if isinstance(fxn, Tensor) else fxn).call(*[t.uop for t in (self,)+lst], **kwargs), device=self.device)
 
   def custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None) -> list[Tensor]:
     """
@@ -271,6 +272,13 @@ class Tensor(OpMixin):
   @disable_gc()
   def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
     """Triggers the computation needed to create these Tensor(s)."""
+    # side-realize pending assigns for buffers referenced by these tensors
+    if _pending_assigns:
+      for buf in {u for t in (self,)+lst for u in t.uop.toposort() if u.op is Ops.BUFFER}:
+        for assign_uop in _pending_assigns.pop(buf, []):
+          becomes_map, schedule, var_vals = complete_create_schedule_with_vars(UOp.sink(assign_uop))
+          _apply_map_to_tensors(becomes_map, name="Apply Pending Assign")
+          run_schedule(schedule, var_vals, do_update_stats=do_update_stats)
     if len(to_realize:=[x for x in (self,)+lst if not x.uop.has_buffer_identity()]):
       run_schedule(*Tensor.schedule_with_vars(*to_realize), do_update_stats=do_update_stats)
     return self
@@ -299,7 +307,11 @@ class Tensor(OpMixin):
     if is_disk:
       self._buffer().copyin(x._data())
       return self
-    return self.replace(self._apply_uop(UOp.assign, x))
+    result = self._apply_uop(UOp.assign, x)
+    # track view assigns (not full-buffer or assign-chain) so they can be side-realized when the buffer is read
+    if (buf_uop:=self.uop.base).op is Ops.BUFFER and self.uop.op is not Ops.ASSIGN and not self.uop.has_buffer_identity():
+      _pending_assigns.setdefault(buf_uop, []).append(result.uop)
+    return self.replace(result)
 
   def detach(self) -> Tensor:
     """
@@ -1279,21 +1291,20 @@ class Tensor(OpMixin):
     return self._getitem(indices)
 
   def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None:
-    if isinstance(self.device, str) and self.device.startswith("DISK"):
-      self.realize()._getitem(indices).assign(v)
-      return
-    # NOTE: check that setitem target is valid first
-    if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
-    if self.requires_grad or v.requires_grad: raise NotImplementedError("setitem with requires_grad is not supported")
-    self.realize()
-    if not self.uop.is_writable_view(): raise RuntimeError("setitem target must be a writable view backed by a buffer")
-    res = self._getitem(indices, v)
-    # if shapes match and data is not shared it's a copy and we assign to self
-    if res.shape == self.shape and res.uop is not self.uop:
-      self.assign(res).realize()
-    else: # no copy, basic setitem
-      v = v.cast(res.dtype)._broadcast_to(_broadcast_shape(res.shape, v.shape)).contiguous()
-      res.assign(v).realize()
+    if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
+    if self.requires_grad or (isinstance(v, Tensor) and v.requires_grad): raise NotImplementedError("setitem with requires_grad is not supported")
+    idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices)
+    is_disk = isinstance(self.device, str) and self.device.startswith("DISK")
+    if any(isinstance(i, (Tensor, list, tuple)) for i in idx): # advanced setitem
+      if is_disk: raise RuntimeError("advanced setitem is not supported for DISK tensors")
+      if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
+      self.assign(self._getitem(indices, v))
+    else: # basic setitem
+      if is_disk: self[indices].assign(v)
+      else:
+        self.realize()
+        if not self.uop.is_writable_view(): raise RuntimeError("setitem target must be a writable view backed by a buffer")
+        self[indices].assign(v).realize()
 
   def __delitem__(self, indices) -> None:
     raise TypeError("Tensor does not support deleting items")

tinygrad/uop/ops.py

@@ -823,10 +823,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     src = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) + (() if device is None else (UOp(Ops.DEVICE, arg=device),))
     return UOp(Ops.PARAM, dtype, src, arg=slot)
 
-  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=()) -> UOp:
+  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), inline=False) -> UOp:
     # TODO: reenable this after ENCDEC is fixed
     #assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
-    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata))
+    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, inline))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
     contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
     placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
@@ -848,9 +848,10 @@ class KernelInfo:
 class CallInfo:
   grad_fxn: Callable|None = None
   metadata: tuple[Metadata, ...] = ()
+  inline: bool = False
   # grad_fxn can't be pickled, but metadata can
-  def __reduce__(self): return (CallInfo, (None, self.metadata))
-  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata})"
+  def __reduce__(self): return (CallInfo, (None, self.metadata, self.inline))
+  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata}, {self.inline})"
 
 # ******** ops in python ********