go

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '16'
+  CACHE_VERSION: '15'
   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: ${{ github.event_name != 'pull_request' && 'type=gha,mode=min' || '' }}
+        cache-to: 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_multitensor.py --durations=20
+        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
       - 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,7 +2,6 @@
 
 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 
@@ -13,7 +12,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:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export DP=8 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=1
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -21,14 +20,13 @@ 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="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="2.5e-4" END_LR="2.5e-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}
+export SEED=5760
 
-export JITBEAM=${JITBEAM:-3}
+export 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:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -21,10 +21,9 @@ 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="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="2.5e-4" END_LR="2.5e-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,6 +1,4 @@
 #!/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,12 +14,7 @@ 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):
-    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}")
+    batch //= len(a.device)
     dname = a.device[0]
   else: dname = a.device
   arch = getattr(Device[dname].renderer, "arch", "")
@@ -65,8 +65,6 @@ 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)
@@ -82,10 +80,9 @@ 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) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+    out = Tensor(Tensor.empty(batch//len(a.device), 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)
 
@@ -96,5 +93,4 @@ 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 self.max_context:
+    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:
       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 self.max_context != 0 and seqlen > 1:
+    if not Tensor.training 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,7 +31,6 @@ 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)")
@@ -73,9 +72,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 if args.top == -1 else items[:args.top]
+      sel = items[:10]
       table = [[name, time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in sel]
-      if args.top != -1 and (other:=items[len(sel):]):
+      if (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
+from tinygrad.device import CompileError, Device, Compiler
 if Device.DEFAULT=="METAL":
   from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
@@ -48,4 +48,28 @@ 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,19 +1,9 @@
 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")
@@ -27,7 +17,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="__kernel void test(__global int* a) { a[0] = 1; }".encode())
+      CLProgram(device, name="", lib=CLCompiler(device, "test").compile("__kernel void test(__global int* a) { a[0] = 1; }"))
     assert str(err.exception) == "OpenCL Error -46: CL_INVALID_KERNEL_NAME"
 
   def test_unaligned_copy(self):

test/null/test_egraph.py

@@ -0,0 +1,502 @@
+import unittest
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import Ops, UOp, GroupOp, PatternMatcher, UPat, graph_rewrite
+from tinygrad.uop.egraph import uf_find, uf_union, rewrite_all, EGraph, egraph_saturate, egraph_extract, node_cost, _rebuild_tree
+
+# *** test union-find ***
+
+class TestUnionFind(unittest.TestCase):
+  def test_find_self(self):
+    a, b = UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)
+    parent = {a: a, b: b}
+    self.assertIs(uf_find(parent, a), a)
+    self.assertIs(uf_find(parent, b), b)
+
+  def test_union_basic(self):
+    a, b = UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)
+    parent = {a: a, b: b}
+    size = {a: 1, b: 1}
+    root = uf_union(parent, size, a, b)
+    self.assertIs(uf_find(parent, a), uf_find(parent, b))
+    self.assertIs(root, uf_find(parent, a))
+
+  def test_union_chain(self):
+    a, b, c = UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2), UOp.const(dtypes.int, 3)
+    parent = {a: a, b: b, c: c}
+    size = {a: 1, b: 1, c: 1}
+    uf_union(parent, size, a, b)
+    uf_union(parent, size, b, c)
+    self.assertIs(uf_find(parent, a), uf_find(parent, c))
+
+  def test_union_idempotent(self):
+    a, b = UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)
+    parent = {a: a, b: b}
+    size = {a: 1, b: 1}
+    r1 = uf_union(parent, size, a, b)
+    r2 = uf_union(parent, size, a, b)
+    self.assertIs(r1, r2)
+
+# *** test rewrite_all ***
+
+class TestRewriteAll(unittest.TestCase):
+  def test_single_match(self):
+    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
+    a = UOp.variable("a", 0, 10)
+    results = rewrite_all(pm, a + 0)
+    self.assertEqual(len(results), 1)
+    self.assertIs(results[0], a)
+
+  def test_no_match(self):
+    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    results = rewrite_all(pm, a + b)
+    self.assertEqual(len(results), 0)
+
+  def test_multiple_matches(self):
+    pm = PatternMatcher([
+      (UPat.var("x") + 0, lambda x: x),
+      (UPat.var("x") * 1, lambda x: x),
+    ])
+    a = UOp.variable("a", 0, 10)
+    results = rewrite_all(pm, a + 0)
+    self.assertEqual(len(results), 1)
+    self.assertIs(results[0], a)
+
+  def test_both_rules_fire(self):
+    pm = PatternMatcher([
+      (UPat.var("x") + UPat.var("x"), lambda x: x * 2),
+      (UPat.var("x") + UPat.var("x"), lambda x: UOp(Ops.SHL, x.dtype, (x, x.const_like(1)))),
+    ])
+    a = UOp.variable("a", 0, 10)
+    results = rewrite_all(pm, a + a)
+    self.assertEqual(len(results), 2)
+
+  def test_const_folding(self):
+    pm = PatternMatcher([
+      (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name="a"),
+       lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else None),
+    ])
+    results = rewrite_all(pm, UOp.const(dtypes.int, 3) + UOp.const(dtypes.int, 4))
+    self.assertEqual(len(results), 1)
+    self.assertEqual(results[0].arg, 7)
+
+# *** test EGraph class ***
+
+class TestEGraphClass(unittest.TestCase):
+  def test_init(self):
+    a = UOp.variable("a", 0, 10)
+    expr = a + 0
+    eg = EGraph(expr)
+    self.assertEqual(len(eg.eclass), len(list(expr.toposort())))
+    self.assertIn(expr, eg.all_nodes)
+
+  def test_add_node(self):
+    a = UOp.variable("a", 0, 10)
+    eg = EGraph(a)
+    b = UOp.variable("b", 0, 10)
+    eg._add_node(b)
+    self.assertIn(b, eg.all_nodes)
+
+  def test_merge(self):
+    a = UOp.variable("a", 0, 10)
+    expr = a + 0
+    eg = EGraph(expr)
+    result = eg._merge(expr, a)
+    self.assertIsNotNone(result)
+    self.assertIs(uf_find(eg.parent, expr), uf_find(eg.parent, a))
+
+  def test_merge_idempotent(self):
+    a = UOp.variable("a", 0, 10)
+    eg = EGraph(a)
+    result = eg._merge(a, a)
+    self.assertIsNone(result)
+
+# *** test egraph_saturate ***
+
+class TestEGraphSaturate(unittest.TestCase):
+  def test_identity_rules(self):
+    pm = PatternMatcher([
+      (UPat.var("x") + 0, lambda x: x),
+      (UPat.var("x") * 1, lambda x: x),
+    ])
+    a = UOp.variable("a", 0, 10)
+    expr = a + 0
+    eclass = egraph_saturate(expr, pm)
+    # a+0 and a should be in the same e-class
+    a_class = expr_class = None
+    for canon, members in eclass.items():
+      if a in members: a_class = canon
+      if expr in members: expr_class = canon
+    self.assertIsNotNone(a_class)
+    self.assertIsNotNone(expr_class)
+    self.assertIs(a_class, expr_class)
+
+  def test_const_fold_saturation(self):
+    from tinygrad.uop.symbolic import symbolic_simple
+    c2, c3 = UOp.const(dtypes.int, 2), UOp.const(dtypes.int, 3)
+    expr = c2 + c3
+    eclass = egraph_saturate(expr, symbolic_simple)
+    c5 = UOp.const(dtypes.int, 5)
+    for canon, members in eclass.items():
+      if expr in members:
+        self.assertIn(c5, members, f"expected CONST(5) in eclass of 2+3, got {members}")
+        return
+    self.fail("expr not found in any eclass")
+
+  def test_no_rules_match(self):
+    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    eclass = egraph_saturate(a + b, pm)
+    for canon, members in eclass.items():
+      self.assertEqual(len(members), 1)
+
+  def test_max_iters_respected(self):
+    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
+    a = UOp.variable("a", 0, 10)
+    expr = a + 0
+    eclass = egraph_saturate(expr, pm, max_iters=1)
+    a_class = expr_class = None
+    for canon, members in eclass.items():
+      if a in members: a_class = canon
+      if expr in members: expr_class = canon
+    self.assertIs(a_class, expr_class)
+
+  def test_rebuilding_propagates(self):
+    """After a*0 merges with 0, rebuilding should create (0+a) which then matches x+0 -> x."""
+    pm = PatternMatcher([
+      (UPat.var("x") * 0, lambda x: x.const_like(0)),
+      (UPat.var("x") + 0, lambda x: x),
+    ])
+    a = UOp.variable("a", 0, 10)
+    expr = (a * 0) + a
+    eclass = egraph_saturate(expr, pm)
+    expr_cls = a_cls = None
+    for canon, members in eclass.items():
+      if expr in members: expr_cls = canon
+      if a in members: a_cls = canon
+    self.assertIsNotNone(expr_cls)
+    self.assertIsNotNone(a_cls)
+    self.assertIs(expr_cls, a_cls)
+
+  def test_rebuilding_chain(self):
+    """((a*0)+0)+b should simplify to b through multiple rebuild steps."""
+    pm = PatternMatcher([
+      (UPat.var("x") * 0, lambda x: x.const_like(0)),
+      (UPat.var("x") + 0, lambda x: x),
+    ])
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    expr = ((a * 0) + 0) + b
+    eclass = egraph_saturate(expr, pm)
+    expr_cls = b_cls = None
+    for canon, members in eclass.items():
+      if expr in members: expr_cls = canon
+      if b in members: b_cls = canon
+    self.assertIsNotNone(expr_cls)
+    self.assertIsNotNone(b_cls)
+    self.assertIs(expr_cls, b_cls)
+
+# *** test egraph_extract ***
+
+class TestEGraphExtract(unittest.TestCase):
+  def test_extract_identity(self):
+    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
+    a = UOp.variable("a", 0, 10)
+    self.assertIs(egraph_extract(a + 0, pm), a)
+
+  def test_extract_mul_identity(self):
+    pm = PatternMatcher([(UPat.var("x") * 1, lambda x: x)])
+    a = UOp.variable("a", 0, 10)
+    self.assertIs(egraph_extract(a * 1, pm), a)
+
+  def test_extract_const_fold(self):
+    from tinygrad.uop.symbolic import symbolic_simple
+    result = egraph_extract(UOp.const(dtypes.int, 2) + UOp.const(dtypes.int, 3), symbolic_simple)
+    self.assertEqual(result.op, Ops.CONST)
+    self.assertEqual(result.arg, 5)
+
+  def test_extract_chain(self):
+    pm = PatternMatcher([
+      (UPat.var("x") + 0, lambda x: x),
+      (UPat.var("x") * 1, lambda x: x),
+    ])
+    a = UOp.variable("a", 0, 10)
+    self.assertIs(egraph_extract((a + 0) * 1, pm), a)
+
+  def test_extract_no_change(self):
+    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    self.assertIs(egraph_extract(a + b, pm), a + b)
+
+  def test_extract_prefers_cheaper(self):
+    pm = PatternMatcher([(UPat.var("x") + UPat.var("x"), lambda x: x * 2)])
+    a = UOp.variable("a", 0, 10)
+    result = egraph_extract(a + a, pm)
+    self.assertEqual(result.op, Ops.ADD)  # ADD cost 1 < MUL cost 2
+
+  def test_extract_with_symbolic_simple(self):
+    from tinygrad.uop.symbolic import symbolic_simple
+    a = UOp.variable("a", 0, 10)
+    self.assertIs(egraph_extract((a + 0) * 1, symbolic_simple), a)
+
+  def test_combine_terms(self):
+    from tinygrad.uop.symbolic import symbolic
+    a = UOp.variable("a", 0, 10)
+    result = egraph_extract(a * 3 + a * 4, symbolic)
+    self.assertEqual(result.op, Ops.MUL)
+    self.assertEqual(result.src[1].arg, 7)
+
+  # *** tests that REQUIRE rebuilding ***
+
+  def test_rebuild_mul_zero_plus(self):
+    pm = PatternMatcher([
+      (UPat.var("x") * 0, lambda x: x.const_like(0)),
+      (UPat.var("x") + 0, lambda x: x),
+    ])
+    a = UOp.variable("a", 0, 10)
+    self.assertIs(egraph_extract((a * 0) + a, pm), a)
+
+  def test_rebuild_nested_zero(self):
+    pm = PatternMatcher([
+      (UPat.var("x") * 0, lambda x: x.const_like(0)),
+      (UPat.var("x") + 0, lambda x: x),
+    ])
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    self.assertIs(egraph_extract(((a * 0) + 0) + b, pm), b)
+
+  def test_rebuild_distribute_then_fold(self):
+    pm = PatternMatcher([(UPat.var("x") * 0, lambda x: x.const_like(0))])
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    result = egraph_extract((a + b) * 0, pm)
+    self.assertEqual(result.op, Ops.CONST)
+    self.assertEqual(result.arg, 0)
+
+  def test_rebuild_symmetric(self):
+    pm = PatternMatcher([
+      (UPat.var("x") * 0, lambda x: x.const_like(0)),
+      (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name="a"),
+       lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else None),
+    ])
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    result = egraph_extract((a * 0) + (b * 0), pm)
+    self.assertEqual(result.op, Ops.CONST)
+    self.assertEqual(result.arg, 0)
+
+  def test_rebuild_with_real_rules(self):
+    from tinygrad.uop.symbolic import symbolic_simple
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    self.assertIs(egraph_extract((a * 0) + (b * 1), symbolic_simple), b)
+
+  def test_rebuild_deep_chain(self):
+    pm = PatternMatcher([
+      (UPat.var("x") * 0, lambda x: x.const_like(0)),
+      (UPat.var("x") + 0, lambda x: x),
+      (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name="a"),
+       lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else None),
+    ])
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    c = UOp.variable("c", 0, 10)
+    self.assertIs(egraph_extract(((a * 0) + (b * 0)) + c, pm), c)
+
+# *** test cost model ***
+
+class TestCostModel(unittest.TestCase):
+  def test_const_is_free(self):
+    self.assertEqual(node_cost(UOp.const(dtypes.int, 0)), 0)
+
+  def test_add_is_cheap(self):
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    self.assertEqual(node_cost(a + b), 1)
+
+  def test_div_is_expensive(self):
+    a = UOp.variable("a", 0, 10).cast(dtypes.index)
+    b = UOp.variable("b", 1, 10).cast(dtypes.index)
+    self.assertEqual(node_cost(a // b), 5)
+
+  def test_mul_more_than_add(self):
+    a = UOp.variable("a", 0, 10)
+    b = UOp.variable("b", 0, 10)
+    self.assertGreater(node_cost(a * b), node_cost(a + b))
+
+# *** test e-graph matches greedy rewrite ***
+
+class TestEGraphVsGreedy(unittest.TestCase):
+  def test_matches_greedy_identity(self):
+    from tinygrad.uop.ops import graph_rewrite
+    from tinygrad.uop.symbolic import symbolic_simple
+    a = UOp.variable("a", 0, 10)
+    greedy = graph_rewrite(a + 0, symbolic_simple)
+    egraph = egraph_extract(a + 0, symbolic_simple)
+    self.assertIs(greedy, egraph)
+
+  def test_matches_greedy_const_fold(self):
+    from tinygrad.uop.ops import graph_rewrite
+    from tinygrad.uop.symbolic import symbolic_simple
+    expr = UOp.const(dtypes.int, 10) + UOp.const(dtypes.int, 20)
+    greedy = graph_rewrite(expr, symbolic_simple)
+    egraph = egraph_extract(expr, symbolic_simple)
+    self.assertEqual(greedy.op, Ops.CONST)
+    self.assertEqual(egraph.op, Ops.CONST)
+    self.assertEqual(greedy.arg, egraph.arg)
+
+  def test_matches_greedy_double_identity(self):
+    from tinygrad.uop.ops import graph_rewrite
+    from tinygrad.uop.symbolic import symbolic_simple
+    a = UOp.variable("a", 0, 10)
+    expr = (a + 0) * 1
+    self.assertIs(graph_rewrite(expr, symbolic_simple), a)
+    self.assertIs(egraph_extract(expr, symbolic_simple), a)
+
+# *** test e-graph beats greedy (phase-ordering problems) ***
+
+# helper PMs that create phase-ordering traps
+_pm_strength_reduce = PatternMatcher([
+  # strength reduction x*2 -> x+x fires FIRST and destroys the x*c form needed by combine-terms
+  (UPat.var('x') * UPat.cvar('c', vec=False), lambda x,c: x+x if c.arg == 2 else None),
+  # combine terms: x*c0 + x*c1 -> x*(c0+c1) can only match if both sides are x*c
+  (UPat.var('x') * UPat.cvar('c0') + UPat.var('x') * UPat.cvar('c1'), lambda x,c0,c1: x*(c0+c1)),
+  # constant folding
+  (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name='a'),
+   lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else
+             a.const_like(a.src[0].arg * a.src[1].arg) if a.op is Ops.MUL else None),
+  (UPat.var('x') + 0, lambda x: x),
+  (UPat.var('x') * 1, lambda x: x),
+])
+
+_pm_shift_reduce = PatternMatcher([
+  # strength reduction x*2 -> x<<1 fires FIRST and destroys the x*c form
+  (UPat.var('x') * UPat.cvar('c', vec=False),
+   lambda x,c: UOp(Ops.SHL, x.dtype, (x, x.const_like(1))) if c.arg == 2 else None),
+  (UPat.var('x') * UPat.cvar('c0') + UPat.var('x') * UPat.cvar('c1'), lambda x,c0,c1: x*(c0+c1)),
+  (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name='a'),
+   lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else
+             a.const_like(a.src[0].arg * a.src[1].arg) if a.op is Ops.MUL else None),
+  (UPat.var('x') + 0, lambda x: x),
+  (UPat.var('x') * 1, lambda x: x),
+])
+
+_pm_strength_fold = PatternMatcher([
+  # strength reduction x*2 -> x+x blocks two-stage folding (x*c1)*c2 -> x*(c1*c2)
+  (UPat.var('x') * UPat.cvar('c', vec=False), lambda x,c: x+x if c.arg == 2 else None),
+  ((UPat.var('x') * UPat.cvar('c1')) * UPat.cvar('c2'), lambda x,c1,c2: x*(c1*c2)),
+  (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name='a'),
+   lambda a: a.const_like(a.src[0].arg * a.src[1].arg) if a.op is Ops.MUL else None),
+])
+
+def _total_cost(u:UOp) -> int:
+  return sum(node_cost(n) for n in u.toposort())
+
+class TestEGraphBeatsGreedy(unittest.TestCase):
+  """Tests where the e-graph finds a cheaper result than the greedy rewriter due to phase-ordering.
+
+  The core problem: when Rule A fires first and transforms a node, it can destroy the pattern
+  that Rule B needs to match. Rule B would have led to a cheaper result, but the greedy rewriter
+  never tries it. The e-graph explores BOTH paths and picks the cheapest.
+  """
+  def test_strength_reduce_blocks_combine(self):
+    """a*2 + a*3: strength reduction x*2->x+x destroys the x*c form needed by combine-terms x*c0+x*c1->x*(c0+c1)."""
+    a = UOp.variable("a", 0, 10)
+    expr = a * 2 + a * 3
+    greedy = graph_rewrite(expr, _pm_strength_reduce)
+    egraph = egraph_extract(expr, _pm_strength_reduce)
+    # greedy: (a+a) + a*3 (cost 4) — strength reduction destroyed the a*2 pattern
+    self.assertEqual(greedy.op, Ops.ADD)
+    self.assertGreater(_total_cost(greedy), _total_cost(egraph))
+    # egraph: a*5 (cost 2) — combine-terms wins because the e-graph explored both paths
+    self.assertEqual(egraph.op, Ops.MUL)
+    self.assertEqual(egraph.src[1].arg, 5)
+
+  def test_shift_reduce_blocks_combine(self):
+    """a*2 + a*3: shift reduction x*2->x<<1 also destroys the combine-terms pattern."""
+    a = UOp.variable("a", 0, 10)
+    expr = a * 2 + a * 3
+    greedy = graph_rewrite(expr, _pm_shift_reduce)
+    egraph = egraph_extract(expr, _pm_shift_reduce)
+    self.assertEqual(greedy.op, Ops.ADD)
+    self.assertGreater(_total_cost(greedy), _total_cost(egraph))
+    self.assertEqual(egraph.op, Ops.MUL)
+    self.assertEqual(egraph.src[1].arg, 5)
+
+  def test_strength_reduce_chain(self):
+    """a*2 + a*3 + a*4: strength reduction causes greedy to miss the combined a*9."""
+    a = UOp.variable("a", 0, 10)
+    expr = a * 2 + a * 3 + a * 4
+    greedy = graph_rewrite(expr, _pm_strength_reduce)
+    egraph = egraph_extract(expr, _pm_strength_reduce)
+    self.assertGreater(_total_cost(greedy), _total_cost(egraph))
+
+  def test_strength_reduce_blocks_two_stage_fold(self):
+    """(a*2)*3: strength reduction x*2->x+x blocks two-stage constant folding (x*c1)*c2->x*(c1*c2)."""
+    a = UOp.variable("a", 0, 10)
+    expr = (a * 2) * 3
+    greedy = graph_rewrite(expr, _pm_strength_fold)
+    egraph = egraph_extract(expr, _pm_strength_fold)
+    # greedy: (a+a)*3 (cost 3) — can't fold constants because *2 was rewritten to +
+    self.assertGreater(_total_cost(greedy), _total_cost(egraph))
+    # egraph: a*6 (cost 2) — two-stage folding path was explored
+    self.assertEqual(egraph.op, Ops.MUL)
+    self.assertEqual(egraph.src[1].arg, 6)
+
+  def test_both_sides_strength_reduced(self):
+    """a*2 + a*2: both sides get strength-reduced, blocking combine-terms."""
+    a = UOp.variable("a", 0, 10)
+    expr = a * 2 + a * 2
+    greedy = graph_rewrite(expr, _pm_strength_reduce)
+    egraph = egraph_extract(expr, _pm_strength_reduce)
+    # greedy: (a+a)+(a+a) — both a*2 were rewritten before combine could fire
+    # egraph: a*4 — combine-terms path was found
+    self.assertEqual(egraph.op, Ops.MUL)
+    self.assertEqual(egraph.src[1].arg, 4)
+    # both have cost 2 here (shared subexpression), but egraph result is canonical
+    self.assertLessEqual(_total_cost(egraph), _total_cost(greedy))
+
+# *** test cycle-breaking in extraction ***
+
+class TestExtractionCycles(unittest.TestCase):
+  def test_self_referencing_eclass(self):
+    """x+0 -> x merges x+0 into x's eclass. Extraction must not recurse on the self-reference."""
+    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
+    a = UOp.variable("a", 0, 10)
+    self.assertIs(egraph_extract(a + 0, pm), a)
+
+  def test_nested_self_referencing_eclass(self):
+    """((a+0)+0)+0 — all merge into a's eclass. Deep self-reference chain."""
+    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
+    a = UOp.variable("a", 0, 10)
+    self.assertIs(egraph_extract(((a + 0) + 0) + 0, pm), a)
+
+  def test_mutual_eclass_cycle(self):
+    """Two eclasses whose best nodes reference each other — extraction must terminate via cycle-breaking cache."""
+    x = UOp.variable("x", 0, 10)
+    y = UOp.variable("y", 0, 10)
+    one = UOp.const(dtypes.index, 1)
+    two = UOp.const(dtypes.index, 2)
+    node1 = x + one  # E1's best, child x is in E2
+    node2 = y + two  # E2's best, child y is in E1
+    eclass_of = {node1: node1, x: node2, node2: node2, y: node1, one: one, two: two}
+    cost_of = {node1: (2, node1), node2: (2, node2), one: (0, one), two: (0, two)}
+    # without cycle-breaking cache, this would recurse: E1->E2->E1->...
+    result = _rebuild_tree(node1, eclass_of, cost_of)
+    self.assertIsNotNone(result)  # just verify it terminates
+
+  def test_mutual_rewrite_cycle(self):
+    """x+x <-> x*2 mutual rewrite. Both forms in same eclass, extraction picks cheaper (ADD)."""
+    pm = PatternMatcher([
+      (UPat.var("x") + UPat.var("x"), lambda x: x * 2),
+      (UPat.var("x") * UPat.cvar("c", vec=False), lambda x,c: x+x if c.arg == 2 else None),
+    ])
+    a = UOp.variable("a", 0, 10)
+    result = egraph_extract(a + a, pm)
+    self.assertEqual(result.op, Ops.ADD)  # ADD cost 1 < MUL cost 2
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)

test/test_edgecases.py

@@ -175,6 +175,7 @@ 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)
@@ -184,6 +185,7 @@ 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

@@ -0,0 +1,68 @@
+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_gc.py

@@ -72,6 +72,7 @@ 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_jit_footguns.py

@@ -111,18 +111,30 @@ class TestJitFootguns(unittest.TestCase):
       self.assertEqual(first.numpy().item(), expected_first)
       buf = new_buf
 
-  def test_slice_assign_works_without_realize(self):
-    """Slice assign then read from same buffer - pending assigns are side-realized."""
+  def test_slice_assign_requires_realize(self):
+    """Slice assign then read from same buffer - assign isn't connected to read without explicit realize()."""
     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(pos):
+    def f_broken(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(v_pos.bind(i)).item(), 4.0)
+      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)
 
   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"max diff {(ref-out).abs().max().item()}"
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
 
 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-5), f"max diff {(ref-out).abs().max().item()}"
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
 
 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"max diff {(ref-out).abs().max().item()}"
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
   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,6 +1,8 @@
-import unittest, random
+import unittest
+import random
+from os import getenv
 from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
-from tinygrad.helpers import Context, getenv
+from tinygrad.helpers import Context
 import numpy as np
 
 class TestSetitem(unittest.TestCase):
@@ -11,7 +13,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.0)),
+      ((4,4,4,4), (Ellipsis, slice(1,3), slice(None)), Tensor(4)),
       ((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),
@@ -48,10 +50,6 @@ 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
@@ -111,6 +109,8 @@ 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. this generates (1==gidx0)|(2==gidx0)|(3==gidx0)|(4==gidx0)|(5==gidx0) ...
+  # TODO: WEBGPU pipeline validation error
   @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, dtype=dtypes.int).contiguous()
+    t = Tensor.zeros(10,20,30,40,50).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), dtype=np.int32)
+    n = np.zeros((10,20,30,40,50))
     n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
-    np.testing.assert_equal(t.numpy(), n)
+    np.testing.assert_allclose(t.numpy(), n)
 
   def test_setitem_2d_tensor_indexing(self):
-    t = Tensor.zeros(2, dtype=dtypes.int).contiguous()
+    t = Tensor.zeros(2).contiguous()
     index = Tensor([[0, 1], [1,0]])
     v = Tensor.arange(2*2).reshape(2, 2).contiguous()
     t[index] = v
-    n = np.zeros((2,), dtype=np.int32)
+    n = np.zeros((2,))
     n[index.numpy()] = v.numpy()
-    np.testing.assert_equal(t.numpy(), n)
+    np.testing.assert_allclose(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 run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None:
+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
   Tensor.manual_seed(0)
-  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)
+  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)
   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=a_shard), b.shard(devs, axis=b_shard)
+  if multi: a, b = a.shard(devs, axis=0), b.shard(devs, axis=None)
   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=a_shard), b_ref.shard(devs, axis=b_shard)
+  if multi: a_ref, b_ref = a_ref.shard(devs, axis=0), b_ref.shard(devs, axis=None)
   with Context(ASM_GEMM=0):
     ref = asm_gemm(a_ref, b_ref)
     ref.sum().backward()
@@ -34,18 +34,10 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=N
 
   # 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.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)
+    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"
 
 # 128x smaller than usual
 # uses the UOp GEMM, runs on non CDNA4 and CI
@@ -58,8 +50,6 @@ 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):
@@ -80,12 +70,6 @@ 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,16 +560,30 @@ 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):
@@ -605,14 +619,26 @@ 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):
@@ -635,9 +661,16 @@ 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):
@@ -701,18 +734,34 @@ 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))
-    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)
+    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)
 
   def test_multiple_slice_assigns_then_read(self):
-    """Multiple non-overlapping slice assigns then read."""
+    """Multiple non-overlapping slice assigns then read - RAW dependencies must ensure all writes complete before 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, inline=True)
+    c = Tensor.call(a, b, fxn=plus_fxn)
     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, inline=True)
+    c = Tensor.call(a, b, fxn=plus_fxn, grad_fxn=grad_fxn)
     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, inline=True)
+    c = Tensor.call(a, b, fxn=plus_fxn)
     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), inline=True)
+    c = Tensor.call(a, b, fxn=a.as_param(0) @ b.as_param(1))
     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, inline=True)
+    c = Tensor.call(a, b, fxn=x@y)
 
     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, inline=True)
+    c = Tensor.call(a, b, fxn=complex_fxn)
     c.mean().backward()
 
     np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)

test/unit/test_disk_tensor.py

@@ -301,11 +301,6 @@ 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

@@ -1,7 +1,7 @@
 from typing import cast
 from dataclasses import replace
 import itertools
-from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, TracingKey, Context
+from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, EGRAPH, TracingKey, Context
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer, ProgramSpec
@@ -22,6 +22,13 @@ from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_s
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 
+def _sym_rewrite(sink:UOp, sym_pm:PatternMatcher, extra_pm:PatternMatcher|None=None, ctx=None, name:str|None=None) -> UOp:
+  """Symbolic rewrite: uses e-graph extraction when EGRAPH is set, otherwise greedy graph_rewrite."""
+  if EGRAPH:
+    from tinygrad.uop.egraph import egraph_rewrite
+    return egraph_rewrite(sink, sym_pm, extra_pm, ctx=ctx, name=name)
+  return graph_rewrite(sink, sym_pm+extra_pm if extra_pm is not None else sym_pm, ctx=ctx, name=name)
+
 def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
   if ren is None: ren = Renderer()
 
@@ -41,7 +48,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
     sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")
 
     # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
-    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
+    sink = _sym_rewrite(sink, sym, pm_flatten_range, name="initial symbolic")
 
     # optimize (schedule) the AST
     sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")
@@ -53,10 +60,10 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
     sink = apply_opts(sink, ren)
 
   # ** expander (expand_rewrite) **
-  sink = graph_rewrite(sink, sym+pm_move_where_on_load, name="postopt symbolic")
+  sink = _sym_rewrite(sink, sym, pm_move_where_on_load, name="postopt symbolic")
 
   # expand
-  sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
+  sink = _sym_rewrite(sink, sym, pm_pre_expander+pm_group_for_reduce+expander, name="expander")
 
   # add locals
   sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")
@@ -74,32 +81,33 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
   sink = graph_rewrite(sink, pm_add_loads, name="** add loads (code)")
 
   # devectorize (TODO: does this need opts?)
-  if DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing
-  elif DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing
-  else: pm_devectorize = sym+load_store_folding+correct_load_store+load_store_indexing
-  if DEVECTORIZE >= 0: sink = graph_rewrite(sink, pm_devectorize, ctx=ren, name="devectorize")
+  if DEVECTORIZE >= 2: pm_devec_extra = load_store_folding+load_store_indexing
+  elif DEVECTORIZE: pm_devec_extra = devectorize+load_store_folding+correct_load_store+load_store_indexing
+  else: pm_devec_extra = load_store_folding+correct_load_store+load_store_indexing
+  if DEVECTORIZE >= 0: sink = _sym_rewrite(sink, sym, pm_devec_extra, ctx=ren, name="devectorize")
 
   # lower the index dtype to a concrete int
   sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing+gep_pushing, ctx=ren.device, name="lower all index dtypes")
-  sink = graph_rewrite(sink, symbolic, name="post index symbolic")
+  sink = _sym_rewrite(sink, symbolic, name="post index symbolic")
 
   # optional pre matcher
   if ren.pre_matcher is not None: sink = graph_rewrite(sink, ren.pre_matcher, name="pre_matcher")
 
   # decompositions
   supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, ren.device, bool(DISABLE_FAST_IDIV))
-  pm_transcendental = symbolic_simple+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
-  sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions")
+  pm_decomp_extra = get_late_rewrite_patterns(supported_ops, ren.device, bool(DISABLE_FAST_IDIV))
+  pm_transcend_extra = get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
+  sink = _sym_rewrite(sink, symbolic_simple, pm_decomp_extra, ctx=ren.device, name="decompositions")
   if not is_dtype_supported(dtypes.long, ren.device) or dtypes.long in EMULATED_DTYPES.tolist(dtypes):
     sink = graph_rewrite(sink, pm_long_decomp, name="decomp long -> int", bottom_up=True)
   for fr, to in [(fr, next((to for to in promo_lattice[fr] if is_dtype_supported(to, ren.device)), dtypes.float))
                  for fr in EMULATED_DTYPES.tolist(dtypes) if fr in dtypes.floats]:
     sink = graph_rewrite(sink, pm_float_decomp, ctx=(fr, to), name=f"decomp {fr} -> {to}", bottom_up=True)
-  sink = graph_rewrite(sink, pm_transcendental, ctx=ren.device, name="transcendental")
+  sink = _sym_rewrite(sink, symbolic_simple, pm_transcend_extra, ctx=ren.device, name="transcendental")
 
   # final rules for the renderer (without sym)
   extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
+  pm_decomp = symbolic_simple+pm_decomp_extra
   pm_final_rewrite = pm_decomp+pm_render+extra_matcher+pm_split_ends
   sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren.device, name="final rewrite")
 
@@ -139,6 +147,7 @@ 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, TYPE_CHECKING
+from typing import Any, Generic, TypeVar, Iterator, Generator
 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
-if TYPE_CHECKING: from tinygrad.renderer import Renderer
+from tinygrad.renderer import Renderer
 
 # **************** Device ****************
 
@@ -278,34 +278,37 @@ class Compiler:
   def disassemble(self, lib:bytes): pass
 
 @dataclass(frozen=True)
-class CompilerSet: cset:list[tuple[type[Renderer]|functools.partial, ContextVar|None]]; ctrl_var:ContextVar|None = None # noqa: E702
+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 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[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)
+    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))
 
   @property
-  def renderer(self) -> Renderer: return self._select_compiler_pair()
+  def renderer(self) -> Renderer: return self._select_compiler_pair()[0]
 
   @property
   def compiler(self) -> Compiler:
-    if (ret:=self.renderer.compiler) is None: raise RuntimeError(f"no compiler for {self.device}")
+    if (ret:=self.renderer.compiler or self._select_compiler_pair()[1]) is None: raise RuntimeError(f"no compiler for {self.device}")
     return ret
 
-  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 _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 _select_compiler_pair(self) -> Renderer:
+  def _select_compiler_pair(self) -> tuple[Renderer, Compiler|None]:
     # 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 [])
 
@@ -394,18 +397,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) in default_comp_pairs.items():
-          d.comp_sets = {k:(None,r)} # env var set to None, so it doesn't interfere
+        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
           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)} to make default)' if cc_ctrl_var is not None else ''
+            set_text = f'({cc_ctrl_var.key}={d._compiler_name(r, c)} 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)} {default_text}")
+            compilers_results.append(f"{colored('+', 'green')} {d._compiler_name(r, c)} {default_text}")
             any_works = True
-          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {d._compiler_name(r)}: {e}")
+          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {d._compiler_name(r, c)}: {e}")
       finally:
         # put the defaults back!
         d.comp_sets, d.comps_ctrl_var = default_comp_pairs, cc_ctrl_var

tinygrad/helpers.py

@@ -180,17 +180,16 @@ SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), Contex
 RING, ALL2ALL = ContextVar("RING", 1), ContextVar("ALL2ALL", 0)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
-CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
+CORRECT_DIVMOD_FOLDING, FUSE_OPTIM, EGRAPH = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0), ContextVar("EGRAPH", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
 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_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)
+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)
 NULL_IR3, NULL_NAK, NULL_ALLOW_COPYOUT = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0), ContextVar("NULL_ALLOW_COPYOUT", 0)
-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)
+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", "")
 # 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,8 +360,7 @@ 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, inline=True)
+    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)
     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
+from typing import Callable, cast, TYPE_CHECKING
 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
-from tinygrad.device import Compiler
+if TYPE_CHECKING: from tinygrad.device import Compiler
 
 @dataclass(frozen=True)
 class Estimates:
@@ -150,8 +150,7 @@ class Renderer:
   pre_matcher: PatternMatcher|None = None
   extra_matcher: PatternMatcher|None = None
   code_for_op: dict[Ops, Callable] = {}
-
-  compiler: Compiler = Compiler()
+  compiler: Compiler|None = None
 
   def __reduce__(self): return self.__class__, ()
   def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")

tinygrad/renderer/cstyle.py

@@ -340,9 +340,7 @@ class IntelRenderer(OpenCLRenderer):
 class MetalRenderer(CStyleLanguage):
   device = "METAL"
   shared_max = 32768
-  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 []
+  def __init__(self): self.tensor_cores = tc.metal if hasattr(os, 'uname') and os.uname().machine == "arm64" else []
 
   # language options
   kernel_typedef = "kernel void"
@@ -384,17 +382,15 @@ 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, 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)
+  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,)
 
   # language options
   # https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html
@@ -551,6 +547,7 @@ 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,9 +205,6 @@ 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,11 +144,9 @@ 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="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 __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 __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", "([] if CPU_CC.value == 'LVP' or bool(CPU_LVP) else ['tinymesa']) + ['tinymesa_cpu']", [
+    case "mesa": return load("mesa", "['tinymesa_cpu', 'tinymesa']", [
         *[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,8 +134,7 @@ 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=["from tinygrad.helpers import CPU_CC, CPU_LVP", "import gzip, base64"],
-  epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
+  prolog=["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,9 +4,8 @@ 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', ([] if CPU_CC.value == 'LVP' or bool(CPU_LVP) else ['tinymesa']) + ['tinymesa_cpu'])
+dll = c.DLL('mesa', ['tinymesa_cpu', 'tinymesa'])
 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
+from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet, CompilerPair
 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, AMD_HIPCC, ceildiv, unwrap
+from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, 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([(functools.partial(AMDHIPRenderer, self.arch), None),
-                             (functools.partial(AMDLLVMRenderer, self.arch), AMD_LLVM),
-                             (functools.partial(AMDHIPCCRenderer, self.arch), AMD_HIPCC)], ctrl_var=AMD_CC)
+    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)
 
     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, CompilerSet
+from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError, CompilerPair, 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, 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(),
+    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(),
                                                         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,9 +125,8 @@ 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
-    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))
-
+    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))
   def synchronize(self):
     check(cl.clFinish(self.queue))
     self.pending_copyin.clear()

tinygrad/runtime/ops_cpu.py

@@ -2,12 +2,13 @@ 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
+from tinygrad.device import BufferSpec, DMACPURef, CompilerSet, CompilerPair
 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
 
@@ -133,5 +134,6 @@ class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
     self.tasks:queue.Queue = queue.Queue()
     CPUWorker(self, self.tasks, thread_id=0).start()
-    compilers = CompilerSet([(ClangJITRenderer, None), (CPULLVMRenderer, CPU_LLVM), (LVPRenderer, CPU_LVP)], ctrl_var=CPU_CC)
+    compilers = CompilerSet([CompilerPair(ClangJITRenderer, None), CompilerPair(CPULLVMRenderer, CPULLVMCompiler, ctrl_var=CPU_LLVM),
+                             CompilerPair(LVPRenderer, None, ctrl_var=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, CUDA_NVCC
-from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerSet
+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.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
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler, NVCCCompiler
 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([(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)
+    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)
     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
+from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, CompilerSet, CompilerPair
 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,8 +45,6 @@ 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*);',
@@ -118,11 +116,28 @@ 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 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}"
+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)
     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']
@@ -131,25 +146,8 @@ class DSPCompiler(Compiler):
         self.link_ld.write(f"SECTIONS {{ . = 0x0; {sections_link}\n /DISCARD/ : {{ *(.note .note.* .gnu.hash .comment) }} }}".encode())
         self.link_ld.flush()
 
-      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))
+      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))
       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)
@@ -270,7 +268,6 @@ 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
+from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet, CompilerPair
 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([(functools.partial(HIPRenderer, self.arch), None)])
+    compilers = CompilerSet([CompilerPair(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, decimal, platform
-from tinygrad.helpers import prod, to_mv, round_up, cache_dir, PROFILE, ProfileRangeEvent, cpu_profile, unwrap, suppress_finalizing
+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 tinygrad.runtime.support.objc as objc
-from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent, CompilerSet
+from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent, CompilerSet, CompilerPair
 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([(MetalRenderer, None)]),
+    super().__init__(device, MetalAllocator(self), CompilerSet([CompilerPair(MetalRenderer, MetalCompiler), CompilerPair(MetalRenderer, Compiler)]),
       functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None)
 
   def synchronize(self):
@@ -54,12 +54,20 @@ 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.
-  import tinygrad.runtime.autogen.llvm as _
+  with contextlib.suppress(FileNotFoundError, ModuleNotFoundError):
+    import tinygrad.runtime.autogen.llvm # noqa: F401
   support = DLL("MTLCompiler", "MTLCompiler")
   support.MTLCodeGenServiceCreate.restype = ctypes.c_void_p
 
@@ -110,9 +118,15 @@ class MetalCompiler(Compiler):
 class MetalProgram:
   def __init__(self, dev:MetalDevice, name:str, lib:bytes, **kwargs):
     self.dev, self.name, self.lib = dev, name, lib
-    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)
+    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
     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, Allocator, CompilerSet
+from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet, CompilerPair
 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([(renderer, None), (functools.partial(IR3Renderer, 0x6030001), NULL_IR3), # adreno 630
-                             (functools.partial(NAKRenderer, "sm_120", 48), NULL_NAK)]) # 5090
+    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
     super().__init__(device, NullAllocator(self), compilers, functools.partial(NullProgram, device), NullGraph)

tinygrad/runtime/ops_nv.py

@@ -6,11 +6,12 @@ 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, CompilerSet
+from tinygrad.device import Compiled, BufferSpec, CompilerPair, 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 CUDARenderer
+from tinygrad.renderer.cstyle import NVRenderer
+from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler
 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
@@ -619,9 +620,10 @@ 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)
 
-    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)])
+    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)])
     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
+from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet, CompilerPair
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer
@@ -214,14 +214,9 @@ 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
@@ -241,6 +236,9 @@ 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
@@ -248,4 +246,4 @@ class PythonAllocator(Allocator['PythonDevice']):
 
 class PythonDevice(Compiled):
   def __init__(self, device:str):
-    super().__init__(device, PythonAllocator(self), CompilerSet([(PythonRenderer, None)]), PythonProgram)
+    super().__init__(device, PythonAllocator(self), CompilerSet([CompilerPair(PythonRenderer, PythonCompiler)]), 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, Device
+from tinygrad.device import BufferSpec, CompilerSet, CompilerPair, 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 CLDevice
+from tinygrad.runtime.ops_cl import CLCompiler, 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, DEBUG
+from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC, PROFILE
 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,6 +49,10 @@ 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})
@@ -227,7 +231,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.NIR = buf_dtypes, name, isinstance(dev.renderer, IR3Renderer)
+    self.buf_dtypes, self.name, self.lib, self.NIR = buf_dtypes, name, lib, isinstance(dev.renderer, IR3Renderer)
 
     if self.NIR:
       from tinygrad.runtime.support.compiler_mesa import IR3Compiler
@@ -248,9 +252,7 @@ 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(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)])
+    else: self._parse_lib()
 
     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
@@ -272,21 +274,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, lib):
+  def _parse_lib(self):
     # Extract image binary
-    self.image_size = _read_lib(lib, 0x100)
-    self.image = bytearray(lib[(image_offset:=_read_lib(lib, 0xc0)):image_offset+self.image_size])
+    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])
 
     # Parse image descriptors
-    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)
+    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)
 
     # Fill up constants and buffers info
     self.consts_info = []
 
     # Collect sampler info.
-    self.samp_cnt = samp_cnt_in_file = _read_lib(lib, image_desc_off + 0xdc)
+    self.samp_cnt = samp_cnt_in_file = _read_lib(self.lib, image_desc_off + 0xdc)
     assert self.samp_cnt <= 1, "Up to one sampler supported"
     if self.samp_cnt:
       self.samp_cnt += 1
@@ -296,8 +298,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(lib):
-      length, _, _, offset_words, _, _, _, typ = struct.unpack("8I", lib[bdoff:bdoff+32])
+    while bdoff + 32 <= len(self.lib):
+      length, _, _, offset_words, _, _, _, typ = struct.unpack("8I", self.lib[bdoff:bdoff+32])
       if length == 0: break
       binfos.append((offset_words * 4, typ))
       bdoff += length
@@ -307,16 +309,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(lib, 0xb0) != 0: # check if we have constants.
-      cdoff = _read_lib(lib, 0xac)
+    if _read_lib(self.lib, 0xb0) != 0: # check if we have constants.
+      cdoff = _read_lib(self.lib, 0xac)
       while cdoff + 40 <= image_offset:
-        cnst, offset_words, _, is32 = struct.unpack("I", lib[cdoff:cdoff+4])[0], *struct.unpack("III", lib[cdoff+16:cdoff+28])
+        cnst, offset_words, _, is32 = struct.unpack("I", self.lib[cdoff:cdoff+4])[0], *struct.unpack("III", self.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(lib, 0x34)
-    self.fregs, self.hregs = _read_lib(lib, reg_desc_off + 0x14), _read_lib(lib, reg_desc_off + 0x18)
+    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)
 
 class QCOMTextureInfo:
   def __init__(self, pitch:int, real_stride:int, desc:list[int], ibo:list[int]):
@@ -383,8 +385,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")
 
-    self.cl_dev = CLDevice(device)
-    compilers = CompilerSet(ctrl_var=QCOM_CC, cset=[(QCOMRenderer, None), (functools.partial(IR3Renderer, info.chip_id), QCOM_IR3)])
+    compilers = CompilerSet(ctrl_var=QCOM_CC, cset=[CompilerPair(QCOMRenderer, functools.partial(QCOMCompiler, device)),
+                                                    CompilerPair(functools.partial(IR3Renderer, info.chip_id), None, 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, BufferSpec, CompilerSet
+from tinygrad.device import  Compiled, Allocator, Compiler, BufferSpec, CompilerSet, CompilerPair
 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([(WGSLRenderer, None)]),
+    super().__init__(device, WebGpuAllocator(self), CompilerSet([CompilerPair(WGSLRenderer, Compiler)]),
       functools.partial(WebGPUProgram, (self.device_res, webgpu.WGPUFeatureName_TimestampQuery in supported)))
 
   def synchronize(self):

tinygrad/runtime/support/compiler_cuda.py

@@ -1,4 +1,5 @@
-import hashlib, tempfile, ctypes, re, pathlib
+import subprocess, hashlib, tempfile, ctypes, re, pathlib
+from typing import Callable
 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
@@ -41,32 +42,36 @@ 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 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}']
+class CUDACompiler(Compiler):
+  def __init__(self, arch:str, cache_key:str="cuda"):
+    self.arch, self.compile_options = 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(self, src:str) -> bytes:
+  def _compile_program(self, src:str, nvrtc_get_content:Callable, nvrtc_get_size:Callable) -> 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.nvrtcGetPTX if self.ptx else nvrtc.nvrtcGetCUBIN,
-                      nvrtc.nvrtcGetPTXSize if self.ptx else nvrtc.nvrtcGetCUBINSize, nvrtc_check)
+    data = _get_bytes(prog, nvrtc_get_content, nvrtc_get_size, nvrtc_check)
     nvrtc_check(nvrtc.nvrtcDestroyProgram(ctypes.byref(prog)))
     return data
-  def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch, ptx=self.ptx)
+  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)
 
 class NVCCCompiler(Compiler):
-  def __init__(self, arch:str, ptx:bool=True, cache_key:str="cuda", extra_options:list[str]=[]):
-    assert ptx, "NVCCCompiler cubin support unimplemented"
+  def __init__(self, arch:str, extra_options:list[str]=[]):
     self.arch, self.extra_options = arch, extra_options
-    super().__init__(f"compile_nvcc_{cache_key}_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}")
+    super().__init__(f"compile_nvcc_{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()
-      system(f"nvcc -arch={self.arch} -ptx -o {libf.name} {srcf.name}" + ' '.join(self.extra_options))
+      subprocess.run(["nvcc", f"-arch={self.arch}", "-ptx", "-o", libf.name, srcf.name] + self.extra_options, check=True)
       return libf.read()
   def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch, ptx=True)
 

tinygrad/schedule/rangeify.py

@@ -33,10 +33,6 @@ 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)
@@ -79,8 +75,9 @@ mop_cleanup = PatternMatcher([
 ])
 
 def resolve_call(c:UOp) -> UOp|None:
-  # we only resolve here if the call is inlined
-  if not c.arg.inline: return 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
   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
@@ -98,8 +95,8 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # resolve calls
   (UPat(Ops.CALL, name="c"), resolve_call),
 
-  # 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)),
+  # 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)),
 
   # split_reduceop
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
@@ -135,9 +132,6 @@ 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,7 +24,6 @@ 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
@@ -239,8 +238,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, **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 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 custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None) -> list[Tensor]:
     """
@@ -272,13 +271,6 @@ 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
@@ -307,11 +299,7 @@ class Tensor(OpMixin):
     if is_disk:
       self._buffer().copyin(x._data())
       return self
-    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)
+    return self.replace(self._apply_uop(UOp.assign, x))
 
   def detach(self) -> Tensor:
     """
@@ -1291,20 +1279,21 @@ class Tensor(OpMixin):
     return self._getitem(indices)
 
   def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None:
-    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()
+    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()
 
   def __delitem__(self, indices) -> None:
     raise TypeError("Tensor does not support deleting items")

tinygrad/uop/egraph.py

@@ -0,0 +1,227 @@
+# e-graph (equality saturation) for UOp rewriting
+# instead of greedy first-match rewriting, we explore ALL equivalent forms and extract the cheapest
+from __future__ import annotations
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, graph_rewrite
+
+# *** union-find (keyed by UOp identity) ***
+
+def uf_find(parent:dict[UOp, UOp], x:UOp) -> UOp:
+  while parent[x] is not x:
+    parent[x] = parent[parent[x]]
+    x = parent[x]  # path compression
+  return x
+
+def uf_union(parent:dict[UOp, UOp], size:dict[UOp, int], a:UOp, b:UOp) -> UOp:
+  a, b = uf_find(parent, a), uf_find(parent, b)
+  if a is b: return a
+  if size[a] < size[b]: a, b = b, a  # merge smaller into larger
+  parent[b] = a
+  size[a] += size[b]
+  return a
+
+# *** e-graph core ***
+
+def rewrite_all(pm:PatternMatcher, uop:UOp, ctx=None) -> list[UOp]:
+  """Apply ALL matching rewrite rules to uop, returning every distinct result."""
+  results: list[UOp] = []
+  seen: dict[UOp, None] = {}
+  for _, match, early_reject in pm.pdict.get(uop.op, []):
+    if not early_reject.issubset({u.op for u in uop.src}): continue
+    try: ret = match(uop, ctx)
+    except Exception: continue  # skip rules that crash on this node (e.g. division by zero in divmod folding)
+    if ret is not None and ret is not uop and ret not in seen:
+      results.append(ret)
+      seen[ret] = None
+  return results
+
+class EGraph:
+  """E-graph with full equality saturation (including rebuilding)."""
+  __slots__ = ("parent", "size", "eclass", "eclass_uses", "all_nodes")
+  def __init__(self, root:UOp):
+    nodes = list(root.toposort())
+    self.parent: dict[UOp, UOp] = {u: u for u in nodes}
+    self.size: dict[UOp, int] = {u: 1 for u in nodes}
+    self.eclass: dict[UOp, dict[UOp, None]] = {u: {u: None} for u in nodes}  # canonical -> members
+    # canonical eclass representative -> dict of nodes that USE this eclass as a child
+    self.eclass_uses: dict[UOp, dict[UOp, None]] = {u: {} for u in nodes}
+    self.all_nodes: dict[UOp, None] = dict.fromkeys(nodes)
+    # build initial parent-child uses
+    for u in nodes:
+      for s in u.src:
+        canon = uf_find(self.parent, s)
+        self.eclass_uses.setdefault(canon, {})[u] = None
+
+  def _add_node(self, u:UOp):
+    """Register a new UOp (and its subtree) in the e-graph."""
+    for sub in u.toposort():
+      if sub in self.parent: continue
+      self.parent[sub] = sub
+      self.size[sub] = 1
+      self.eclass[sub] = {sub: None}
+      self.all_nodes[sub] = None
+      self.eclass_uses[sub] = {}
+      for s in sub.src:
+        canon = uf_find(self.parent, s)
+        self.eclass_uses.setdefault(canon, {})[sub] = None
+
+  def _merge(self, a:UOp, b:UOp) -> UOp|None:
+    """Merge two e-classes. Returns the winner, or None if already merged."""
+    ra, rb = uf_find(self.parent, a), uf_find(self.parent, b)
+    if ra is rb: return None
+    winner = uf_union(self.parent, self.size, ra, rb)
+    loser = rb if winner is ra else ra
+    self.eclass[winner] = {**self.eclass[winner], **self.eclass[loser]}
+    # merge uses
+    winner_uses = self.eclass_uses.setdefault(winner, {})
+    winner_uses.update(self.eclass_uses.pop(loser, {}))
+    del self.eclass[loser]
+    return winner
+
+  def _canonical(self, u:UOp) -> UOp:
+    """Rebuild node with canonical representative for each child's eclass."""
+    if not u.src: return u
+    new_src = []
+    for s in u.src:
+      canon = uf_find(self.parent, s)
+      members = self.eclass.get(canon)
+      if members is not None:
+        best = min(members, key=lambda m: (len(m.src), m.op.value, m.arg if isinstance(m.arg, (int, float, str)) else 0))
+        new_src.append(best)
+      else:
+        new_src.append(s)
+    new_src_tuple = tuple(new_src)
+    if new_src_tuple == u.src: return u
+    return UOp(u.op, u.dtype, new_src_tuple, u.arg, u.tag)
+
+  def _rebuild(self, dirty:dict[UOp, None]) -> list[tuple[UOp, UOp]]:
+    """Rebuild parents of dirty eclasses, creating canonical versions."""
+    new_equalities: list[tuple[UOp, UOp]] = []
+    affected: dict[UOp, None] = {}
+    for d in dirty:
+      canon = uf_find(self.parent, d)
+      affected.update(self.eclass_uses.get(canon, {}))
+    for u in affected:
+      rebuilt = self._canonical(u)
+      if rebuilt is not u:
+        if rebuilt in self.parent and uf_find(self.parent, rebuilt) is uf_find(self.parent, u): continue
+        self._add_node(rebuilt)
+        new_equalities.append((u, rebuilt))
+    return new_equalities
+
+def egraph_saturate(root:UOp, pm:PatternMatcher, max_iters:int=10, ctx=None) -> dict[UOp, dict[UOp, None]]:
+  """Build an e-graph with full equality saturation (with rebuilding). Returns eclass map."""
+  eg = EGraph(root)
+  node_limit = len(eg.all_nodes) * 3  # stop growing at 3x initial size to prevent combinatorial blowup
+  worklist: dict[UOp, None] = dict(eg.all_nodes)  # nodes to match rules on
+  for _ in range(max_iters):
+    # phase 1: match rules only on worklist nodes
+    new_equalities: list[tuple[UOp, UOp]] = []
+    next_worklist: dict[UOp, None] = {}
+    prev_nodes = dict(eg.all_nodes)
+    for u in list(worklist):
+      if len(eg.all_nodes) >= node_limit: break
+      for new in rewrite_all(pm, u, ctx):
+        if new in eg.parent and uf_find(eg.parent, new) is uf_find(eg.parent, u): continue
+        eg._add_node(new)
+        new_equalities.append((u, new))
+    # all newly added nodes (including sub-nodes of rewrite results) go on next worklist
+    for u in eg.all_nodes:
+      if u not in prev_nodes: next_worklist[u] = None
+    if not new_equalities: break
+
+    # phase 2: merge eclasses, then rebuild canonical forms (no rule matching in rebuild)
+    while new_equalities:
+      dirty: dict[UOp, None] = {}
+      for a, b in new_equalities:
+        merged = eg._merge(a, b)
+        if merged is not None: dirty[merged] = None
+      if not dirty: break
+      new_equalities = eg._rebuild(dirty)
+      for _, b in new_equalities: next_worklist[b] = None
+    worklist = next_worklist
+
+  return eg.eclass
+
+# *** cost model ***
+
+OP_COST: dict[Ops, int] = {
+  Ops.CONST: 0, Ops.VCONST: 0, Ops.DEFINE_VAR: 0,
+  Ops.ADD: 1, Ops.MUL: 2, Ops.SUB: 1, Ops.NEG: 1,
+  Ops.IDIV: 5, Ops.MOD: 5, Ops.FDIV: 3,
+  Ops.SHL: 1, Ops.SHR: 1,
+  Ops.AND: 1, Ops.OR: 1, Ops.XOR: 1,
+  Ops.MAX: 1, Ops.CMPLT: 1, Ops.CMPNE: 1, Ops.CMPEQ: 1,
+  Ops.CAST: 1, Ops.BITCAST: 1,
+  Ops.WHERE: 2, Ops.MULACC: 2,
+  Ops.EXP2: 8, Ops.LOG2: 8, Ops.SIN: 8, Ops.SQRT: 4, Ops.RECIPROCAL: 3,
+  Ops.POW: 10, Ops.TRUNC: 1,
+}
+
+def node_cost(u:UOp) -> int:
+  c = OP_COST.get(u.op, 3)
+  # tiebreaker: penalize non-canonical operand order (consts should be on the right for commutative ops)
+  if len(u.src) == 2 and u.src[0].op is Ops.CONST and u.src[1].op is not Ops.CONST: c += 1
+  return c
+
+# *** extraction ***
+
+def egraph_extract(root:UOp, pm:PatternMatcher, max_iters:int=10, ctx=None) -> UOp:
+  """Run equality saturation on root, then extract the cheapest equivalent expression."""
+  eclass = egraph_saturate(root, pm, max_iters, ctx)
+
+  # build eclass lookup: node -> canonical eclass representative
+  eclass_of: dict[UOp, UOp] = {}
+  for canon, members in eclass.items():
+    for u in members: eclass_of[u] = canon
+
+  all_nodes: list[UOp] = [u for members in eclass.values() for u in members]
+
+  # bottom-up DP: for each eclass, find the cheapest representative
+  cost_of: dict[UOp, tuple[int, UOp]] = {}  # eclass_canon -> (cost, best_uop)
+
+  depth_cache: dict[UOp, int] = {}
+  def _depth(u:UOp) -> int:
+    if u in depth_cache: return depth_cache[u]
+    depth_cache[u] = 0  # break cycles
+    depth_cache[u] = (1 + max((_depth(s) for s in u.src), default=0)) if u.src else 0
+    return depth_cache[u]
+
+  for u in sorted(all_nodes, key=_depth):
+    canon = eclass_of[u]
+    child_cost = 0
+    for s in u.src:
+      if (s_canon := eclass_of.get(s)) is not None and s_canon in cost_of: child_cost += cost_of[s_canon][0]
+      else: child_cost += node_cost(s)
+    total = node_cost(u) + child_cost
+    if canon not in cost_of or total < cost_of[canon][0]:
+      cost_of[canon] = (total, u)
+
+  root_canon = eclass_of.get(root)
+  if root_canon is not None and root_canon in cost_of: return _rebuild_tree(cost_of[root_canon][1], eclass_of, cost_of)
+  return root
+
+def _rebuild_tree(u:UOp, eclass_of:dict[UOp, UOp], cost_of:dict[UOp, tuple[int, UOp]], cache:dict[UOp, UOp]|None=None) -> UOp:
+  """Recursively rebuild a UOp tree, picking the cheapest representative for each child's eclass."""
+  if not u.src: return u
+  if cache is None: cache = {}
+  new_src = []
+  for s in u.src:
+    s_canon = eclass_of.get(s)
+    if s_canon is not None and s_canon in cost_of:
+      if s_canon in cache: new_src.append(cache[s_canon])
+      else:
+        cache[s_canon] = s  # placeholder breaks cycles
+        cache[s_canon] = _rebuild_tree(cost_of[s_canon][1], eclass_of, cost_of, cache)
+        new_src.append(cache[s_canon])
+    else:
+      new_src.append(_rebuild_tree(s, eclass_of, cost_of, cache))
+  new_src_tuple = tuple(new_src)
+  return u if new_src_tuple == u.src else UOp(u.op, u.dtype, new_src_tuple, u.arg, u.tag)
+
+# *** graph-level rewrite: drop-in replacement for graph_rewrite when EGRAPH is set ***
+
+def egraph_rewrite(sink:UOp, sym_pm:PatternMatcher, extra_pm:PatternMatcher|None=None, ctx=None, name:str|None=None) -> UOp:
+  """Replace graph_rewrite(sink, sym+extra, ctx) with e-graph extraction for sym, then greedy for the rest."""
+  combined = sym_pm+extra_pm if extra_pm is not None else sym_pm
+  sink = egraph_extract(sink, combined, ctx=ctx)
+  return graph_rewrite(sink, combined, ctx=ctx, name=name)

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, ...]=(), inline=False) -> UOp:
+  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=()) -> 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, inline))
+    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata))
   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,10 +848,9 @@ 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, self.inline))
-  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata}, {self.inline})"
+  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})"
 
 # ******** ops in python ********