Merge branch 'master' into sym_work

* work

* 2

* better color

.github/workflows/autogen.yml

@@ -32,6 +32,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen'
         opencl: 'true'
         amd: 'true'
         cuda: 'true'
@@ -81,6 +82,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen-mac'
         llvm: 'true'
     - name: Regenerate autogen files
       run: |
@@ -110,6 +112,8 @@ jobs:
       uses: actions/checkout@v4
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
+      with:
+        key: 'autogen-comgr'
     - name: Install autogen support packages
       run: |
         wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null

.github/workflows/benchmark.yml

@@ -520,9 +520,8 @@ jobs:
       run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run full CIFAR training steps w 6 GPUS
       run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    # this needs to be mocked and testable on a local machine
-    #- name: Test full tinyfs load
-    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
+    - name: Test full tinyfs load
+      run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 

examples/mlperf/model_train.py

@@ -1335,6 +1335,9 @@ def train_llama3():
   model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
   # vocab_size from the mixtral tokenizer
   if not SMALL: model_params |= {"vocab_size": 32000}
+  real_vocab_size = model_params['vocab_size']
+  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
+  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
   if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
   print(f"model parameters: {model_params}")
 
@@ -1352,6 +1355,8 @@ def train_llama3():
     for v in get_parameters(model):
       v.shard_(device, axis=None)
 
+    vocab_mask.shard_(device, axis=None)
+
   if (MP := getenv("MP", 1)) > 1:
     device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
     for k,v in get_state_dict(model).items():
@@ -1359,6 +1364,7 @@ def train_llama3():
       elif '.attention.wq' in k: v.shard_(device, axis=0)
       elif '.attention.wk' in k: v.shard_(device, axis=0)
       elif '.attention.wv' in k: v.shard_(device, axis=0)
+      elif '.attention.wqkv' in k: v.shard_(device, axis=0)
       elif '.attention.wo' in k: v.shard_(device, axis=1)
       elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
       elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
@@ -1371,6 +1377,8 @@ def train_llama3():
       # prevents memory spike on device 0
       v.realize()
 
+    vocab_mask.shard_(device, axis=2).realize()
+
   optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
   optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
                            eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
@@ -1401,7 +1409,7 @@ def train_llama3():
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
       tokens = tokens.shard(device)
     logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
     loss.backward()
     assert all(p.grad is g for p,g in zip(optim.params, grads))
     Tensor.realize(loss, *grads)
@@ -1431,7 +1439,7 @@ def train_llama3():
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
       tokens = tokens.shard(device)
     logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
     return loss.flatten().float().to("CPU")
 
   # ** data iters **
@@ -1469,29 +1477,28 @@ def train_llama3():
       st = time.perf_counter()
 
       stopped = False
+      losses, data_time, dev_time = [], 0, 0
       for _ in range(grad_acc):
         ist = time.perf_counter()
         try: tokens = next(train_iter)
         except StopIteration:
           stopped = True
           break
-        dt = time.perf_counter()
-        loss = minibatch(tokens)
+        mst = time.perf_counter()
+        data_time += mst - ist
+        losses.append(minibatch(tokens).item())
+        dev_time += time.perf_counter() - mst
       if stopped: break
 
       gt = time.perf_counter()
-      lr = optim_step()
-      ot = time.perf_counter()
-
-      loss = loss.float().item()
-      lr = lr.item()
-
+      lr = optim_step().item()
       et = time.perf_counter()
+
+      loss = sum(losses) / len(losses)
+      optim_time = et - gt
+      dev_time += optim_time
       step_time = et - st
       gbs_time = gt - st
-      optim_time = ot - gt
-      data_time = dt - ist
-      dev_time = step_time - data_time * grad_acc
       if BENCHMARK: step_times.append(step_time)
 
       i += 1

examples/mlperf/optim.py

@@ -21,12 +21,13 @@ class GradAccClipAdamW(Optimizer):
       total_norm = grads[0].float().square().sum().sqrt()
       grads[0] = (grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype)
     else:
+      for i in range(len(grads)):
+        grads[i] = grads[i] / self.grad_acc
       total_norm = Tensor.zeros((), dtype=dtypes.float32, device=self.device)
       for g in grads:
         total_norm += g.float().square().sum()
       total_norm = total_norm.sqrt()
       for i in range(len(grads)):
-        grads[i] = grads[i] / self.grad_acc
         grads[i] = (grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)
 
     ret = []

extra/gemm/asm/cdna/gemm.py

@@ -11,12 +11,12 @@ from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
 WORKGROUP_SIZE = 256
 
 @functools.cache
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
   batch, M, K = A.shape
   K2, N = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2
   lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
-  gidx = UOp.special(wg, "gidx0")
+  gidx = UOp.special(NUM_WG, "gidx0")
   insts = build_kernel(batch, M, N, K, A.dtype.base)
   lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
   sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
@@ -94,7 +94,7 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
   renderer = Device[a.device[0] if is_multi else a.device].renderer
   dname, arch = renderer.device, getattr(renderer, "arch", "")
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=NUM_WG, arch=arch), grad_fxn=custom_gemm_bw)[0]
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), 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)

extra/models/llama.py

@@ -56,7 +56,10 @@ class Attention:
   def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]=None) -> Tensor:
     if getenv("WQKV"):
       xqkv = self.wqkv(x)
-      xq, xk, xv = xqkv.split([self.n_heads * self.head_dim, self.n_kv_heads * self.head_dim, self.n_kv_heads * self.head_dim], dim=2)
+      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
     else:
       xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)
 

test/external/external_test_hive_reset.py

@@ -2,8 +2,9 @@
 import subprocess, sys
 from tinygrad.helpers import getenv
 
-LOOPS = getenv("LOOPS", 10)
+LOOPS = getenv("LOOPS", 50)
 BROKEN = getenv("BROKEN", 0)
+ONLY_RESET = getenv("ONLY_RESET", 0)
 
 BROKEN_KERNEL_SCRIPT = """
 from tinygrad.device import Device
@@ -36,7 +37,7 @@ for i in range(LOOPS):
     print(f"=== Running broken kernel ({i+1}/{LOOPS}) ===")
     ret = subprocess.run([sys.executable, "-c", BROKEN_KERNEL_SCRIPT])
     print(f"=== broken kernel exited with code {ret.returncode} ===")
-
-  print(f"=== Running test_tiny.py ({i+1}/{LOOPS}) ===")
-  ret = subprocess.run([sys.executable, "test/test_tiny.py", "TestTiny.test_plus"])
-  print(f"=== test_tiny.py exited with code {ret.returncode} ===")
+  elif not ONLY_RESET:
+    print(f"=== Running test_tiny.py ({i+1}/{LOOPS}) ===")
+    ret = subprocess.run([sys.executable, "test/test_tiny.py", "TestTiny.test_plus"])
+    print(f"=== test_tiny.py exited with code {ret.returncode} ===")

test/null/test_graph_rewrite.py

@@ -349,5 +349,184 @@ class TestStopEarly(unittest.TestCase):
     ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
     assert ret == cn+d
 
+class TestWalkRewrite(unittest.TestCase):
+  """Tests for graph_rewrite with walk=True (MLIR Walk Pattern Rewrite Driver semantics).
+  walk=True gives a single-pass traversal that does NOT revisit or re-traverse into rewritten subtrees.
+  Supports both top-down (default) and bottom-up (bottom_up=True) modes."""
+
+  # *** top-down walk (default): process children first, then try pm on rebuilt node ***
+
+  def test_walk_topdown_simple_substitute(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    ret = graph_rewrite(a + 4, _substitute, {a:b}, walk=True)
+    self.assertIs(ret, b+4)
+
+  def test_walk_topdown_does_not_traverse_into_replacement(self):
+    """Top-down walk: replacement subtrees are NOT re-entered."""
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    c = UOp.variable('c', 0, 10)
+    d = UOp.variable('d', 0, 10)
+    # a is replaced by b+c, but b inside the replacement is NOT further substituted to d
+    ret_walk = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, walk=True)
+    self.assertIs(ret_walk, (b+c)+4)
+    # contrast: greedy bottom_up WOULD replace b inside the replacement
+    ret_greedy = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, bottom_up=True)
+    self.assertIs(ret_greedy, (d+c)+4)
+
+  def test_walk_topdown_no_fixed_point(self):
+    """A bouncing pattern applies once and stops instead of looping."""
+    a = UOp.const(dtypes.int, 3)
+    pm = PatternMatcher([
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+    ])
+    with self.assertRaises(RuntimeError):
+      graph_rewrite(a, pm, bottom_up=True)
+    ret = graph_rewrite(a, pm, walk=True)
+    self.assertIs(ret, UOp.const(dtypes.int, 4))
+
+  def test_walk_topdown_rewrites_children(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    c = UOp.variable('c', 0, 10)
+    ret = graph_rewrite((a + 4) + (b + 5), _substitute, {a:c, b:c}, walk=True)
+    self.assertIs(ret, (c + 4) + (c + 5))
+
+  def test_walk_topdown_diamond(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    ret = graph_rewrite((a + 4) + (a + 5), _substitute, {a:b}, walk=True)
+    self.assertIs(ret, (b + 4) + (b + 5))
+
+  def test_walk_topdown_children_rewritten_before_parent(self):
+    """Top-down walk processes children first: child substitution changes the rebuilt parent."""
+    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
+    n1 = a.sin()          # sin(a)
+    ret = n1.sin()         # sin(sin(a))
+    # sin(a)->sqrt(a) fires first (child), parent rebuilds to sin(sqrt(a)), which doesn't match sin(sin(a)) in dvars
+    ret_walk = graph_rewrite(ret, _substitute, {a.sin():a.sqrt(), n1.sin():n1.sqrt()}, walk=True)
+    self.assertIs(ret_walk, a.sqrt().sin())
+
+  def test_walk_topdown_self_referential_replacement(self):
+    """Replacement containing the replaced node works without infinite recursion."""
+    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
+    ret = graph_rewrite(a.sin() + 4, _substitute, {a.sin(): a.sin().sqrt()}, walk=True)
+    self.assertIs(ret, a.sin().sqrt() + 4)
+
+  def test_walk_topdown_visit_order(self):
+    """Top-down walk fires pm after children are processed (post-order)."""
+    visited = []
+    def track_visit(ctx, x):
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
+      return None
+    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    graph_rewrite(a + b, pm, ctx=visited, walk=True)
+    self.assertEqual(visited, [1, 2, Ops.ADD])
+
+  # *** bottom-up walk: try bpm on node first, skip children if it matches ***
+
+  def test_walk_bottomup_simple_substitute(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    ret = graph_rewrite(a + 4, _substitute, {a:b}, bottom_up=True, walk=True)
+    self.assertIs(ret, b+4)
+
+  def test_walk_bottomup_does_not_traverse_into_replacement(self):
+    """Bottom-up walk: replacement subtrees are NOT entered."""
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    c = UOp.variable('c', 0, 10)
+    d = UOp.variable('d', 0, 10)
+    ret = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, bottom_up=True, walk=True)
+    self.assertIs(ret, (b+c)+4)
+
+  def test_walk_bottomup_parent_match_skips_children(self):
+    """Bottom-up walk matches parent first: if it matches, children are never visited."""
+    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
+    n1 = a.sin()
+    ret = n1.sin()         # sin(sin(a))
+    # sin(sin(a)) matches n1.sin()->n1.sqrt() immediately, children never visited, sin(a) inside replacement untouched
+    ret_walk = graph_rewrite(ret, _substitute, {a.sin():a.sqrt(), n1.sin():n1.sqrt()}, bottom_up=True, walk=True)
+    self.assertIs(ret_walk, a.sin().sqrt())
+
+  def test_walk_bottomup_no_fixed_point(self):
+    """Bottom-up walk also applies once per node, no fixed-point iteration."""
+    a = UOp.const(dtypes.int, 3)
+    pm = PatternMatcher([
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+    ])
+    ret = graph_rewrite(a, pm, bottom_up=True, walk=True)
+    self.assertIs(ret, UOp.const(dtypes.int, 4))
+
+  def test_walk_bottomup_visit_order(self):
+    """Bottom-up walk fires bpm before descending (pre-order)."""
+    visited = []
+    def track_visit(ctx, x):
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
+      return None
+    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    graph_rewrite(a + b, pm, ctx=visited, bottom_up=True, walk=True)
+    # bpm fires on each node before children: +, 1, 2
+    self.assertEqual(visited, [Ops.ADD, 1, 2])
+
+  def test_walk_bottomup_unmatched_falls_through_to_children(self):
+    """Bottom-up walk: if bpm doesn't match a node, its children are still processed."""
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    c = UOp.variable('c', 0, 10)
+    # only a is in dvars, not a+4. bpm won't match a+4, so it descends and finds a.
+    ret = graph_rewrite((a + 4) + (b + 5), _substitute, {a:c, b:c}, bottom_up=True, walk=True)
+    self.assertIs(ret, (c + 4) + (c + 5))
+
+  # *** bidirectional walk: bpm fires before children, pm fires after rebuild ***
+
+  def test_walk_bidirectional_visit_order(self):
+    """Bidirectional walk: bpm fires pre-order, pm fires post-order."""
+    visited = []
+    def bpm_visit(ctx, x):
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
+      return None
+    def pm_visit(ctx, x):
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
+      return None
+    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)])
+    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)])
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
+    # bpm fires pre-order, pm fires post-order
+    self.assertEqual(visited, [
+      (Ops.ADD, "bpm"), (1, "bpm"), (1, "pm"), (2, "bpm"), (2, "pm"), (Ops.ADD, "pm"),
+    ])
+
+  def test_walk_bidirectional_bpm_short_circuits(self):
+    """If bpm matches, children are skipped and pm never fires on that node."""
+    visited = []
+    def bpm_match(ctx, x):
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
+      # rewrite const(1) -> const(10), short-circuiting its subtree
+      if x.op is Ops.CONST and x.arg == 1: return x.replace(arg=10)
+      return None
+    def pm_match(ctx, x):
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
+      return None
+    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)])
+    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)])
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    ret = graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
+    # bpm matches const(1) and short-circuits it, so pm never fires on const(1)
+    self.assertNotIn((1, "pm"), visited)
+    # but pm still fires on const(2) and the rebuilt ADD
+    self.assertIn((2, "pm"), visited)
+    self.assertIs(ret, UOp.const(dtypes.int, 10) + b)
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_viz.py

@@ -282,9 +282,10 @@ class TestVizIntegration(BaseTestViz):
     ast = Tensor.schedule(Tensor.empty(4)+Tensor.empty(4))[0].ast
     prg = get_program(ast, Device[Device.DEFAULT].renderer)
     lst = get_viz_list()
-    self.assertEqual(len(lst), 2)
-    self.assertEqual(lst[0]["name"], "Schedule 1 Kernel n1")
-    self.assertEqual(lst[1]["name"], prg.name)
+    self.assertEqual(len(lst), 3)
+    self.assertEqual(lst[0]["name"], "Process 1 Buffer n1")
+    self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1")
+    self.assertEqual(lst[2]["name"], prg.name)
 
   # schedule graph CALL nodes have a link to jump to codegen
   def test_link_sched_codegen(self):
@@ -293,8 +294,9 @@ class TestVizIntegration(BaseTestViz):
     sched = Tensor.schedule(c1, c2)
     prgs = [si.lower().prg.p.name for si in sched]
     lst = get_viz_list()
-    viz_kernel = next(i for i,s in enumerate(lst[0]["steps"]) if s["name"] == "View Kernel Graph")
-    graph = next(get_viz_details(0, viz_kernel))["graph"]
+    sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule"))
+    viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph")
+    graph = next(get_viz_details(sched_idx, viz_kernel))["graph"]
     call_nodes = [n for n in graph.values() if n["label"].startswith("CALL")]
     for i,n in enumerate(call_nodes):
       assert n["ref"] is not None

test/unit/test_assign.py

@@ -269,6 +269,16 @@ class TestAssign(unittest.TestCase):
     out = attn.cache_k.flatten().numpy()
     np.testing.assert_allclose(out, [1.,1.,1.,1.,1.,1.,0.,0.,1.,1.,1.,1.,1.,1.,0.,0.])
 
+  def test_assign_after(self):
+    t = Tensor.zeros(10).contiguous().realize()
+    t.uop = t.uop.after(t.uop.assign((t+1).uop))
+    np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,1.,1.,1.,1.,1.])
+
+  def test_assign_after_partial(self):
+    t = Tensor.zeros(10).contiguous().realize()
+    t.uop = t.uop.after(t[:5].uop.assign(Tensor.ones(5).uop))
+    np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,0.,0.,0.,0.,0.])
+
   def test_assign_contiguous(self):
     b = Tensor.arange(16).reshape(4,4).contiguous().realize()
     a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
@@ -485,10 +495,10 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
 
   def test_assign_bitcast_different_size(self):
-    # different-size bitcast creates a new tensor, not a view, so assign doesn't modify the original
+    # assign to a shape-changing bitcast view (only works on DISK currently)
     a = Tensor([0]*8, dtype=dtypes.uint8).realize()
     a.bitcast(dtypes.int64).assign(Tensor([12345], dtype=dtypes.int64)).realize()
-    np.testing.assert_equal(a.numpy(), [0]*8)
+    np.testing.assert_equal(a.numpy(), [0]*8)  # TODO: should be [57, 48, 0, 0, 0, 0, 0, 0] (little-endian 12345)
 
   @unittest.skip("don't use output buffer, and mismatch dtype no longer supported")
   def test_cast_assignment(self):

test/unit/test_callify.py

@@ -0,0 +1,111 @@
+import unittest
+from tinygrad import Tensor, dtypes
+
+class TestCallify(unittest.TestCase):
+  def test_basic(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    out = a + b
+    out.callify()
+    self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0])
+
+  def test_const(self):
+    out = Tensor(2.0) + Tensor(3.0)
+    out.callify()
+    self.assertEqual(out.item(), 5.0)
+
+  def test_sum(self):
+    out = Tensor.ones(16).contiguous().sum()
+    out.callify()
+    self.assertEqual(out.item(), 16.0)
+
+  def test_multi_output(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    c = a + b
+    d = a * b
+    c.callify(d)
+    self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
+    self.assertListEqual(d.tolist(), [4.0, 10.0, 18.0])
+
+  def test_two_callify_independent(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    c = a + b
+    c.callify()
+
+    d = Tensor([10.,20,30])
+    e = Tensor([1.,1,1])
+    f = d - e
+    f.callify()
+
+    self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
+    self.assertListEqual(f.tolist(), [9.0, 19.0, 29.0])
+
+  def test_two_callify_shared_input(self):
+    a = Tensor([1.,2,3]).contiguous().realize()
+    b = a + 1
+    b.callify()
+    c = a * 2
+    c.callify()
+    self.assertListEqual(b.tolist(), [2.0, 3.0, 4.0])
+    self.assertListEqual(c.tolist(), [2.0, 4.0, 6.0])
+
+  def test_chained_callify(self):
+    a = Tensor([1.,2,3])
+    b = a + 1
+    b.callify()
+    b.realize()
+    c = b + 1
+    c.callify()
+    self.assertListEqual(c.tolist(), [3.0, 4.0, 5.0])
+
+  def test_gemm(self):
+    a = Tensor.ones(8, 8).contiguous()
+    b = Tensor.eye(8).contiguous()
+    out = a @ b
+    out.callify()
+    lst = out.tolist()
+    for y in range(8):
+      for x in range(8):
+        self.assertEqual(lst[y][x], 1.0)
+
+  def test_int_dtype(self):
+    a = Tensor([1,2,3], dtype=dtypes.int)
+    b = Tensor([4,5,6], dtype=dtypes.int)
+    out = a + b
+    out.callify()
+    self.assertListEqual(out.tolist(), [5, 7, 9])
+
+  def test_reduce(self):
+    out = Tensor([1.,2,3,4]).sum()
+    out.callify()
+    self.assertEqual(out.item(), 10.0)
+
+  def test_multiple_ops(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    out = (a + b) * (a - b)
+    out.callify()
+    self.assertListEqual(out.tolist(), [-15.0, -21.0, -27.0])
+
+  def test_double_callify(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    out = a + b
+    out.callify()
+    out.callify()
+    self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0])
+
+  def test_double_callify_multi_output(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    c = a + b
+    d = a * b
+    c.callify(d)
+    c.callify(d)
+    self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
+    self.assertListEqual(d.tolist(), [4.0, 10.0, 18.0])
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_disk_tensor.py

@@ -74,13 +74,13 @@ class TestRawDiskBuffer(unittest.TestCase):
     _test_bitcasted(t, dtypes.float32, 0.0)
     _test_bitcasted(t, dtypes.uint32, 0)
     # pi in float16 stored via int16
-    t.assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16).bitcast(dtypes.uint8)).realize()
+    t.bitcast(dtypes.uint16).assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16)).realize()
     _test_bitcasted(t, dtypes.float16, 3.140625)
     _test_bitcasted(t, dtypes.float32, 50.064727)
     _test_bitcasted(t, dtypes.uint16, 0x4248)
     _test_bitcasted(t, dtypes.uint32, 0x42484248)
     # pi in float32 stored via float32
-    t.assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32).bitcast(dtypes.uint8)).realize()
+    t.bitcast(dtypes.float32).assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32)).realize()
     _test_bitcasted(t, dtypes.float32, 3.1415927)
     _test_bitcasted(t, dtypes.uint32, 0x40490FDB)
     # doesn't suport normal cast
@@ -178,6 +178,13 @@ class TestSafetensors(TempDirTestCase):
     import json
     assert json.loads(dat[8:8+sz])['__metadata__']['hello'] == 'world'
 
+  def test_safe_save_only_copy(self):
+    from tinygrad.helpers import GlobalCounters
+    t = Tensor.rand(10, 10).realize()
+    GlobalCounters.reset()
+    safe_save({"t": t}, self.tmp("test_copy.safetensors"))
+    assert GlobalCounters.global_ops == 0, f"safe_save should have no compute, got {GlobalCounters.global_ops} ops"
+
   def test_save_all_dtypes(self):
     for dtype in dedup(DTYPES_DICT.values()):
       if dtype in [dtypes.bfloat16]: continue # not supported in numpy
@@ -357,15 +364,10 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_assign_with_bitcast(self):
     # bitcast assign is used in safe_save for writing header length
-    # bitcast on source side works, bitcast on target side raises
     t = Tensor.empty(16, device=f"disk:{self.tmp('dt_assign_bitcast')}", dtype=dtypes.uint8)
-    # correct way: bitcast the source to match target dtype
-    t[0:8].assign(Tensor([12345], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
+    t[0:8].bitcast(dtypes.int64).assign([12345])
     val = int.from_bytes(t[0:8].data(), 'little')
     self.assertEqual(val, 12345)
-    # bitcast on target with non-broadcastable dtype raises
-    with self.assertRaises(RuntimeError):
-      t[0:4].bitcast(dtypes.int32).assign(Tensor([12345], dtype=dtypes.int64))
 
   def test_assign_to_bitcast_view(self):
     # assign float values to a float32 view of a uint8 disk buffer (used by safe_save)

test/unit/test_function.py

@@ -0,0 +1,197 @@
+import numpy as np
+import unittest
+from tinygrad.function import function
+from tinygrad import Tensor
+from tinygrad.uop.ops import UOp
+
+class TestFunction(unittest.TestCase):
+  def test_simple(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
+
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    np.testing.assert_equal(f(a,b).numpy(), [5,7,9])
+
+  def test_simple_same(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
+
+    a = Tensor([1,2,3])
+    np.testing.assert_equal(f(a,a).numpy(), [2,4,6])
+
+  def test_implicit(self):
+    inp = Tensor([7,8,9])
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
+
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    np.testing.assert_equal(f(a,b).numpy(), [12,15,18])
+
+  def test_implicit_same_as_input(self):
+    inp = Tensor([7,8,9])
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
+
+    a = Tensor([1,2,3])
+    np.testing.assert_equal(f(a, inp).numpy(), [15,18,21])
+
+  def test_implicit_2(self):
+    inp = Tensor([7,8,9])
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor:
+      return a+b+inp
+    inp2 = Tensor([7,8,10])
+    @function
+    def g(a:Tensor, b:Tensor) -> Tensor:
+      return a+b+inp2
+
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    c = f(a,b)
+    d = g(a,b)
+    c.realize(d)
+    np.testing.assert_equal(c.numpy(), [12,15,18])
+    np.testing.assert_equal(d.numpy(), [12,15,19])
+
+  def test_implicit_unrealized(self):
+    inp = Tensor([1,2,3]) + Tensor([4,5,6])
+    @function
+    def f(a:Tensor) -> Tensor: return a + inp
+
+    np.testing.assert_equal(f(Tensor([10,20,30])).numpy(), [15,27,39])
+
+  def test_detach(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a.detach() + b
+
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
+
+  def test_method(self):
+    class Foo:
+      def __init__(self): self.w = Tensor([10,20,30])
+      @function
+      def __call__(self, x:Tensor) -> Tensor: return x + self.w
+
+    foo = Foo()
+    np.testing.assert_equal(foo(Tensor([1,2,3])).numpy(), [11,22,33])
+
+  def test_grad_gemm(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a @ b
+
+    a = Tensor([[1.,2.],[3.,4.]], requires_grad=True)
+    b = Tensor([[5.,6.],[7.,8.]], requires_grad=True)
+    (f(a, b).contiguous() * b).sum().backward()
+    Tensor.realize(a, b, a.grad, b.grad)
+    # L = sum((a@b) * b), dL/d(a@b) = b, dL/da = b @ b^T, dL/db = a^T @ b + (a@b)
+    na, nb = a.numpy(), b.numpy()
+    np.testing.assert_allclose(a.grad.numpy(), nb @ nb.T)
+    np.testing.assert_allclose(b.grad.numpy(), na.T @ nb + na @ nb)
+
+  def test_grad_implicit(self):
+    w = Tensor([1., 2., 3.], requires_grad=True)
+    w.realize() # TODO: this is required
+    @function
+    def f(x:Tensor) -> Tensor: return x * w
+
+    x = Tensor([4., 5., 6.])
+    f(x).sum().backward()
+    np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6.])
+
+  def test_symbolic_index(self):
+    table = Tensor([10,20,30,40]).contiguous().realize()
+    @function
+    def f(x:Tensor, start_pos:int|UOp) -> Tensor:
+      return x + table[start_pos]
+
+    v = UOp.variable("start_pos", 0, 3)
+    np.testing.assert_equal(f(Tensor([1,2,3]), v.bind(0)).numpy(), [11,12,13])
+
+  def test_symbolic_shape_input(self):
+    table = Tensor([10,20,30,40]).contiguous().realize()
+    @function
+    def f(x:Tensor) -> Tensor: return x * 2
+    sz = UOp.variable("sz", 1, 3)
+    slic = table[:sz.bind(2)]
+    np.testing.assert_equal(f(slic)[:2].numpy(), [20,40])
+
+  def test_nested_calls(self):
+    w = Tensor([10., 20., 30.])
+    @function
+    def f(a:Tensor) -> Tensor: return a + w
+    @function
+    def g(a:Tensor) -> Tensor: return a * w
+
+    a = Tensor([1., 2., 3.])
+    np.testing.assert_allclose(g(f(a)).numpy(), [110., 440., 990.])
+
+  def test_name(self):
+    @function
+    def f(a:Tensor) -> Tensor: return a + 1
+    assert f(Tensor([1])).uop.arg.name.endswith("f")
+
+  def test_method_name(self):
+    class Foo:
+      @function
+      def __call__(self, x:Tensor) -> Tensor: return x + 1
+    assert Foo()(Tensor([1])).uop.arg.name.endswith("Foo.__call__")
+
+  def test_callable_instance(self):
+    class Foo:
+      def __init__(self): self.w = Tensor([10,20,30])
+      def __call__(self, x:Tensor) -> Tensor: return x + self.w
+    foo = Foo()
+    f = function(foo)
+    np.testing.assert_equal(f(Tensor([1,2,3])).numpy(), [11,22,33])
+    assert f(Tensor([1,2,3])).uop.arg.name.endswith("Foo")
+
+  def test_iadd(self):
+    @function
+    def f(x:Tensor) -> Tensor:
+      x += 1
+      return x
+
+    a = Tensor([1,2,3]).realize()
+    np.testing.assert_equal(f(a).numpy(), [2,3,4])
+    np.testing.assert_equal(a.numpy(), [3,4,5])  # TODO: should be [1,2,3]
+
+  def test_implicit_assign(self):
+    a = Tensor([1,2,3])
+    a += 1
+    c = Tensor([2,2,2]).contiguous()
+    @function
+    def f(b:Tensor) -> Tensor: return a+b+c
+    b = Tensor([10,20,30]).realize()
+    np.testing.assert_equal(f(b).numpy(), [14,25,36])
+
+  def test_assign_input(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor:
+      a.assign(b+1)
+      return a
+
+    a = Tensor([1,2,3]).realize()
+    b = Tensor([10,20,30]).realize()
+    np.testing.assert_equal(f(a,b).numpy(), [11,21,31])
+    np.testing.assert_equal(a.numpy(), [11,21,31])  # TODO: should be [1,2,3]
+    np.testing.assert_equal(b.numpy(), [10,20,30])
+
+  @unittest.expectedFailure
+  def test_assign_slice(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor:
+      a[1:] = b[1:]+1
+      return a
+
+    a = Tensor([1,2,3]).realize()
+    b = Tensor([10,20,30]).realize()
+    np.testing.assert_equal(f(a,b).numpy(), [1,21,31])
+    np.testing.assert_equal(a.numpy(), [1,2,3])
+    np.testing.assert_equal(b.numpy(), [10,20,30])
+
+if __name__ == '__main__':
+  unittest.main()

test/unit/test_tinyfs.py

@@ -45,37 +45,31 @@ class TestTinyFS(unittest.TestCase):
     cls._server.shutdown()
     cls._server.server_close()
 
-  @unittest.expectedFailure
   def test_store(self):
     h = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
     self.assertEqual(h.shape, (16,))
     self.assertEqual(h.dtype, dtypes.uint8)
 
-  @unittest.expectedFailure
   def test_store_deterministic(self):
     a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
     b = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
     np.testing.assert_array_equal(a.numpy(), b.numpy())
 
-  @unittest.expectedFailure
   def test_store_different_data(self):
     a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
     b = Tensor([5.0, 6.0, 7.0, 8.0]).fs_store().realize()
     self.assertNotEqual(a.tolist(), b.tolist())
 
-  @unittest.expectedFailure
   def test_roundtrip_uint8(self):
     arr = np.arange(256, dtype=np.uint8)
-    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr))
+    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
     np.testing.assert_array_equal(loaded.numpy(), arr)
 
-  @unittest.expectedFailure
   def test_roundtrip_multichunk_uint8(self):
     arr = np.random.default_rng(42).integers(0, 256, size=Tensor.CHUNK_SIZE + 1024, dtype=np.uint8)
-    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr))
+    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
     np.testing.assert_array_equal(loaded.numpy(), arr)
 
-  @unittest.expectedFailure
   def test_hash_matches_python_impl(self):
     arr = np.arange(256, dtype=np.uint8)
     h = Tensor(arr).fs_store().realize()

tinygrad/__init__.py

@@ -4,6 +4,7 @@ if int(os.getenv("TYPED", "0")):
   install_import_hook(__name__)
 from tinygrad.tensor import Tensor                                    # noqa: F401
 from tinygrad.engine.jit import TinyJit                               # noqa: F401
+from tinygrad.function import function                                # noqa: F401
 from tinygrad.uop.ops import UOp
 Variable = UOp.variable
 from tinygrad.dtype import dtypes                                     # noqa: F401

tinygrad/apps/llm.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import sys, argparse, typing, re, unicodedata, json, uuid, time, functools
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
 from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
 
@@ -116,6 +116,7 @@ class TransformerBlock:
       self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
       self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)
 
+  @function
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     x_norm = self.attn_norm(x)                       # (B,T,D)
     q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
@@ -131,11 +132,15 @@ class TransformerBlock:
     q = apply_rope(q, freqs_cis)
     k = apply_rope(k, freqs_cis)
 
-    if not hasattr(self, "cache_kv"):
-      self.cache_kv = Tensor.zeros(2, B, self.n_kv_heads, self.max_context, self.head_dim, dtype=k.dtype, device=k.device).contiguous().realize()
-    self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
-    k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-    v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+    # TODO: fix assign to behave like this
+    assigned_kv = self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.assign(Tensor.stack(k, v).contiguous().uop))
+    tensor_assigned_kv = Tensor(assigned_kv, device=assigned_kv.device)
+    k = tensor_assigned_kv[0, :, :, 0:start_pos+T, :]
+    v = tensor_assigned_kv[1, :, :, 0:start_pos+T, :]
+
+    #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
+    #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
+    #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
 
     # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
     mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(int(start_pos)+1) if T > 1 else None
@@ -144,6 +149,7 @@ class TransformerBlock:
     attn = self.attn_output(attn)
     return x + attn
 
+  @function
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
     if hasattr(self, 'ffn_gate_exps'):
@@ -156,6 +162,9 @@ class TransformerBlock:
     return h + self.ffn_down(gated)
 
   def __call__(self, x: Tensor, start_pos: int|UOp):
+    if not hasattr(self, "cache_kv"):
+      # TODO: how is the dtype of this determined?
+      self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).contiguous().realize()
     return self._feed_forward(self._attention(x, start_pos)).contiguous()
 
 class Transformer:

tinygrad/engine/allocations.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
-from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, identity_element, profile_matches
+from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, identity_element, track_rewrites
 from tinygrad.dtype import ImageDType
-from tinygrad.helpers import prod, DEBUG, argsort, VIZ
+from tinygrad.helpers import prod, DEBUG, argsort, VIZ, pluralize
 
 @dataclass
 class AllocCtx:
@@ -18,10 +18,10 @@ def tag_uop(ctx:AllocCtx, x:UOp):
 
 def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
   # copies to disk are replaced with the disk buffer
-  to_disk = isinstance(u._device, str) and u._device.startswith("DISK")
+  to_disk = isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS"))
   if to_disk: ctx.buffer_map[u] = UOp.new_buffer(u.device, u.shard_size, u.dtype).reshape(u.max_shard_shape)
   # all copies from disk/numpy are realized into a real buffer
-  from_creation = isinstance(u.src[0]._device, str) and any(u.src[0]._device.startswith(x) for x in ["NPY", "DISK", "PYTHON"])
+  from_creation = isinstance(u.src[0]._device, str) and any(u.src[0]._device.startswith(x) for x in ["NPY", "DISK", "PYTHON", "TINYFS"])
   if from_creation: return tag_uop(ctx, u)
 
 def apply_after(ctx:AllocCtx, u:UOp):
@@ -41,8 +41,8 @@ add_tags = PatternMatcher([
 def replace_contig_with_assign(u:UOp):
   # if size is 0, remove the contig
   if u.size == 0: return u.src[0]
-  # no real contig for DISK tensors, they are left alone
-  if isinstance(u._device, str) and u._device.startswith("DISK"): return u.rtag(None)
+  # no real contig for DISK/TINYFS tensors, they are left alone
+  if isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS")): return u.rtag(None)
   dtype = u.dtype
   if isinstance(dtype, ImageDType):
     if prod(dtype.shape) != prod(u.max_shard_shape) or ([x for x in u.max_shard_shape if x != 1] or [1])[-1] % 4 != 0:
@@ -113,7 +113,7 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
 pm_finalize_call = PatternMatcher([
   (UPat(Ops.ASSIGN, name="x"), untag_and_append),
   (UPat(Ops.AFTER, name="x"), append_after),
-  (UPat(Ops.COPY, name="x"), lambda ctx,x: append_after(ctx,x) if isinstance(x.device, str) and x.device.startswith("DISK") else None),
+  (UPat(Ops.COPY, name="x"), lambda ctx,x: append_after(ctx,x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
   # replace UNIQUE with LUNIQUE for CONST cache key normalization
   (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(d,))),
 ])
@@ -125,8 +125,9 @@ pm_replace_buf = PatternMatcher([
   (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), replace_input_buffer),
 ])
 
-@profile_matches
+@track_rewrites(lambda _,ret: f"Process {pluralize('Buffer', len(ret[1]))}")
 def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
+  if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
   # uop list is a list in the original_sink graph and we can map to the tags later
   # here we build buffer map
   dont_realize = {Ops.CONST, Ops.BUFFER, Ops.BIND, Ops.DEFINE_VAR, Ops.AFTER}

tinygrad/engine/schedule.py

@@ -1,12 +1,12 @@
 import time, inspect
 from typing import cast
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink
+from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo
+from tinygrad.uop.ops import _remove_all_tags
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR
 from tinygrad.engine.realize import ExecItem
-from tinygrad.engine.allocations import transform_to_call
 
 # **** schedule linearizer
 
@@ -23,7 +23,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
     for u in sched_sink.toposort(gate_kernel_sink):
       if u.op is not Ops.AFTER: continue
       k = u.src[1]
-      assert k.op in {Ops.CALL, Ops.END}, f"AFTER src[1] should be KERNEL or END, not {k.op}"
+      assert k.op in {Ops.CALL, Ops.END, Ops.LINEAR}, f"AFTER src[1] should be CALL or END, not {k.op}"
       in_degree.setdefault(k, 0)
       if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
       # WAR deps from rangeify are stored in AFTER src[2:]
@@ -50,61 +50,20 @@ def create_schedule(sched_sink:UOp) -> UOp:
     linearized: list[UOp] = []
     while len(queue):
       rk = queue.popleft()
-      k = rk.src[0] if rk.op is Ops.END else rk
-      assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
-      buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
-      linearized.append(k.src[0].call(*buf_uops, metadata=k.arg.metadata))
+      if rk.op is Ops.LINEAR:
+        linearized.extend(rk.src)
+      else:
+        k = rk.src[0] if rk.op is Ops.END else rk
+        assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
+        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
+        linearized.append(k.src[0].call(*buf_uops, metadata=k.arg.metadata))
       for x in children.get(rk, []):
         in_degree[x] -= 1
         if in_degree[x] == 0: queue.append(x)
   return UOp(Ops.LINEAR, src=tuple(linearized))
 
-from tinygrad.engine.memory import memory_planner
-from tinygrad.schedule.rangeify import get_kernel_graph
-from tinygrad.uop.ops import PatternMatcher, UPat
-
-def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
-  if (ret:=ctx[0].get(b, None)) is None: ctx[0][b] = ret = UOp.new_buffer(b.device, b.arg, b.dtype)
-  return ret
-
-pm_post_sched_cache = PatternMatcher([
-  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg]),
-  # create new BUFFERs for LUNIQUE BUFFERs from rangeify
-  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer),
-])
-
-schedule_cache: dict[bytes, UOp] = {}
-@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[1]))}")
-def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], list[ExecItem], dict[str, int]]:
-  # big_sink srcs are all the Tensors
-  st = time.perf_counter()
-  big_sink, buffer_map = transform_to_call(big_sink)
-  function = big_sink.src[0]
-
-  if not SCACHE or (sc_ret:=schedule_cache.get(function.key, None)) is None:
-    if SPEC: type_verify(big_sink, tensor_spec)
-    linear = create_schedule(get_kernel_graph(function))
-    if SCACHE: schedule_cache[function.key] = linear
-  else:
-    # schedule cache hit
-    linear = sc_ret
-
-  # it's a call that we late apply
-  linear = graph_rewrite(linear, pm_post_sched_cache, ctx=({}, big_sink.src[1:]), name="params to buffers")
-
-  # vars used in the schedule
-  used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
-  # get var_vals
-  var_vals: dict[str, int] = {}
-  for b in big_sink.src[1:]:
-    if b.op is Ops.BIND:
-      nm = b.src[0].expr
-      if nm not in used_vars: continue
-      val = b.src[1].arg
-      assert nm not in var_vals or var_vals[nm] == val, f"bind mismatch on {nm}, {var_vals[nm]} != {val}"
-      var_vals[nm] = val
-
-  # convert LINEAR to ExecItems
+def linear_to_schedule(linear:UOp) -> list[ExecItem]:
+  """Convert a LINEAR UOp to a list of ExecItems."""
   schedule: list[ExecItem] = []
   for si in linear.src:
     ast, buf_uops = si.src[0], si.src[1:]
@@ -121,17 +80,78 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
       for j, bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
         schedule.append(ExecItem(ast, list(bufs), metadata, {dnums[0].expr:j} if len(dnums) else {}))
     else:
-      schedule.append(ExecItem(ast, list(ubufs), metadata))
-  with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
+      schedule.append(ExecItem(ast, cast(list[Buffer|None], ubufs), metadata))
+  return schedule
 
-  if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
+from tinygrad.engine.memory import memory_planner
+from tinygrad.schedule.rangeify import get_kernel_graph
+from tinygrad.uop.ops import PatternMatcher, UPat
+
+def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
+  if (ret:=ctx[0].get(b, None)) is None: ctx[0][b] = ret = UOp.new_buffer(b.device, b.arg, b.dtype)
+  return ret
+
+pm_post_sched_cache = PatternMatcher([
+  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg].rtag() if x.tag is None else None),
+  # create new BUFFERs for LUNIQUE BUFFERs from rangeify
+  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer),
+])
+
+# the AFTER structure is already in LINEAR
+pm_collapse_after = PatternMatcher([
+  (UPat(Ops.AFTER, name="x"), lambda x: x.src[0])
+])
+
+schedule_cache: dict[bytes, UOp] = {}
+def lower_schedule_to_linear(big_sink:UOp) -> UOp|None:
+  st = time.perf_counter()
+  function = big_sink.src[0]
+  if isinstance(function.arg, KernelInfo): return None
+  if not SCACHE or (sc_ret:=schedule_cache.get(function.key, None)) is None:
+    if SPEC: type_verify(big_sink, tensor_spec)
+    # support recursive CALLs
+    function = graph_rewrite(function, pm_schedule, name="inner schedule to linear")
+    linear = create_schedule(get_kernel_graph(function))
+    if SCACHE: schedule_cache[function.key] = linear
+  else:
+    # schedule cache hit
+    linear = sc_ret
+  if (DEBUG >= 1 and len(linear.src) > 1) or DEBUG >= 3:
     for frm in inspect.stack():
+      if frm.filename == "<string>": continue
       if frm.filename.startswith(str(BASEDIR / "apps")): break
       if not frm.filename.startswith(str(BASEDIR)) and not frm.filename.endswith("/contextlib.py"): break
     else:
       frm = None
-    print(f"scheduled {len(schedule):5d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
+    print(f"scheduled {len(linear.src):5d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
           f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {function.key.hex()[:8]}"+\
           f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
+  # TODO: use walk and avoid the remove tags
+  linear = graph_rewrite(linear, pm_post_sched_cache, ctx=({}, big_sink.src[1:]), walk=True, name="params to buffers")
+  return graph_rewrite(linear, pm_collapse_after+_remove_all_tags, name="remove tags/after")
 
-  return buffer_map, schedule, var_vals
+pm_schedule = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK),), allow_any_len=True, name="big_sink"), lower_schedule_to_linear),
+])
+
+@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0]))}")
+def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[list[ExecItem], dict[str, int]]:
+  # big_sink srcs are all the Tensors
+  linear = graph_rewrite(big_sink, pm_schedule, name="schedule to linear")
+
+  # vars used in the schedule
+  used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
+  # get var_vals
+  var_vals: dict[str, int] = {}
+  for b in big_sink.src[1:]:
+    if b.op is Ops.BIND:
+      nm = b.src[0].expr
+      if nm not in used_vars: continue
+      val = b.src[1].arg
+      assert nm not in var_vals or var_vals[nm] == val, f"bind mismatch on {nm}, {var_vals[nm]} != {val}"
+      var_vals[nm] = val
+
+  # convert LINEAR to ExecItems
+  schedule: list[ExecItem] = linear_to_schedule(linear)
+  with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
+  return schedule, var_vals

tinygrad/function.py

@@ -0,0 +1,62 @@
+import functools
+from typing import Generic, TypeVar, Callable, cast
+from tinygrad.helpers import Context, dedup, getenv
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
+from tinygrad.tensor import Tensor
+
+def add_to_ctx(ctx, x:UOp):
+  ret = x.param_like(len(ctx))
+  ctx.append(x)
+  return ret
+
+pm_ctx = PatternMatcher([
+  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
+  (UPat((Ops.ASSIGN, Ops.CONTIGUOUS), name="x"),
+   lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) else None),
+])
+
+ReturnType = TypeVar('ReturnType')
+class function(Generic[ReturnType]):
+  def __init__(self, fxn:Callable[..., ReturnType]):
+    self.fxn = fxn
+
+  def __get__(self, obj, objtype=None): return functools.partial(self.__call__, obj) if obj is not None else self
+
+  def __call__(self, *args, **kwargs) -> ReturnType:
+    input_uops: list[UOp] = [(t.uop if isinstance(t, Tensor) else t)
+                             for name,t in list(enumerate(args))+sorted(kwargs.items()) if isinstance(t, (Tensor, UOp))]
+
+    # use the base
+    #input_uops = [x.multibase for x in input_uops]
+
+    # deduplicate input_uops, keeping the first occurrence index for each unique uop
+    call_uops: list[UOp] = dedup(input_uops)
+
+    # disable realize/schedule while this is running
+    # run it and do surgery later
+    with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
+      ret = self.fxn(*args, **kwargs)
+    assert isinstance(ret, Tensor), "only supports one tensor return for now"
+
+    # replace the known inputs with params (using deduplicated slots)
+    subs = {}
+    for i,x in enumerate(call_uops): subs[x] = x.param_like(i)
+    uret = ret.uop.substitute(subs)
+
+    # add contiguous to call_uops
+    #call_uops = [x.contiguous() for x in call_uops]
+
+    # the BUFFERs that are left are the implicit inputs
+    uret = graph_rewrite(uret, pm_ctx, call_uops, bottom_up=True, name="get_implicit_inputs")
+    name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
+
+    # assign output
+    #pbuffer = uret.param_like(len(call_uops))
+    #assigned = pbuffer.assign(uret).sink()
+    #buffer = UOp.new_buffer(pbuffer.device, pbuffer.size, pbuffer.dtype).reshape(uret.shape)
+    #call = assigned.call(*call_uops, buffer, name=name)
+    #ret = buffer.after(call)
+
+    ret = uret.call(*call_uops, name=name)
+    return cast(ReturnType, Tensor(ret, device=ret.device))
+

tinygrad/gradient.py

@@ -13,14 +13,20 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
     return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
   if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
-def call_gradient(ctx:UOp, k:UOp):
+def call_gradient(ctx:UOp, k:UOp) -> tuple[UOp|None, ...]:
   if k.arg.grad_fxn is not None: return (None,) + k.arg.grad_fxn(ctx, k)
   # auto-differentiate the function
   fxn, args = k.src[0], k.src[1:]
   params = sorted([x for x in fxn.toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
-  grads = compute_gradient(fxn, ctx, set(params))
-  subst = dict(zip(params, args))
-  return (None,) + tuple(grads[p].substitute(subst) if p in grads else None for p in params)
+  grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params))
+  ret: list[UOp|None] = [None]
+  for i,p in enumerate(params):
+    if p in grads:
+      # TODO: compact the args and remove unused ones
+      ret.append(grads[p].call(*args, ctx, name=(k.arg.name or "")+f"_backward_{i}"))
+    else:
+      ret.append(None)
+  return tuple(ret)
 
 # ctx is grad_output
 pm_gradient = PatternMatcher([

tinygrad/nn/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 import math
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import prod, make_tuple, flatten, USE_ATOMICS
 from tinygrad.nn import optim, state, datasets  # noqa: F401
 
@@ -36,7 +35,7 @@ class BatchNorm:
     self.weight: Tensor|None = Tensor.ones(sz) if affine else None
     self.bias: Tensor|None = Tensor.zeros(sz) if affine else None
 
-    self.num_batches_tracked = Tensor.zeros(dtype='long' if is_dtype_supported(dtypes.long) else 'int', requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(dtype='long', requires_grad=False)
     if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, requires_grad=False), Tensor.ones(sz, requires_grad=False)
 
   def calc_stats(self, x:Tensor) -> tuple[Tensor, Tensor]:

tinygrad/nn/state.py

@@ -78,7 +78,7 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=No
   j += "\x20"*(round_up(len(j),8)-len(j))
   pathlib.Path(fn).unlink(missing_ok=True)
   t = Tensor.empty(8+len(j)+offset, dtype=dtypes.uint8, device=f"disk:{fn}")
-  t[0:8].assign(Tensor([len(j)], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
+  t[0:8].bitcast(dtypes.int64).assign([len(j)])
   t[8:8+len(j)].assign(list(j.encode('utf-8')))
   for k,v in safe_load(t).items(): v.assign(tensors[k])
 

tinygrad/runtime/ops_amd.py

@@ -865,23 +865,25 @@ class PCIIface(PCIIfaceBase):
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=pv,
       read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), params=rcvr_params)
 
-  def sleep(self, timeout):
-    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
-      self.pci_dev.irq_fd.read(8 * events_cnt)
-    self.dev_impl.ih.interrupt_handler()
-    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")
-
-  def on_device_hang(self):
+  def _collect_faults(self, reset=False):
     devs:list[AMDDevice] = [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, AMDDevice) and d.is_am()]
-    for d in devs: d.iface.dev_impl.ih.interrupt_handler()
-    faults = [f for d in devs if (f:=d.iface.dev_impl.gmc.check_fault())]
     for d in devs:
-      if d.iface.dev_impl.recover():
+      d.iface.dev_impl.ih.interrupt_handler()
+      if reset and d.iface.dev_impl.recover():
         d.compute_queue.put_value, _ = d.iface.dev_impl.gfx.setup_ring(*d.compute_queue.params)
         d.compute_queue.read_ptr[0] = d.compute_queue.write_ptr[0] = d.compute_queue.put_value
         d.timeline_signal.value = d.timeline_value - 1
         d.error_state = None
-    raise RuntimeError(f"Device hang detected: {'; '.join(faults)}" if faults else "Device hang detected")
+
+  def sleep(self, timeout):
+    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
+      self.pci_dev.irq_fd.read(8 * events_cnt)
+    self._collect_faults()
+    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")
+
+  def on_device_hang(self):
+    self._collect_faults(reset=True)
+    raise RuntimeError("Device hang detected")
 
   def device_fini(self): self.dev_impl.fini()
 

tinygrad/runtime/support/am/amdev.py

@@ -254,8 +254,8 @@ class AMDev(PCIDevImplBase):
     else: self.mmio[reg] = val
 
   def wreg_pair(self, reg_base:str, lo_suffix:str, hi_suffix:str, val:int, inst:int=0):
-    self.reg(f"{reg_base}{lo_suffix}").write(val & 0xffffffff, inst=inst)
-    self.reg(f"{reg_base}{hi_suffix}").write(val >> 32, inst=inst)
+    self.reg(f"{reg_base}{lo_suffix}").write(lo32(val), inst=inst)
+    self.reg(f"{reg_base}{hi_suffix}").write(hi32(val), inst=inst)
 
   def indirect_rreg(self, reg:int) -> int:
     self.reg("regBIF_BX_PF0_RSMU_INDEX").write(reg * 4)
@@ -268,9 +268,9 @@ class AMDev(PCIDevImplBase):
   def indirect_wreg_pcie(self, reg:int, val:int, aid:int=0):
     reg_addr = reg * 4 + ((((aid & 0b11) << 32) | (1 << 34)) if aid > 0 else 0)
     self.reg("regBIF_BX0_PCIE_INDEX2").write(lo32(reg_addr))
-    if reg_addr >> 32: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(hi32(reg_addr) & 0xff)
+    if hi32(reg_addr) > 0: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(hi32(reg_addr) & 0xff)
     self.reg("regBIF_BX0_PCIE_DATA2").write(val)
-    if reg_addr >> 32: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(0)
+    if hi32(reg_addr) > 0: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(0)
 
   def _read_vram(self, addr, size) -> bytes:
     assert addr % 4 == 0 and size % 4 == 0, f"Invalid address {addr:#x} or size {size:#x}"

tinygrad/runtime/support/am/ip.py

@@ -171,12 +171,6 @@ class AM_GMC(AM_IP):
     if self.adev.ip_ver[am.GC_HWIP] < (10,0,0): return (pte & am.AMDGPU_PDE_PTE) if pte_lv != am.AMDGPU_VM_PDB0 else not (pte & am.AMDGPU_PTE_TF)
     return pte & (am.AMDGPU_PDE_PTE_GFX12 if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else am.AMDGPU_PDE_PTE)
 
-  def check_fault(self) -> str|None:
-    va = (self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_HI32').read()<<32) | self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_LO32').read()
-    if self.adev.reg(self.pf_status_reg("GC")).read():
-      return f"am {self.adev.devfmt}: GCVM_L2_PROTECTION_FAULT_STATUS: {self.adev.reg(self.pf_status_reg('GC')).read_bitfields()} {va<<12:#x}"
-    return None
-
 class AM_SMU(AM_IP):
   def init_sw(self):
     self.smu_mod = self.adev._ip_module("smu", am.MP1_HWIP, prever_prefix='v')
@@ -461,6 +455,11 @@ class AM_IH(AM_IP):
         err_info = f" ({['EDC_FUE', 'ILLEGAL_INST', 'MEMVIOL', 'EDC_FED'][err_type]})" if enc_type == 2 else ""
         print(f"am {self.adev.devfmt}: sq_intr: {['auto', 'wave', 'error'][enc_type]}{err_info}")
         self.adev.is_err_state |= enc_type == 2
+      elif src_name == "UTCL2_FAULT" or (self.adev.ip_ver[am.GC_HWIP][0] == 9 and client == am.SOC15_IH_CLIENTID_UTCL2):
+        bf = self.adev.reg(self.adev.gmc.pf_status_reg('GC')).read_bitfields()
+        va = (self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_HI32').read()<<32) | self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_LO32').read()
+        print(f"am {self.adev.devfmt}: GCVM_L2_PROTECTION_FAULT_STATUS: {bf} {va<<12:#x}")
+        self.adev.is_err_state = True
       else: self.adev.is_err_state = True
 
       rptr = (rptr + 8) % (self.ring_size // 4)

tinygrad/schedule/indexing.py

@@ -18,6 +18,10 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
 def realize_assign_src(ctx:dict[UOp, None], buf:UOp, x:UOp):
+  # don't realize COPY/BUFFER_VIEW/ENCDEC when they are the direct source of ASSIGN — the ASSIGN target buffer is the output
+  if x.op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC} and x in ctx \
+     and not buf.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
+    del ctx[x]
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
   if buf.base in x.backward_slice_with_self: ctx[x] = None
 
@@ -52,7 +56,7 @@ class IndexingContext:
     return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
 
 def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
-  if x.op in {Ops.BUFFERIZE, Ops.INDEX, Ops.AFTER}: return None
+  if x.op in {Ops.BUFFERIZE, Ops.INDEX}: return None
   new_srcs = []
   for s in x.src:
     new_src = s
@@ -118,8 +122,6 @@ pm_apply_rangeify = PatternMatcher([
   (UPat(GroupOp.All, name="x"), create_bufferize_and_index_based_on_ranges),
   # remove movement op
   (UPat(GroupOp.Movement, name="x"), remove_movement_op_after_rangeify),
-  # const/define_var shouldn't have src
-  (UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"), lambda ctx,c: c.replace(src=()) if c in ctx.range_map else None),
 ])
 
 @functools.cache
@@ -146,10 +148,9 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
     case Ops.FLIP:    rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg))
     case Ops.EXPAND:  rngs = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, arg))
     case Ops.PAD:
-      # TODO: why is multiple graph_rewrites faster than one here?
-      # TODO: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local`
+      # NOTE: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local`
       #       wraps the pad with only the newly added valid
-      rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh+s))),
+      rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite((r >= s) & (r < (sh+s)),
         symbolic+pm_simplify_valid, name="pad").where(r-s, UOp.invalid()) for r,sh,(s,e) in zip(rngs, in_shape, arg))
     case Ops.RESHAPE:
       sink = UOp.sink(*rngs)
@@ -164,12 +165,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   rctx = IndexingContext()
 
   # get ops to realize
-  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, bottom_up=True, name="get realize")
-  # don't realize COPY/BUFFER_VIEW/ENCDEC when they are the direct source of ASSIGN — the ASSIGN target buffer is the output
-  for u in tsink.toposort():
-    if u.op is Ops.ASSIGN and u.src[1].op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC} and u.src[1] in rctx.realize_map \
-       and not u.src[0].op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
-      del rctx.realize_map[u.src[1]]
+  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize")
 
   # get the consumer map
   with cpu_profile("consumer map in rangeify", "TINY"):
@@ -181,7 +177,13 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
 
     # no ranges on kernels, they are internal
-    if x.op is Ops.CALL: continue
+    if x.op in {Ops.CALL, Ops.LINEAR}: continue
+
+    # no range on after
+    if x.op is Ops.AFTER: continue
+
+    # treat MSTACK/MSELECT like SINK
+    if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
 
     if x.dtype.scalar() == dtypes.index: continue  # TODO: why do I need this?
     ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
@@ -200,9 +202,6 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       # mark all ranges as ended
       assert rctx.realize_map[x] is None
       rctx.realize_map[x] = list(range(len(x.shape)))
-    elif x.op in {Ops.MSTACK, Ops.MSELECT}:
-      # treat MSTACK/MSELECT like SINK
-      continue
     elif len(consumer_rngs) == 0:
       # if no consumers have ranges and this isn't realized, this doesn't have ranges either.
       continue
@@ -237,7 +236,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
 
     # if this element is a reduce and there's ended ranges, we might have to end some other ranges
     if len(ending_ranges[x]) and x.op in GroupOp.Elementwise.union({Ops.REDUCE_AXIS}):
-      _realize_axis = rctx.realize_map.get(x, []) or []
+      _realize_axis = rctx.realize_map.get(x) or []
       for i,r in enumerate(out_rngs):
         if i in _realize_axis: continue
         if not (PCONTIG > 1) or any(any(rr.arg > e.arg for e in ending_ranges[x]) for rr in r.ranges):

tinygrad/schedule/multi.py

@@ -1,6 +1,6 @@
 import functools, itertools
 from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, getenv
-from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
+from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite, should_resolve_call
 from tinygrad.dtype import dtypes
 
 # *** allreduce implementation ***
@@ -163,6 +163,9 @@ def assign_multi(dest:UOp, src:UOp):
 def passthrough_multi(root:UOp, multi:UOp):
   return UOp(root.op, root.dtype, (multi.src[0],)+tuple(x.src[0] if x.op is Ops.MULTI else x for x in root.src[1:]), root.arg).multi(multi.axis)
 
+def rewrite_into_call(call:UOp):
+  return call.replace(src=(graph_rewrite(call.src[0], multi_pm, name="subcall"),)+call.src[1:]) if should_resolve_call(call) else None
+
 # NOTE: this is the same pattern as Ops.UNROLL
 multi_pm = PatternMatcher([
   (UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.MULTI])), alu_multi),
@@ -177,6 +180,8 @@ multi_pm = PatternMatcher([
   (UPat(Ops.COPY, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.DEVICE, name="device"))), copy_multi),
   (UPat(Ops.ALLREDUCE, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.DEVICE, name="device")), name="red"),
     lambda multi,device,red: multi.src[0].allreduce(red.arg, device).multi(axis=multi.axis)),
+  # rewrite into calls explicitly for MULTI
+  (UPat(Ops.CALL, name="call"), rewrite_into_call),
   (UPat(Ops.CALL, src=(UPat(Ops.MULTI, name="multi"), ), name="root", allow_any_len=True), passthrough_multi),
   # we just remove the MULTI from CALLs with dtypes.void and assume they are handled by the user for custom kernels
   (UPat(Ops.CALL, dtype=dtypes.void, name="root", custom_early_reject=set([Ops.MULTI])), lambda root:

tinygrad/schedule/rangeify.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field, replace
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo
-from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches
+from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, should_resolve_call
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
 from tinygrad.helpers import PCONTIG, partition, get_single_element
@@ -76,22 +76,23 @@ mop_cleanup = PatternMatcher([
 ])
 
 pm_gather_params = PatternMatcher([ (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.append(p)), ])
-def resolve_call(c:UOp, allow_param_mismatch=False) -> UOp|None:
-  # don't resolve real kernel calls, sink or program
-  if c.src[0].op is Ops.SINK and isinstance(c.src[0].arg, KernelInfo): return None
-  if c.src[0].op is Ops.PROGRAM: return None
+def resolve_call(c:UOp, allow_param_mismatch=True) -> UOp|None:
+  if not should_resolve_call(c): return None
   params: list[UOp] = []
-  graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params)
+  graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params, name="gather params")
   params = sorted(params, key=lambda x: x.arg)
   args = c.src[1:]
-  # TODO: this check belongs in spec, not here
+
+  # NOTE: this isn't really needed. it's okay if there's unused args in the function
   if not allow_param_mismatch:
     if [x.arg for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg for x in params]}")
     if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")
-  for i, (p, a) in enumerate(zip(params, args)):
-    if p.shape != a.shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
+
+  dict_map = {x:args[x.arg] for x in params}
+  for i, (p, a) in enumerate(dict_map.items()):
+    if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
     if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
-  return c.src[0].substitute(dict(zip(params, args)))
+  return c.src[0].substitute(dict_map, walk=True)
 
 earliest_rewrites = mop_cleanup+PatternMatcher([
   # resolve calls
@@ -100,6 +101,9 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # split_reduceop
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
 
+  # remove DETACH/CONTIGUOUS_BACKWARD (TODO: this is copied in allocations)
+  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
+
   # remove contiguous on movement ops before a copy on disk
   (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, allow_any_len=True, name="copy"),
    lambda x,copy: copy.replace(src=(x,)+copy.src[1:]) if isinstance(x.device, str) and x.device.startswith("DISK") else None),
@@ -360,6 +364,11 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
 
   # remove any RESHAPEs on KERNEL
   (UPat(Ops.CALL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
+
+  # remove MOP on AFTER
+  (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(GroupOp.Movement, name="y"))), lambda x,y: x.after(y.src[0])),
+  # remove double AFTER
+  (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.AFTER, name="y"))), lambda x,y: x.after(*y.src[1:]))
 ])
 
 pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
@@ -467,7 +476,8 @@ def split_store(x:UOp) -> UOp|None:
   if ret.op is Ops.STORE: stored = ret.src[1]
   elif ret.op is Ops.END and ret.src[0].op is Ops.STORE: stored = ret.src[0].src[1]
   else: raise RuntimeError(f"unknown kernel type {ret.op}")
-  if stored.op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC}: ret = stored
+  if stored.op in {Ops.COPY, Ops.BUFFER_VIEW}: ret = stored.replace(src=stored.src + ret.ended_ranges)
+  elif stored.op is Ops.ENCDEC: ret = stored
   else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
 
   kernel = ret.call(*lctx.map.values(), *lctx.vars.keys())
@@ -481,7 +491,7 @@ split_kernels = PatternMatcher([
 
 @profile_matches
 def get_kernel_graph(sink:UOp) -> UOp:
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm", rewrite_into_calls=True)
+  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
   tsink = graph_rewrite(tsink, pm_syntactic_sugar+pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
 
   # convert movement ops to ranges
@@ -511,4 +521,4 @@ def get_kernel_graph(sink:UOp) -> UOp:
       assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
   if assign_rep: tsink = graph_rewrite(tsink, _substitute, ctx=assign_rep, bottom_up=True, name="fix_assign")
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
-  return tsink
+  return tsink

tinygrad/tensor.py

@@ -16,6 +16,7 @@ from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_eleme
 from tinygrad.engine.schedule import ExecItem, complete_create_schedule_with_vars
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
+from tinygrad.engine.allocations import transform_to_call
 
 # TODO: this should be the only usage of Device
 def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
@@ -249,17 +250,23 @@ class Tensor(OpMixin):
     """
     return [Tensor(u, device=u.device) for u in UOp.custom_kernel(*[t.uop for t in (self,)+lst], fxn=fxn, grad_fxn=grad_fxn)]
 
+  def callify(self, *lst:Tensor) -> Tensor:
+    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
+    big_sink, buffer_map = transform_to_call(big_sink)
+    _apply_map_to_tensors({x:y.after(big_sink) for x,y in buffer_map.items()}, name="callify")
+    return self
+
   def schedule_with_vars(self, *lst:Tensor) -> tuple[list[ExecItem], dict[str, int]]:
     """
     Creates the schedule needed to realize these Tensor(s), with Variables.
 
     NOTE: A Tensor can only be scheduled once.
     """
-    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
+    big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
+    _apply_map_to_tensors(becomes_map, name="buffers")
 
     # this is where the schedule cache should go
-    becomes_map, schedule, var_vals = complete_create_schedule_with_vars(big_sink)
-    _apply_map_to_tensors(becomes_map, name="buffers")
+    schedule, var_vals = complete_create_schedule_with_vars(big_sink)
     return schedule, var_vals
 
   def schedule(self, *lst:Tensor) -> list[ExecItem]:
@@ -278,7 +285,8 @@ class Tensor(OpMixin):
           # recursively realize pending assigns that this assign's value depends on
           for u in assign_uop.toposort():
             if u.op is Ops.BUFFER and u in _pending_assigns: _realize_pending(u)
-          becomes_map, schedule, var_vals = complete_create_schedule_with_vars(UOp.sink(assign_uop))
+          big_sink, becomes_map = transform_to_call(UOp.sink(assign_uop))
+          schedule, var_vals = complete_create_schedule_with_vars(big_sink)
           _apply_map_to_tensors(becomes_map, name="Apply Pending Assign")
           run_schedule(schedule, var_vals, do_update_stats=do_update_stats)
           # update remaining pending assigns so they reference realized buffers instead of stale lazy graphs
@@ -308,7 +316,7 @@ class Tensor(OpMixin):
     if self.shape != x.shape: x = x._broadcast_to(self.shape)
     if self.shape != x.shape: raise RuntimeError(f"assign shape mismatch {self.shape} != {x.shape}")
     if not is_disk and self.device != x.device: raise RuntimeError(f"assign device mismatch {self.device} != {x.device}")
-    if self.dtype != x.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
+    if not is_disk and self.dtype != x.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
     if isinstance(self.device, tuple) and self.uop.axis != x.uop.axis: raise RuntimeError(f"multi axis mismatch {self.uop.axis} != {x.uop.axis}")
 
     # TODO: this is a hack for writing to DISK. remove with working assign
@@ -3561,10 +3569,7 @@ class Tensor(OpMixin):
 
   def bitcast(self, dtype:DTypeLike) -> Tensor:
     """
-    Bitcasts `self` to the given `dtype`.
-
-    When the target dtype has the same itemsize, this is a view of the same memory.
-    When itemsizes differ, the last dimension is adjusted and a new Tensor is created.
+    Bitcasts `self` to the given `dtype` of the same itemsize.
 
     `self` must not require a gradient.
 

tinygrad/uop/ops.py

@@ -26,7 +26,7 @@ axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL:
 axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
                AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
-range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1}
+range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.COPY: 2, Ops.BUFFER_VIEW: 1}
 
 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> PyConst: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
@@ -212,7 +212,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
         return None
 
       case Ops.CAST:
-        # when PTX cases from ptr to non ptr, remove the shape
+        # when PTX casts from ptr to non ptr, remove the shape
         if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType):
           return None
 
@@ -329,11 +329,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def shard_size(self) -> int: return prod(self.max_shard_shape)
 
   @functools.cached_property
-  def ended_ranges(self):
+  def ended_ranges(self) -> tuple[UOp, ...]:
     if self.op in range_start: return self.src[range_start[self.op]:]
     if self.op is Ops.AFTER: return tuple(flatten([x.ended_ranges for x in self.src[1:]]))
-    # TODO: copy isn't using range properly and isn't ending the range it uses, remove this
-    if self.op in {Ops.COPY, Ops.BUFFER_VIEW}: return self.src[0].ranges
     return ()
 
   # determine what ranges this is in
@@ -375,11 +373,12 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def __bool__(self): return self._eval((dtypes.bool,), bool)
   def __int__(self): return self._eval(dtypes.ints, int)
   def __float__(self): return float(self._eval(dtypes.floats, float))
-  def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None):
+  def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None, walk:bool=False):
     dvars = {k:v for k,v in dvars.items() if k is not v}
     if len(dvars) == 0: return self
     with Context(TRACK_MATCH_STATS=(0 if name is None else TRACK_MATCH_STATS.value)):
-      return graph_rewrite(self, (extra_pm+_substitute) if extra_pm is not None else _substitute, dvars, bottom_up=True, name=name)
+      return graph_rewrite(self, (extra_pm+_substitute) if extra_pm is not None else _substitute, dvars,
+                           bottom_up=True, walk=walk, name=name)
   # NOTE: this is not called by Tensor slice (Tensor handles UOps directly), but satisfies SupportsIndex for type checking
   def __index__(self): return self.__int__()
 
@@ -865,11 +864,15 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if vmin_vmax is not None: src += (UOp.const(dtype, vmin_vmax[0]), UOp.const(dtype.scalar(), vmin_vmax[1]))
     if name is not None: src += (UOp(Ops.NOOP, arg=name),)
     return UOp(Ops.PARAM, dtype, src, arg=slot)
+  def param_like(self, slot:int):
+    if self.op is Ops.BIND:
+      return UOp.param(slot, self.dtype, self._shape, self._device, self._min_max, self.src[0].arg[0])
+    return UOp.param(slot, self.dtype, self._shape, self._device)
 
-  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=()) -> UOp:
+  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None) -> UOp:
     # TODO: reenable this after ENCDEC is fixed
     #assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
-    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata))
+    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, name))
   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)]
@@ -891,9 +894,17 @@ class KernelInfo:
 class CallInfo:
   grad_fxn: Callable|None = None
   metadata: tuple[Metadata, ...] = ()
+  name: str|None = None
   # grad_fxn can't be pickled, but metadata can
-  def __reduce__(self): return (CallInfo, (None, self.metadata))
-  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata})"
+  def __reduce__(self): return (CallInfo, (None, self.metadata, self.name))
+  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata}, {repr(self.name)})"
+
+def should_resolve_call(c:UOp) -> bool:
+  # 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 False
+  if c.src[0].op is Ops.PROGRAM: return False
+  if c.src[0].op is Ops.COPY: return False
+  return True
 
 # ******** ops in python ********
 
@@ -1239,13 +1250,12 @@ if TRACK_MATCH_STATS or PROFILE:
 SENTINEL: Final[UOp] = cast(UOp, object())
 class BottomUpGate(Exception): pass
 class RewriteContext:
-  def __init__(self, pm, bpm, ctx=None, rewrite_into_calls=False):
+  def __init__(self, pm, bpm, ctx=None):
     self.pm: PatternMatcher|None = pm
     self.bpm: PatternMatcher|None = bpm
     self.bpm_cache: dict[UOp, UOp|None] = {}
     self.ctx = ctx
     self.replace: dict[UOp, UOp] = {}
-    self.rewrite_into_calls = rewrite_into_calls
 
   # no cache needed: pm_rewrite is called at most once per UOp due to the replace dict check in unified_rewrite
   def pm_rewrite(self, x:UOp) -> UOp|None: return unwrap(self.pm).rewrite(x, self.ctx)
@@ -1255,6 +1265,31 @@ class RewriteContext:
     ret = self.bpm_cache[x] = unwrap(self.bpm).rewrite(x, self.ctx)
     return ret
 
+  def walk_rewrite(self, root:UOp) -> UOp:
+    """MLIR-style Walk Pattern Rewrite Driver: single-pass, no re-traversal into rewritten subtrees."""
+    stack: list[tuple[UOp, bool]] = [(root, False)]
+    while stack:
+      n, processed = stack.pop()
+      if n in self.replace: continue
+      if not processed:
+        # bottom-up: try bpm on original node first, if it rewrites, use result as-is (no traversal into replacement)
+        if self.bpm is not None and (rewritten:=self.cached_bpm_rewrite(n)) is not None:
+          self.replace[n] = rewritten
+          continue
+        # no rewrite, process children then come back to rebuild
+        stack.append((n, True))
+        if n.op is Ops.CALL: self.replace[n.src[0]] = n.src[0]
+        for x in reversed(n.src):
+          if x not in self.replace: stack.append((x, False))
+      else:
+        # rebuild node with rewritten srcs
+        new_src = tuple(self.replace.get(x, x) for x in n.src)
+        new_n = UOp(n.op, n.dtype, new_src, n.arg, n.tag) if new_src != n.src else n
+        # top-down: try pm on rebuilt node, use result as-is (no re-traversal)
+        if self.pm is not None and (rewritten:=self.pm_rewrite(new_n)) is not None: new_n = rewritten
+        self.replace[n] = new_n
+    return self.replace.get(root, root)
+
   def unified_rewrite(self, root:UOp) -> UOp:
     stack: collections.deque[tuple[UOp, int, UOp]] = collections.deque([(root, 0, root)])
     on_stack = {root}  # all UOps either on the stack or in self.replace, i.e. dont have to be placed again
@@ -1283,7 +1318,7 @@ class RewriteContext:
         # NOTE: CALL is handled as a special case.
         # The function that is called is not included in the graph_rewrite.
         # If you want to graph_rewrite a call, you can
-        if new_n.op is Ops.CALL and not self.rewrite_into_calls: self.replace[new_n.src[0]] = new_n.src[0]
+        if new_n.op is Ops.CALL: self.replace[new_n.src[0]] = new_n.src[0]
         for x in reversed(new_n.src):
           if x in on_stack: continue
           stack.append((x, 0, x))
@@ -1322,9 +1357,9 @@ class RewriteContext:
     return self.replace[root]
 
 @profile_matches
-def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, rewrite_into_calls=False) -> UOp:
-  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, rewrite_into_calls=rewrite_into_calls)
-  return rewrite_ctx.unified_rewrite(sink)
+def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, walk=False) -> UOp:
+  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx)
+  return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)
 
 def sint_to_uop(x:sint, dtype=dtypes.index) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype)
 

tinygrad/uop/spec.py

@@ -205,6 +205,12 @@ kernel_spec = PatternMatcher([
 
   # reduce must be on ranges
   (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype in (dtypes.index, dtypes.int) for y in x.src[1:])),
+
+  # COPY/BUFFER_VIEW can have ranges appended
+  (UPat(Ops.COPY, name="x", src=(UPat.var("s"), UPat(Ops.DEVICE)), allow_any_len=True, arg=None),
+   lambda x,s: x.dtype == s.dtype and all(u.op is Ops.RANGE for u in x.src[2:])),
+  (UPat(Ops.BUFFER_VIEW, src=(UPat((Ops.INDEX, Ops.LOAD)),), allow_any_len=True, name="x"),
+   lambda x: all(u.op is Ops.RANGE for u in x.src[1:])),
 ])+movement_ops+shared_codegen_spec+shared_spec
 
 tensor_spec = PatternMatcher([

tinygrad/uop/symbolic.py

@@ -258,7 +258,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   ((UPat.var("x", dtypes.index) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
   # only RANGE/IF/STORE/KERNEL have side effects
   (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
-    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.UNROLL} else y.src for y in x.src[1:]])))),
+    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR, Ops.BUFFERIZE}
+                   else y.src for y in x.src[1:]])))),
   # after with 1 src is just src[0]
   (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
   # VECTORIZE/CONST

tinygrad/viz/index.html

@@ -102,6 +102,14 @@
     fill: #FFD700;
     stroke: #B8860B;
   }
+  g.tag.collapsed circle {
+    fill: #5CD68D;
+    stroke: #4a4b57;
+  }
+  g.tag.expanded circle {
+    fill: #9FDDE6;
+    stroke: #4a4b57;
+  }
   g.port circle {
     fill: #b3dcc2;
   }
@@ -109,6 +117,7 @@
     stroke-width: 0.8;
   }
   g.tag text, #edge-labels text {
+    font-family: monospace;
     text-anchor: middle;
     font-size: 6px;
     fill: #08090e;

tinygrad/viz/js/index.js

@@ -59,9 +59,14 @@ const drawGraph = (data) => {
   const g = dagre.graphlib.json.read(data);
   // draw nodes
   d3.select("#graph-svg").on("click", () => d3.selectAll(".highlight").classed("highlight", false));
+  const callCount = g.graph().callCount;
   const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g").attr("class", d => d.className ?? "node")
-    .attr("transform", d => `translate(${d.x},${d.y})`).classed("clickable", d => d.ref != null).on("click", (e,d) => {
-      if (d.ref != null) return switchCtx(d.ref);
+    .attr("transform", d => `translate(${d.x},${d.y})`).on("click", (e,d) => {
+      if (d.label.startsWith("CALL")) {
+        if (state.callSrcMask.has(d.id)) state.callSrcMask.delete(d.id); else state.callSrcMask.add(d.id);
+        if (state.callSrcMask.size >= callCount) { showCallSrc.toggle.checked = !showCallSrc.toggle.checked; state.callSrcMask.clear(); }
+        return setState({});
+      }
       const parents = g.predecessors(d.id);
       const children = g.successors(d.id);
       if (parents == null && children == null) return;
@@ -105,6 +110,9 @@ const drawGraph = (data) => {
   });
   addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
     .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag));
+  addTags(nodes.selectAll("g.type").data(d => d.label.startsWith("CALL\n") ? [d] : []).join("g")
+    .attr("class", d => `tag ${d.collapsed ? 'collapsed' : 'expanded'}`)
+    .attr("transform", d => `translate(${-d.width/2}, ${0})`).datum(d => d.collapsed ? "+" : "−"));
   // draw edges
   const line = d3.line().x(d => d.x).y(d => d.y).curve(d3.curveBasis), edges = g.edges();
   d3.select("#edges").selectAll("path.edgePath").data(edges).join("path").attr("class", "edgePath").attr("d", (e) => {
@@ -707,7 +715,7 @@ const evtSources = [];
 // rewrite: a single UOp transformation
 // step: collection of rewrites
 // context: collection of steps
-const state = {currentCtx:-1, currentStep:0, currentRewrite:0, expandSteps:false};
+const state = {currentCtx:-1, currentStep:0, currentRewrite:0, expandSteps:false, callSrcMask:new Set()};
 function setState(ns) {
   saveToHistory(state);
   const { ctx:prevCtx, step:prevStep } = select(state.currentCtx, state.currentStep);
@@ -755,7 +763,7 @@ const createToggle = (id, text) => {
   return { toggle, label };
 }
 const showIndexing = createToggle("show-indexing", "Show indexing (r)");
-const showCallSrc = createToggle("show-call-src", "Show CALL src (c)");
+const showCallSrc = createToggle("show-call-src", "Show all CALL src (c)"); showCallSrc.toggle.checked = false;
 const showSink = createToggle("show-sink", "Show SINK (s)");
 showSink.toggle.checked = false;
 const showGraph = createToggle("show-graph", "Show graph (g)");
@@ -907,10 +915,10 @@ async function main() {
   // ** center graph
   const data = ret[currentRewrite];
   const render = (opts) => renderDag({ data, opts }, { recenter:currentRewrite === 0 });
-  const getOpts = () => ({ showIndexing:showIndexing.toggle.checked, showCallSrc:showCallSrc.toggle.checked, showSink:showSink.toggle.checked });
+  const getOpts = () => ({ showIndexing:showIndexing.toggle.checked, showCallSrc:showCallSrc.toggle.checked, showSink:showSink.toggle.checked, callSrcMask:state.callSrcMask });
   render(getOpts());
   showIndexing.toggle.onchange = () => render(getOpts());
-  showCallSrc.toggle.onchange = () => render(getOpts());
+  showCallSrc.toggle.onchange = () => { state.callSrcMask.clear(); render(getOpts()); }
   showSink.toggle.onchange = () => render(getOpts());
   // ** right sidebar metadata
   metadata.innerHTML = "";
@@ -942,7 +950,7 @@ async function main() {
       metadata.appendChild(codeBlock(upat[1], "python", { loc:upat[0], wrap:true }));
       const diffCode = metadata.appendChild(document.createElement("pre")).appendChild(document.createElement("code"));
       for (const line of diff) {
-        diffCode.appendChild(colored([{st:line, color:line.startsWith("+") ? "#3aa56d" : line.startsWith("-") ? "#d14b4b" : "#f0f0f5"}]));
+        diffCode.appendChild(colored([{st:line, color:line.startsWith("+") ? "#3aa56d" : line.startsWith("−") ? "#d14b4b" : "#f0f0f5"}]));
         diffCode.appendChild(document.createElement("br"));
       }
       diffCode.className = "wrap";

tinygrad/viz/js/worker.js

@@ -46,6 +46,7 @@ const layoutUOp = (g, { graph, change }, opts) => {
   g.setGraph({ rankdir: "LR", font:"sans-serif", lh:lineHeight });
   ctx.font = `350 ${lineHeight}px ${g.graph().font}`;
   if (change?.length) g.setNode("overlay", {label:"", labelWidth:0, labelHeight:0, className:"overlay"});
+  let callCount = 0;
   for (const [k, {label, src, ref, color, tag }] of Object.entries(graph)) {
     // adjust node dims by label size (excluding escape codes) + add padding
     let [width, height] = [0, 0];
@@ -53,11 +54,13 @@ const layoutUOp = (g, { graph, change }, opts) => {
       width = Math.max(width, ctx.measureText(line).width);
       height += lineHeight;
     }
+    if (label.startsWith("CALL\n")) callCount++;
     g.setNode(k, {...rectDims(width, height), label, ref, id:k, color, tag});
     // add edges
     const edgeCounts = {};
     for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
-    for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port}});
+    for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port},
+      ...(label.startsWith("CALL\n") && port === 0 && {color:"#a0a1b8"})});
     if (change?.includes(parseInt(k))) g.setParent(k, "overlay");
   }
   // optionally hide nodes from the layout
@@ -73,12 +76,13 @@ const layoutUOp = (g, { graph, change }, opts) => {
       if (node.label.includes("dtypes.index")) g.removeNode(n);
     }
   }
-  if (!opts.showCallSrc) {
+  if (!opts.showCallSrc || opts.callSrcMask.size > 0) {
     // remove edges from src[0] to CALL nodes, track affected nodes
     const disconnected = new Set();
     for (const n of g.nodes()) {
       const node = g.node(n);
-      if (node.label.startsWith("CALL\n")) {
+      if (node.label.startsWith("CALL\n") && (opts.showCallSrc ? opts.callSrcMask.has(n) : !opts.callSrcMask.has(n))) {
+        node.collapsed = true;
         for (const pred of (g.predecessors(n) || [])) {
           const edge = g.edge(pred, n);
           if (edge?.label?.text === 0) {
@@ -102,6 +106,7 @@ const layoutUOp = (g, { graph, change }, opts) => {
       }
     }
   }
+  g.graph().callCount = callCount;
   dagre.layout(g);
   // remove overlay node if it's empty
   if (!g.node("overlay")?.width) g.removeNode("overlay");