CALL with return value is FUNCTION

* fix symbolic contiguous_view_offset

* flatten

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

@@ -17,7 +17,7 @@ export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

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

@@ -17,7 +17,7 @@ export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

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

@@ -18,7 +18,7 @@ export FP8=1
 export ALLREDUCE_CAST=1
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=8 MP=1 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=4
+export DP=8 MP=1 BS=16 EVAL_BS=16 GRADIENT_ACC_STEPS=2
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

extra/viz/cli.py

@@ -52,16 +52,15 @@ def get(data:dict, key:str):
   raise RuntimeError(f'item "{key}" not found in list'+(f", did you mean {match[0]!r}?" if match else ''))
 
 def main(args) -> None:
-  data = viz.VizData(viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {})))
-  viz.load_rewrites(data)
+  viz.load_rewrites(viz_data:=viz.VizData(viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))))
 
   def format_colored(s:str) -> str: return ansistrip(s) if args.no_color else s
 
   if args.profile:
     events:list = viz.load_pickle(args.profile_path, default=[])
-    if (profile_bytes:=viz.get_profile(data, events)) is None: raise RuntimeError(f"empty profile in {args.profile_path}")
+    if (profile_bytes:=viz.get_profile(viz_data, events)) is None: raise RuntimeError(f"empty profile in {args.profile_path}")
     profile = decode_profile(profile_bytes)
-    profile["layout"].update([(f'{c["name"][5:]}{" SQTT" if s["name"].endswith("PKTS") else ""} {s["name"]}', s["data"]) for c in data.ctxs
+    profile["layout"].update([(f'{c["name"][5:]}{" SQTT" if s["name"].endswith("PKTS") else ""} {s["name"]}', s["data"]) for c in viz_data.ctxs
                               if c["name"].startswith("SQTT") for s in c["steps"] if s["name"].endswith(("PMC", "PKTS"))])
     if args.src is None:
       for k in profile["layout"]:
@@ -103,10 +102,10 @@ def main(args) -> None:
 
     # ** PMC printer
     if "PMC" in args.src:
-      table = viz.unpack_pmc(data[0])
-      cols = table["cols"]
+      pmc = viz.unpack_pmc(data)
+      cols = pmc["cols"]
       rows:list = []
-      for r in table["rows"]:
+      for r in pmc["rows"]:
         if args.item is None: rows.append(r[:2])
         elif args.item == r[0]:
           rows = r[2]["rows"] if len(r) > 2 else [r[:2]]
@@ -132,17 +131,17 @@ def main(args) -> None:
     if agg and total > 0:
       from tabulate import tabulate
       items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
-      rows = 20
-      table = [[format_colored(name), time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in items[:rows]]
-      if items[rows:]:
-        other_t = sum(t for _,(t,_) in items[rows:])
-        other_c = sum(c for _,(_,c) in items[rows:])
+      num_rows = 20
+      table = [[format_colored(name), time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in items[:num_rows]]
+      if items[num_rows:]:
+        other_t = sum(t for _,(t,_) in items[num_rows:])
+        other_c = sum(c for _,(_,c) in items[num_rows:])
         table.append(["Other", time_to_str(other_t, w=9), other_c, f"{(other_t/total*100.0):.2f}%"])
       print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
     return None
 
   # ** Graph rewrites printer
-  rewrites = {c["name"]:{s["name"]:s for s in c["steps"]} for c in data.ctxs if c.get("steps")}
+  rewrites = {c["name"]:{s["name"]:s for s in c["steps"]} for c in viz_data.ctxs if c.get("steps")}
   if args.src is None:
     for k in rewrites: print(f"  {format_colored(k)}")
     return None

pyproject.toml

@@ -50,8 +50,8 @@ tinygrad = ["py.typed"]
 
 
 [project.optional-dependencies]
-arm = ["unicorn"]
-triton = ["triton-nightly>=2.1.0.dev20231014192330"]
+# arm = ["unicorn"]
+# triton = ["triton-nightly>=2.1.0.dev20231014192330"]
 linting = [
   "pylint",
   "mypy==1.19.1",

test/amd/test_custom_kernel.py

@@ -15,7 +15,7 @@ from extra.gemm.amd_asm_matmul import Kernel
 def custom_add_one(A:UOp) -> UOp:
   A = A.flatten()
   assert dtypes.is_float(A.dtype.base), f"buffer dtype must be float32, got {A.dtype}"
-  threads = UOp.special(A.size, "lidx0")
+  threads = UOp.special(A.numel(), "lidx0")
   insts = [
     s_load_b64(s[0:1], s[0:1], soffset=NULL),
     s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
@@ -27,13 +27,13 @@ def custom_add_one(A:UOp) -> UOp:
     global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]),
     s_endpgm(),
   ]
-  sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.size}", estimates=Estimates(ops=A.size, mem=A.size*4*2)))
+  sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.numel()}", estimates=Estimates(ops=A.numel(), mem=A.numel()*4*2)))
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
 def custom_add_var(A:UOp, B:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
   assert A.dtype.base == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
-  threads = UOp.special(A.size, "lidx0")
+  threads = UOp.special(A.numel(), "lidx0")
   var = UOp.variable("var", 0, 10)
   insts = [
     s_load_b128(s[4:7], s[0:1]),
@@ -46,7 +46,7 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
     global_store_b32(addr=v[0], data=v[1], saddr=s[4:5]),
     s_endpgm(),
   ]
-  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.size}"))
+  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.numel()}"))
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
 def custom_wave_sync(A:UOp, arch:str) -> UOp:
@@ -132,7 +132,7 @@ def custom_handwritten(A:UOp, arch:str) -> UOp:
 
 def custom_data_deps(A:UOp, arch:str) -> UOp:
   A = A.flatten()
-  threads = UOp.special(A.size, "lidx0")
+  threads = UOp.special(A.numel(), "lidx0")
   k = Kernel(arch)
   k.emit(s_load_b64(s[0:1], s[0:1], soffset=NULL))
   k.emit(s_waitcnt_lgkmcnt(sdst=NULL, simm16=0))

test/backend/test_custom_kernel.py

@@ -6,8 +6,8 @@ from tinygrad.uop.ops import KernelInfo, AxisType
 # **** kernels ****
 
 def custom_arange_kernel(C:UOp) -> UOp:
-  i = UOp.range(C.size, 0)
-  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.size}"))
+  i = UOp.range(C.shape[0], 0)
+  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))
 
 def custom_eye_kernel(C:UOp) -> UOp:
   i = UOp.range(C.shape[0], 0)
@@ -16,22 +16,22 @@ def custom_eye_kernel(C:UOp) -> UOp:
 
 def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
-  assert B.size == A.size
-  i = UOp.range(A.size, 0)
-  return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.size}"))
+  assert B.numel() == A.numel()
+  i = UOp.range(A.numel(), 0)
+  return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.numel()}"))
 
 def custom_elementwise_add_kernel(C:UOp, A:UOp, B:UOp) -> UOp:
   C,A,B = C.flatten(), A.flatten(), B.flatten()
-  i = UOp.range(C.size, 0)
-  return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.size}")).simplify()
+  i = UOp.range(C.numel(), 0)
+  return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.numel()}")).simplify()
 
 def custom_elementwise_addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
   C,D,A,B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
-  assert C.size == D.size
-  i = UOp.range(C.size, 0)
+  assert C.numel() == D.numel()
+  i = UOp.range(C.numel(), 0)
   store_c = C[i].store(A[i]+B[i])
   store_d = D[i].store(A[i]*B[i])
-  return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.size}")).simplify()
+  return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.numel()}")).simplify()
 
 def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   assert A.shape[1] == B.shape[0]
@@ -291,10 +291,10 @@ class TestCustomKernel(unittest.TestCase):
 
     def custom_add_with_tmp(o1:UOp, o2:UOp, A:UOp, B:UOp) -> UOp:
       o1,o2,A,B = o1.flatten(), o2.flatten(), A.flatten(), B.flatten()
-      i = UOp.range(o1.size, 0)
+      i = UOp.range(o1.numel(), 0)
       store_o1 = o1[i].store(A[i]+B[i])
       store_o2 = o2[i].store(A[i]+B[i]+2)
-      return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.size}")).simplify()
+      return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.numel()}")).simplify()
 
     from tinygrad import function
     @function(precompile=True)

test/backend/test_graph.py

@@ -2,13 +2,12 @@ import numpy as np
 import functools, unittest, ctypes
 
 from tinygrad.device import Device, Buffer
-from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import Context, dedup, from_mv
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import Context, from_mv
 from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.engine.realize import BufferXfer, get_runner, CompiledRunner
-from tinygrad.schedule import ExecItem
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.schedule import linear_to_schedule
+from tinygrad.uop.ops import UOp, Ops, buffers
 
 from test.helpers import needs_second_gpu
 
@@ -17,77 +16,46 @@ Tensor.manual_seed(1337)
 BUF_SIZE = 4096
 RUN_CNT = 5
 
-cached_prgs = {}
-def helper_exec_op(device, outbuf, inbufs):
-  if (device, len(inbufs)) not in cached_prgs:
+# cache AST by (device, num_inputs)
+cached_asts: dict[tuple[str, int], UOp] = {}
+def get_ast(device:str, num_inputs:int) -> UOp:
+  if (device, num_inputs) not in cached_asts:
     with Context(DEBUG=0):
-      fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for i in range(len(inbufs))]
+      fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for _ in range(num_inputs)]
       s = fst[0]
-      for i in range(1, len(inbufs)): s = s.bitwise_xor(fst[i])
+      for i in range(1, num_inputs): s = s.bitwise_xor(fst[i])
+      cached_asts[(device, num_inputs)] = s.schedule()[-1].ast
+  return cached_asts[(device, num_inputs)]
 
-      si = s.schedule()[-1]
-      prg = get_runner(device, si.ast)
-    cached_prgs[(device, len(inbufs))] = prg
-
-  return ExecItem(UOp(Ops.NOOP), [outbuf] + inbufs, prg=cached_prgs[(device, len(inbufs))])
-
-def helper_copy_op(device, dest, src):
-  prg = BufferXfer(dest.nbytes, device, src.device)
-  return ExecItem(UOp(Ops.NOOP), [dest, src], prg=prg)
-
-def helper_alloc_rawbuffer(device, fill=False):
-  rawbuf = Buffer(device, BUF_SIZE, dtypes.int).ensure_allocated()
+def make_buffer(device, size=BUF_SIZE, fill=False):
+  buf = Buffer(device, size, dtypes.int).ensure_allocated()
   if fill:
     with Context(DEBUG=0):
-      data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
-  return rawbuf
+      buf.copyin(Tensor(np.random.randint(-10000, 10000, size=size, dtype=np.int32)).realize().uop.base.realized.as_memoryview())
+  return buf
 
-def helper_create_offset_rawbuffer(base, offset=0):
-  x = Buffer(base.device, base.size-offset, base.dtype, base=base, offset=offset)
-  return x.ensure_allocated()
-
-def helper_alloc_rawbuffer_sized(device, size, fill=False):
-  rawbuf = Buffer(device, size, dtypes.int).ensure_allocated()
-  if fill:
-    with Context(DEBUG=0):
-      data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
-  return rawbuf
-
-def helper_make_view(base, offset_elems, size_elems):
+def make_view(base, offset_elems, size_elems):
   return Buffer(base.device, size_elems, base.dtype, base=base, offset=offset_elems * base.dtype.itemsize).ensure_allocated()
 
-def helper_run_jit(jis, bufs, out_buffers):
-  for rawbuf in out_buffers:
-    mv = memoryview(bytearray(rawbuf.nbytes))
+def get_buf_uop(buf:Buffer, cache:dict[Buffer,UOp]) -> UOp:
+  if buf not in cache:
+    cache[buf] = u = UOp.new_buffer(buf.device, buf.size, buf.dtype)
+    buffers[u] = buf
+  return cache[buf]
+
+def make_graph(graph_cls, calls:list[UOp]):
+  linear = UOp(Ops.LINEAR, src=tuple(calls))
+  cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(linear,), arg="graph")
+  return graph_cls(cf, [])
+
+def run_schedule(calls:list[UOp]):
+  for ei in linear_to_schedule(UOp(Ops.LINEAR, src=tuple(calls))): ei.lower().run({})
+
+def zero_bufs(bufs):
+  for b in bufs:
+    mv = memoryview(bytearray(b.nbytes))
     ctypes.memset(from_mv(mv), 0, len(mv))
-    rawbuf.copyin(mv)
-
-  for ei in jis: ei.run({}, jit=True)
-  return [rawbuf.as_memoryview() for rawbuf in bufs]
-
-def helper_test_graphs(graph_impl, graphs, runs=RUN_CNT):
-  reg_ji = []
-  bufs = []
-  out_buffers = set()
-  for graph in graphs:
-    for ji in graph:
-      out_buffers.update([ji.bufs[i] for i in (ji.prg.p.outs if isinstance(ji.prg, CompiledRunner) else [0])])
-      bufs += ji.bufs
-      reg_ji.append(ji)
-  bufs = dedup(bufs)
-
-  ground_thruth_bufs = helper_run_jit(reg_ji, bufs, out_buffers)
-  ground_truth_np = [np.frombuffer(x, _to_np_dtype(bufs[i].dtype)) for i,x in enumerate(ground_thruth_bufs)]
-
-  # Build graphs
-  gr_ji = [ExecItem(UOp(Ops.NOOP), [], prg=graph_impl(None, None, graph)) for graph in graphs]
-
-  for _ in range(runs):
-    test_bufs = helper_run_jit(gr_ji, bufs, out_buffers)
-    test_bufs_np = [np.frombuffer(x, _to_np_dtype(bufs[i].dtype)) for i,x in enumerate(test_bufs)]
-    for i in range(len(ground_thruth_bufs)): np.testing.assert_equal(ground_truth_np[i], test_bufs_np[i])
+    b.copyin(mv)
 
 @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
 class TestGraph(unittest.TestCase):
@@ -101,236 +69,251 @@ class TestGraph(unittest.TestCase):
 
   def test_order_2_writes_to_same_buf(self):
     d0 = Device.DEFAULT
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(5)]
+    b = [make_buffer(d0, fill=True) for _ in range(5)]
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_exec_op(d0, b0[0], [b0[1], b0[2]]), helper_exec_op(d0, b0[0], [b0[3], b0[4]])]
+    calls = [
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()),
     ]
 
-    helper_test_graphs(Device[d0].graph, graphs)
+    zero_bufs([b[0]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[0]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
 
   def test_order_read_write_same_buf(self):
     d0 = Device.DEFAULT
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(5)]
+    b = [make_buffer(d0, fill=True) for _ in range(5)]
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_exec_op(d0, b0[0], [b0[1], b0[2]]), helper_exec_op(d0, b0[1], [b0[3], b0[4]])]
+    calls = [
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()),
     ]
 
-    helper_test_graphs(Device[d0].graph, graphs)
+    zero_bufs([b[0], b[1]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[0], b[1]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
 
   def test_order_write_read_same_buf(self):
     d0 = Device.DEFAULT
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(5)]
+    b = [make_buffer(d0, fill=True) for _ in range(5)]
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_exec_op(d0, b0[0], [b0[1], b0[2]]), helper_exec_op(d0, b0[1], [b0[0], b0[4]])]
+    calls = [
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), get_buf_uop(b[4],c), metadata=()),
     ]
 
-    helper_test_graphs(Device[d0].graph, graphs)
+    zero_bufs([b[0], b[1]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[0], b[1]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
 
   def test_order_copy_writed(self):
     self.skip_if_not_multigraph()
-
     d0 = Device.DEFAULT
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(4)]
+    b = [make_buffer(d0, fill=True) for _ in range(4)]
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_exec_op(d0, b0[0], [b0[1], b0[2]]), helper_copy_op(d0, b0[3], b0[0])]
+    calls = [
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+      UOp(Ops.COPY).call(get_buf_uop(b[3],c), get_buf_uop(b[0],c), metadata=()),
     ]
 
-    helper_test_graphs(Device[d0].graph, graphs)
+    zero_bufs([b[0], b[3]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[0], b[3]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
 
   def test_order_copy_then_read(self):
     self.skip_if_not_multigraph()
-
     d0 = Device.DEFAULT
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(4)]
+    b = [make_buffer(d0, fill=True) for _ in range(4)]
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_copy_op(d0, b0[1], b0[0]), helper_exec_op(d0, b0[3], [b0[1], b0[2]])]
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
     ]
 
-    helper_test_graphs(Device[d0].graph, graphs)
+    zero_bufs([b[1], b[3]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[1], b[3]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
 
   def test_read_write_several_graphs(self):
     d0 = Device.DEFAULT
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(8)]
+    b = [make_buffer(d0, fill=True) for _ in range(8)]
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_exec_op(d0, b0[3], [b0[1], b0[2]])],
-      [helper_exec_op(d0, b0[4], [b0[1], b0[3]])],
-      [helper_exec_op(d0, b0[5], [b0[4], b0[2]])]
-    ]
+    calls1 = [get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=())]
+    calls2 = [get_ast(d0, 2).call(get_buf_uop(b[4],c), get_buf_uop(b[1],c), get_buf_uop(b[3],c), metadata=())]
+    calls3 = [get_ast(d0, 2).call(get_buf_uop(b[5],c), get_buf_uop(b[4],c), get_buf_uop(b[2],c), metadata=())]
 
-    helper_test_graphs(Device[d0].graph, graphs)
+    out = [b[3], b[4], b[5]]
+    zero_bufs(out)
+    run_schedule(calls1 + calls2 + calls3)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
 
-    graphs = [
-      [helper_exec_op(d0, b0[3], [b0[1], b0[2]]), helper_exec_op(d0, b0[4], [b0[1], b0[2]]), helper_exec_op(d0, b0[5], [b0[1], b0[2]])],
-      [helper_exec_op(d0, b0[2], [b0[6], b0[7]])]
-    ]
-
-    helper_test_graphs(Device[d0].graph, graphs)
+    for _ in range(RUN_CNT):
+      zero_bufs(out)
+      make_graph(Device[d0].graph, calls1)([], {})
+      make_graph(Device[d0].graph, calls2)([], {})
+      make_graph(Device[d0].graph, calls3)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
 
   @needs_second_gpu
   def test_copies_2_devs(self):
     self.skip_if_not_multigraph()
-
     d0, d1 = Device.DEFAULT, f"{Device.DEFAULT}:1"
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(3)]
-    b1 = [helper_alloc_rawbuffer(d1, fill=True) for _ in range(1)]
+    b0 = [make_buffer(d0, fill=True) for _ in range(3)]
+    b1 = [make_buffer(d1, fill=True)]
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_copy_op(d0, b1[0], b0[0]), helper_exec_op(d0, b0[2], [b0[0], b0[1]])]
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(b1[0],c), get_buf_uop(b0[0],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b0[2],c), get_buf_uop(b0[0],c), get_buf_uop(b0[1],c), metadata=()),
     ]
 
-    helper_test_graphs(Device[d0].graph, graphs)
+    out = [b1[0], b0[2]]
+    zero_bufs(out)
+    run_schedule(calls)
+    expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in b0 + b1}
 
-  @needs_second_gpu
-  def test_copies_after_graph_global(self):
-    self.skip_if_not_multigraph()
-
-    d0, d1, d2, d3 = Device.DEFAULT, f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3"
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(8)]
-    b1 = [helper_alloc_rawbuffer(d1, fill=True) for _ in range(6)]
-    b2 = [helper_alloc_rawbuffer(d2, fill=True) for _ in range(6)]
-    b3 = [helper_alloc_rawbuffer(d3, fill=True) for _ in range(6)]
-
-    graphs = [
-      [helper_exec_op(d0, b0[2], [b0[0], b0[1]]), helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
-       helper_exec_op(d0, b0[5], [b0[0], b0[2]]), helper_exec_op(d0, b0[6], [b0[1], b0[2]]), helper_exec_op(d0, b0[7], [b0[0], b0[2]])],
-      [helper_copy_op(d1, b0[2], b1[0])],
-      [helper_exec_op(d0, b0[2], [b0[0], b0[1]]), helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
-       helper_exec_op(d0, b0[5], [b0[0], b0[2]]), helper_exec_op(d0, b0[6], [b0[1], b0[2]]), helper_exec_op(d0, b0[7], [b0[0], b0[2]])],
-      [helper_copy_op(d3, b0[2], b3[0])],
-    ]
-
-    helper_test_graphs(Device[d0].graph, graphs)
-
-    graphs = [
-      [helper_exec_op(d0, b0[2], [b0[0], b0[1]]), helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
-       helper_exec_op(d0, b0[5], [b0[0], b0[2]]), helper_copy_op(d0, b2[0], b0[2]), helper_copy_op(d0, b2[1], b0[5]),
-       helper_exec_op(d0, b0[7], [b0[0], b0[2]])],
-      [helper_copy_op(d1, b0[2], b1[0])],
-      [helper_exec_op(d0, b0[2], [b0[0], b0[1]])],
-      [helper_copy_op(d3, b0[2], b3[0])],
-    ]
-
-    helper_test_graphs(Device[d0].graph, graphs)
-
-    graphs = [
-      [helper_exec_op(d0, b0[2], [b0[0], b0[1]]), helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
-       helper_exec_op(d0, b0[5], [b0[0], b0[2]]), helper_copy_op(d0, b2[0], b0[2]), helper_copy_op(d0, b2[1], b0[5]),
-       helper_exec_op(d0, b0[7], [b0[0], b0[2]])],
-      [helper_copy_op(d1, b0[5], b1[0])],
-      [helper_copy_op(d3, b0[5], b3[0])],
-    ]
-
-    helper_test_graphs(Device[d0].graph, graphs)
-
-    graphs = [
-      [helper_copy_op(d1, b0[5], b1[0])],
-      [helper_copy_op(d3, b0[5], b3[0])],
-    ]
-
-    helper_test_graphs(Device[d0].graph, graphs)
-
-  @needs_second_gpu
-  def test_graph_after_copies_devs(self):
-    self.skip_if_not_multigraph()
-
-    d0, d1, d2, d3 = Device.DEFAULT, f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3"
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(8)]
-    b1 = [helper_alloc_rawbuffer(d1, fill=True) for _ in range(1)]
-    b2 = [helper_alloc_rawbuffer(d2, fill=True) for _ in range(2)]
-    b3 = [helper_alloc_rawbuffer(d3, fill=True) for _ in range(2)]
-
-    graphs = [
-      [helper_copy_op(d1, b0[0], b1[0])],
-      [helper_copy_op(d2, b0[1], b2[0]), helper_copy_op(d3, b0[2], b3[0])],
-      [helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
-       helper_exec_op(d0, b0[5], [b0[0], b0[2]])],
-    ]
-
-    helper_test_graphs(Device[d0].graph, graphs)
-
-    graphs = [
-      [helper_copy_op(d1, b0[0], b1[0])],
-      [helper_exec_op(d0, b0[2], [b0[0], b0[1]])],
-      [helper_copy_op(d2, b0[1], b2[0]), helper_copy_op(d3, b0[2], b3[0])],
-      [helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
-       helper_exec_op(d0, b0[5], [b0[0], b0[2]])],
-    ]
-
-    helper_test_graphs(Device[d0].graph, graphs)
+    for _ in range(RUN_CNT):
+      zero_bufs(out)
+      make_graph(Device[d0].graph, calls)([], {})
+      for buf in b0 + b1: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
 
   def test_graph_offset_bufs(self):
     self.skip_if_not_multigraph()
-
     d0 = Device.DEFAULT
     if not hasattr(Device[d0].allocator, "_offset"): self.skipTest("device does not support _offset")
 
-    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(1)]
-    b0 += [helper_create_offset_rawbuffer(b0[0]), helper_create_offset_rawbuffer(b0[0])]
+    b0 = make_buffer(d0, fill=True)
+    b1 = make_view(b0, 0, b0.size)
+    b2 = make_view(b0, 0, b0.size)
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_copy_op(d0, b0[0], b0[2]), helper_exec_op(d0, b0[1], [b0[0], b0[2]])],
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b1,c), get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()),
     ]
 
-    helper_test_graphs(Device[d0].graph, graphs)
+    zero_bufs([b0])
+    run_schedule(calls)
+    expected = np.frombuffer(b0.as_memoryview(), np.int32).copy()
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b0])
+      make_graph(Device[d0].graph, calls)([], {})
+      np.testing.assert_equal(expected, np.frombuffer(b0.as_memoryview(), np.int32))
 
   def test_partial_write_preserves_write_dep(self):
     self.skip_if_not_multigraph()
     self.skip_if_no_offset()
     d0 = Device.DEFAULT
 
-    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
-    copy_src_full = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
-    copy_src_lo = helper_alloc_rawbuffer(d0, fill=True)
-    v_lo = helper_make_view(base, 0, BUF_SIZE)
-    v_hi = helper_make_view(base, BUF_SIZE, BUF_SIZE)
-    a, c = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(2)]
+    base = make_buffer(d0, BUF_SIZE * 2, fill=True)
+    copy_src_full = make_buffer(d0, BUF_SIZE * 2, fill=True)
+    copy_src_lo = make_buffer(d0, fill=True)
+    v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
+    a, out = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_copy_op(d0, base, copy_src_full), helper_copy_op(d0, v_lo, copy_src_lo), helper_exec_op(d0, c, [v_hi, a])]
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()),
+      UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(out,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()),
     ]
-    helper_test_graphs(Device[d0].graph, graphs)
+
+    zero_bufs([base, out])
+    run_schedule(calls)
+    expected = {base: np.frombuffer(base.as_memoryview(), np.int32).copy(), out: np.frombuffer(out.as_memoryview(), np.int32).copy()}
+
+    for _ in range(RUN_CNT):
+      zero_bufs([base, out])
+      make_graph(Device[d0].graph, calls)([], {})
+      for buf in [base, out]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
 
   def test_partial_write_preserves_read_dep(self):
     self.skip_if_not_multigraph()
     self.skip_if_no_offset()
     d0 = Device.DEFAULT
 
-    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
-    copy_dst = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
-    copy_src_lo = helper_alloc_rawbuffer(d0, fill=True)
-    v_lo = helper_make_view(base, 0, BUF_SIZE)
-    v_hi = helper_make_view(base, BUF_SIZE, BUF_SIZE)
-    a, b = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(2)]
+    base = make_buffer(d0, BUF_SIZE * 2, fill=True)
+    copy_dst = make_buffer(d0, BUF_SIZE * 2, fill=True)
+    copy_src_lo = make_buffer(d0, fill=True)
+    v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
+    a, b = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_copy_op(d0, copy_dst, base), helper_copy_op(d0, v_lo, copy_src_lo), helper_exec_op(d0, v_hi, [a, b])]
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(copy_dst,c), get_buf_uop(base,c), metadata=()),
+      UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(v_hi,c), get_buf_uop(a,c), get_buf_uop(b,c), metadata=()),
     ]
-    helper_test_graphs(Device[d0].graph, graphs)
+
+    zero_bufs([copy_dst, base])
+    run_schedule(calls)
+    expected = {copy_dst: np.frombuffer(copy_dst.as_memoryview(), np.int32).copy(), base: np.frombuffer(base.as_memoryview(), np.int32).copy()}
+
+    for _ in range(RUN_CNT):
+      zero_bufs([copy_dst, base])
+      make_graph(Device[d0].graph, calls)([], {})
+      for buf in [copy_dst, base]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
 
   def test_middle_write_splits_write_dep(self):
     self.skip_if_not_multigraph()
     self.skip_if_no_offset()
     d0 = Device.DEFAULT
 
-    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 3, fill=True)
-    copy_src_full = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 3, fill=True)
-    copy_src_mid = helper_alloc_rawbuffer(d0, fill=True)
-    v_lo = helper_make_view(base, 0, BUF_SIZE)
-    v_mid = helper_make_view(base, BUF_SIZE, BUF_SIZE)
-    v_hi = helper_make_view(base, BUF_SIZE * 2, BUF_SIZE)
-    a, c, e = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(3)]
+    base = make_buffer(d0, BUF_SIZE * 3, fill=True)
+    copy_src_full = make_buffer(d0, BUF_SIZE * 3, fill=True)
+    copy_src_mid = make_buffer(d0, fill=True)
+    v_lo, v_mid, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE), make_view(base, BUF_SIZE * 2, BUF_SIZE)
+    a, out1, out2 = make_buffer(d0, fill=True), make_buffer(d0, fill=True), make_buffer(d0, fill=True)
+    c: dict[Buffer,UOp] = {}
 
-    graphs = [
-      [helper_copy_op(d0, base, copy_src_full), helper_copy_op(d0, v_mid, copy_src_mid),
-       helper_exec_op(d0, c, [v_lo, a]), helper_exec_op(d0, e, [v_hi, a])]
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()),
+      UOp(Ops.COPY).call(get_buf_uop(v_mid,c), get_buf_uop(copy_src_mid,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(out1,c), get_buf_uop(v_lo,c), get_buf_uop(a,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(out2,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()),
     ]
-    helper_test_graphs(Device[d0].graph, graphs)
+
+    outs = [base, out1, out2]
+    zero_bufs(outs)
+    run_schedule(calls)
+    expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in outs}
+
+    for _ in range(RUN_CNT):
+      zero_bufs(outs)
+      make_graph(Device[d0].graph, calls)([], {})
+      for buf in outs: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
 
 if __name__ == '__main__':
   unittest.main()

test/external/process_replay/process_replay.py

@@ -43,6 +43,7 @@ class ProcessReplayWarning(Warning): pass
 # *** replay the function and convert return values to string
 
 def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+  if ast.op is Ops.BEAM: ast = ast.src[0]
   # the ast.arg is non None if we are inside of search.py
   sink_arg = ast.arg or KernelInfo()
   if opts is not None: sink_arg = replace(sink_arg, opts_to_apply=tuple(opts))

test/null/test_device.py

@@ -26,8 +26,17 @@ class TestDevice(unittest.TestCase):
 
   @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
   def test_nonexistent_renderer(self):
-    with self.assertRaisesRegex(AssertionError, "No renderer"):
+    with self.assertRaisesRegex(RuntimeError, "has no renderer"):
       with Context(DEV="CPU:TYPO"): Device[Device.DEFAULT].renderer
+    with self.assertRaisesRegex(RuntimeError, "did you mean: 'CLANGJIT'"):
+      with Context(DEV="CPU:CLANG"): Device[Device.DEFAULT].renderer
+
+  @unittest.skipIf(Device.DEFAULT != "AMD", "only run on AMD")
+  def test_nonexistent_iface(self):
+    result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device[Device.DEFAULT].iface'],
+                            env={**os.environ, "DEV":"USA+AMD"}, capture_output=True)
+    self.assertNotEqual(result.returncode, 0)
+    self.assertIn(b"did you mean: 'USB'", result.stderr)
 
   def test_lowercase_canonicalizes(self):
     device = Device.DEFAULT
@@ -118,10 +127,11 @@ class TestDevVar(unittest.TestCase):
                  ("AMD:LLVM:gfx1100", Target(device="AMD", renderer="LLVM", arch="gfx1100")), ("::gfx1100", Target(arch="gfx1100")),
                  ("USB+", Target(interface="USB")), ("USB+AMD", Target(device="AMD", interface="USB")),
                  ("PCI:0+AMD", Target(device="AMD", interface="PCI", indices="0")), (":0+AMD", Target(device="AMD", indices="0")),
-                 ("PCI:0,1+AMD", Target(device="AMD", interface="PCI", indices="0,1"))]:
+                 ("PCI:0,1+AMD", Target(device="AMD", interface="PCI", indices="0,1")),
+                 ("QCOM;USB+AMD", [Target(device="QCOM"), Target(device="AMD", interface="USB")])]:
       with Context(DEV=d):
-        self.assertEqual(DEV.value, t)
-        self.assertEqual(str(DEV.value), d)
+        self.assertEqual(DEV.value, t if isinstance(t, list) else [t])
+        self.assertEqual(str(DEV), d)
 
   def test_target(self):
     with Context(DEV="CPU"): self.assertEqual(DEV.target("CPU"), Target("CPU"))
@@ -129,6 +139,10 @@ class TestDevVar(unittest.TestCase):
     with Context(DEV=":LLVM"): self.assertEqual(DEV.target("CPU"), Target("CPU", "LLVM"))
     with Context(DEV="AMD:LLVM"): self.assertEqual(DEV.target("CPU"), Target("CPU"))
     with Context(DEV=""): self.assertEqual(DEV.target("CPU"), Target("CPU"))
+    with Context(DEV="QCOM:IR3;AMD:LLVM"):
+      self.assertEqual(DEV.target("QCOM"), Target("QCOM", "IR3"))
+      self.assertEqual(DEV.target("AMD"), Target("AMD", "LLVM"))
+      self.assertEqual(DEV.target("CPU"), Target("CPU"))
 
   def test_dev_arch_override(self):
     with Context(DEV="NULL:HIP:gfx1100"):

test/null/test_process_replay.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, Device
+from tinygrad import Tensor, Device, Context
 from tinygrad.engine.realize import get_program
 from tinygrad.codegen.opt import Opt, OptOps
 from test.external.process_replay.process_replay import replay_get_program
@@ -30,5 +30,12 @@ class TestProcessReplay(unittest.TestCase):
     good, compare, _ = replay_get_program(p, self.ast, self.renderer, opts=opts)
     self.assertEqual(good, compare)
 
+  @Context(BEAM=1)
+  def test_beam(self):
+    si = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1]
+    p = get_program(si.ast, self.renderer)
+    good, compare, _ = replay_get_program(p, self.ast, self.renderer)
+    self.assertEqual(good, compare)
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/null/test_schedule.py

@@ -146,7 +146,7 @@ class TestSchedule(unittest.TestCase):
   def test_create_schedule_handles_multi_kernel_after_and_after_deps(self):
     def named_copy(name:str):
       def fxn(out:UOp, src:UOp) -> UOp:
-        i = UOp.range(src.size, 0)
+        i = UOp.range(src.shape[0], 0)
         return out[i].store(src[i]).end(i).sink(arg=KernelInfo(name=name))
       return fxn
 

test/null/test_symbolic_tensor.py

@@ -97,5 +97,16 @@ class TestSymbolicShrink(unittest.TestCase):
     t = Tensor.rand(3, 5).shrink(((0, 2), (vi, vi+1)))
     assert t.shape == (2, 1)
 
+class TestSymbolicContiguousViewOffset(unittest.TestCase):
+  def test_shrink_from_start(self):
+    v = Variable("v", 1, 10).bind(5)
+    t = Tensor.rand(10).realize().shrink(((0, v),))
+    self.assertEqual(t.uop.contiguous_view_offset(), 0)
+
+  def test_shrink_with_offset(self):
+    v = Variable("v", 1, 7).bind(4)
+    t = Tensor.rand(10).realize().shrink(((3, 3+v),))
+    self.assertEqual(t.uop.contiguous_view_offset(), 3)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_function.py

@@ -418,8 +418,8 @@ class TestFunctionTuple(unittest.TestCase):
 
   def test_custom_kernel_save_unused_output(self):
     def my_kernel(C:UOp, D:UOp, A:UOp) -> UOp:
-      i = UOp.range(A.size, 0)
-      j = UOp.range(D.size, 1)
+      i = UOp.range(A.shape[0], 0)
+      j = UOp.range(D.shape[0], 1)
       store_c = C[i].store(A[i] * 2.0).end(i)
       store_d = D[j].store(A[j]).end(j)
       return UOp.group(store_c, store_d).sink(arg=KernelInfo(name="my_kernel"))
@@ -444,7 +444,7 @@ class TestFunctionTuple(unittest.TestCase):
 
   def test_custom_kernel_both_outputs_used(self):
     def my_kernel(C:UOp, D:UOp, A:UOp) -> UOp:
-      i = UOp.range(A.size, 0)
+      i = UOp.range(A.shape[0], 0)
       store_c = C[i].store(A[i] * 2.0)
       store_d = D[i].store(A[i] * 3.0)
       return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name="my_kernel"))

tinygrad/callify.py

@@ -94,13 +94,13 @@ def contiguous_mops_to_view(c:UOp, src:UOp):
   if (view := _make_buffer_view(src)) is None: return None
   return view.contiguous(tag=c.tag)
 
-def transform_precompiled_call(c:UOp) -> UOp|None:
+def transform_precompiled_function(c:UOp) -> UOp|None:
   if not c.arg.precompile: return None
   if c.src[0].op is Ops.SINK: return None
-  assert c.src[0].op is Ops.TUPLE, f"expected TUPLE body for precompiled call, got {c.src[0].op}"
+  assert c.src[0].op is Ops.TUPLE, f"expected TUPLE body for precompiled function, got {c.src[0].op}"
   input_buffers = tuple(x.contiguous() if x.op not in {Ops.AFTER, Ops.BIND} else x for x in c.src[1:])
 
-  # add the outputs to the call
+  # add the outputs to the function
   srcs = c.src[0].src
   resolved = [c.gettuple(i) for i in range(len(srcs))]
   outs = tuple(_buffer_like(r) for r in resolved)
@@ -108,7 +108,7 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
   fxn = UOp.sink(*[t.after(t.store(s)) for t,s in zip(targets, srcs)])
 
   # create the new thing for the big graph
-  new_call = c.replace(src=(fxn, *input_buffers, *outs), tag=None)
+  new_call = c.replace(op=Ops.CALL, src=(fxn, *input_buffers, *outs), tag=None)
   rets = tuple(o.after(new_call) for o in outs)
 
   # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
@@ -119,8 +119,8 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
 
 # NOTE: adding rules to here is bad. these all need to run before the schedule cache
 pm_early_transform_tensor_graph = PatternMatcher([
-  # transform precompiled CALLs
-  (UPat(Ops.CALL, name="c"), transform_precompiled_call),
+  # transform precompiled FUNCTIONs -> CALLs
+  (UPat(Ops.FUNCTION, name="c"), transform_precompiled_function),
 
   # resolve TUPLE+GETTUPLE (for precompiled calls)
   (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),

tinygrad/device.py

@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import BENCHMARKS, CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
-from tinygrad.helpers import select_first_inited, DEV, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target
+from tinygrad.helpers import select_by_name, select_first_inited, DEV, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target
 from tinygrad.dtype import DType, PtrDType, dtypes, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
@@ -292,9 +292,8 @@ class Compiled:
     assert (rn:=next((self._renderer_name(r) for r in self.renderers if getenv(f"{self.device}_{self._renderer_name(r)}")), None)) is None, \
       f"{self.device}_{rn}=1 is deprecated, use DEV={self.device}:{rn} or {self.device}_CC={rn} instead"
     t = DEV.target(self.device.split(':')[0], **({"arch":self.arch} if self.arch else {}))
-    renderers = [r for r in self.renderers if self._renderer_name(r) == t.renderer] if t.renderer else self.renderers
-    assert renderers, f"No renderer for {self.device} " + (f"matches request {t.renderer!r}" if t.renderer else "is available")
-    return select_first_inited(renderers, f"No renderer for {self.device} is available", self.cached_renderer, target=t)
+    return select_first_inited(select_by_name(self.renderers, self._renderer_name, t.renderer, f"{self.device} has no renderer {t.renderer!r}"),
+                               f"No renderer for {self.device} is available", self.cached_renderer, target=t)
 
   def synchronize(self):
     """

tinygrad/gradient.py

@@ -95,15 +95,15 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
     if t0 not in grads or grads[t0].op is Ops.NOOP: continue
     # GETTUPLE: accumulate gradient into a TUPLE UOp on the CALL, process when we hit the CALL
     if t0.op is Ops.GETTUPLE:
-      k = t0.src[0]  # the CALL
-      assert k.op is Ops.CALL and k.src[0].op is Ops.TUPLE
+      k = t0.src[0]  # the FUNCTION
+      assert k.op is Ops.FUNCTION and k.src[0].op is Ops.TUPLE
       n_outputs = len(k.src[0].src)
       prev = grads[k].src if k in grads else tuple(UOp(Ops.NOOP) for _ in range(n_outputs))
       grads[k] = UOp.maketuple(*(prev[i] + grads[t0] if i == t0.arg and prev[i].op is not Ops.NOOP else
                                  grads[t0] if i == t0.arg else prev[i] for i in range(n_outputs)))
       continue
-    # CALL: pass needed param set so backward only computes required gradients
-    if t0.op is Ops.CALL:
+    # FUNCTION: pass needed param set so backward only computes required gradients
+    if t0.op is Ops.FUNCTION:
       needed = {i for i, arg in enumerate(t0.src[1:]) if arg in targets or in_target_path.get(arg, False)}
       lgrads:tuple[UOp|None, ...]|None = call_gradient(grads[t0], t0, needed)
     else:

tinygrad/helpers.py

@@ -3,7 +3,7 @@ import time
 START_TIME = time.perf_counter()
 import os, functools, platform, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
 from collections import defaultdict
-import subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
+import subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
 from dataclasses import dataclass, field, replace
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload
 
@@ -121,6 +121,11 @@ def suppress_finalizing(func):
       if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
   return wrapper
 
+def select_by_name(candidates:Sequence[T], get_name:Callable[...,str], query:str, err_msg:str) -> list[T]:
+  if len(ret:=[c for c in candidates if not query or get_name(c) == query]) == 0:
+    raise RuntimeError(err_msg + (f", did you mean: {m[0]!r}?" if (m:=difflib.get_close_matches(query, map(get_name, candidates))) else ""))
+  return ret
+
 def select_first_inited(candidates:Sequence[Callable[...,T]], err_msg:str, cache:dict|None=None, **kwargs):
   excs = []
   for typ in candidates:
@@ -130,7 +135,7 @@ def select_first_inited(candidates:Sequence[Callable[...,T]], err_msg:str, cache
       if cache is not None: cache[typ] = x
       return x
     except Exception as e: excs.append(e)
-  raise excs[0] if len(excs) == 1 else ExceptionGroup(err_msg, excs)
+  raise excs[0] if len(excs) == 1 else ExceptionGroup(err_msg + " is available", excs)
 
 def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's')
 
@@ -202,20 +207,19 @@ class Target:
   def replacedefault(self, **kwargs) -> Target: return replace(self, **{k:v for k,v in kwargs.items() if not getattr(self, k)})
 
 class _DEV(ContextVar):
-  _value = Target()
+  _value: list[Target] = [Target()]
   @property
-  def value(self) -> Target: return self._value
+  def value(self) -> list[Target]: return self._value
   @value.setter
-  def value(self, v:str|Target): self._value = v if isinstance(v, Target) else Target.parse(v)
-  def __getattr__(self, k): return getattr(self.value, k)
+  def value(self, v:str|Target|list[Target]):
+    self._value = v if isinstance(v, list) else [v] if isinstance(v, Target) else [Target.parse(t) for t in v.split(';')]
+  def __repr__(self) -> str: return ";".join([repr(t) for t in self._value])
+  def __getattr__(self, k): return getattr(self._value[0], k)
   # get target for device string, kwargs are passed if not already specified
   def target(self, dev:str, **kwargs) -> Target:
-    t = self.value.replacedefault(**kwargs) if self.device == dev or not self.device else Target(device=dev, **kwargs)
-    # TODO: remove this once DEV supports secondary targets
-    if (cv:=ContextVar._cache.get(f"{dev}_CC", None)) is not None and cv.value:
-      assert not t.renderer, f"renderer set in DEV and {dev}_CC"
-      return replace(t, renderer=cv.value.upper())
-    return replace(t, device=dev)
+    assert (v:=getenv(k:=f"{dev}_CC", "")) == "", \
+      f"{k}={v} is deprecated, use DEV='{';'.join([repr(t) for t in self._value if t.device != dev] + [f'{dev}:{v}'])}' instead"
+    return replace(next((t for t in self._value if not t.device or t.device == dev), Target(device=dev)).replacedefault(**kwargs), device=dev)
 
 DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
 IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
@@ -233,10 +237,7 @@ MAX_KERNEL_BUFFERS = ContextVar("MAX_KERNEL_BUFFERS", 0)
 EMULATED_DTYPES = ContextVar("EMULATED_DTYPES", "")
 CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0)
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
-# Compilers
-CPU_CC, NV_CC, CUDA_CC, NULL_CC = ContextVar("CPU_CC", ""), ContextVar("NV_CC", ""), ContextVar("CUDA_CC", ""), ContextVar("NULL_CC", "")
 NULL_ALLOW_COPYOUT = ContextVar("NULL_ALLOW_COPYOUT", 0)
-AMD_CC, QCOM_CC = ContextVar("AMD_CC", ""), ContextVar("QCOM_CC", "")
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
 VIZ = ContextVar("VIZ", 0)
 PROFILE = ContextVar("PROFILE", abs(VIZ.value))

tinygrad/runtime/autogen/__init__.py

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

tinygrad/runtime/autogen/mesa.py

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

tinygrad/runtime/support/hcq.py

@@ -4,8 +4,8 @@ import contextlib, decimal, statistics, time, ctypes, array, os, struct, collect
 from dataclasses import replace
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
-from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, select_first_inited, unwrap, suppress_finalizing
-from tinygrad.helpers import TracingKey
+from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, select_first_inited, select_by_name, unwrap
+from tinygrad.helpers import suppress_finalizing, TracingKey
 from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
@@ -61,7 +61,7 @@ if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.mockgpu import MockFileIOInterf
 # **************** for HCQ Compatible Devices ****************
 
 def hcq_filter_visible_devices(devs, device):
-  assert (v:=getenv("HCQ_VISIBLE_DEVICES", "")) == "", f"HCQ_VISIBLE_DEVICES={v} is deprecated, use DEV={replace(DEV.value, indices=v)} instead"
+  assert (v:=getenv("HCQ_VISIBLE_DEVICES", "")) == "", f"HCQ_VISIBLE_DEVICES={v} is deprecated, use DEV={DEV.target(device, indices=v)} instead"
   return [devs[x] for x in ids] if (ids:=[int(x) for x in DEV.target(device).indices.split(',') if x.strip()]) else devs
 
 SignalType = TypeVar('SignalType', bound='HCQSignal')
@@ -489,9 +489,9 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   def _select_iface(self, *ifaces:Type):
     assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
       f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
-    if (iface:=DEV.target(dev:=type(self).__name__[:-6]).interface): ifaces = tuple(x for x in ifaces if x.__name__.startswith(iface.upper()))
-    assert len(ifaces), f"No interface for {dev} " + (f"matches request {iface!r}" if iface else "is available")
-    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in ifaces],
+    t = DEV.target(dev:=type(self).__name__[:-6])
+    filtered = select_by_name(ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
+    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered],
                                f"No interface for {dev}:{self.device_id} is available")
 
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"

tinygrad/schedule/multi.py

@@ -1,5 +1,5 @@
 from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST
-from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite, should_resolve_call
+from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite, KernelInfo
 from tinygrad.dtype import dtypes
 from tinygrad.schedule.allreduce import handle_allreduce
 
@@ -116,6 +116,15 @@ def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0
 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)
 
+# TODO: this is all junk
+
+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 in {Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.CUSTOM_FUNCTION}: return False
+  if c.arg.precompile: return False
+  return True
+
 def rewrite_into_call(call:UOp):
   if not should_resolve_call(call): return None
   new_body = graph_rewrite(call.src[0], multi_pm, name="subcall")

tinygrad/schedule/rangeify.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field, replace
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, AddrSpace, Invalid
 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, should_resolve_call, identity_element
+from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element
 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, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
@@ -126,8 +126,7 @@ 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=True) -> UOp|None:
-  if not should_resolve_call(c): return None
+def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None:
   params: list[UOp] = []
   graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params, name="gather params")
   params = sorted(params, key=lambda x: x.arg)
@@ -150,8 +149,8 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.COPY, src=(UPat.var("s"), UPat.var("d"))),
    lambda s,d: s.substitute({UOp(Ops.DEVICE, arg=s.device):d}) if s.base.op is Ops.CONST else None),
 
-  # resolve calls
-  (UPat(Ops.CALL, name="c"), resolve_call),
+  # resolve functions
+  (UPat(Ops.FUNCTION, name="c"), resolve_function),
 
   # resolve TUPLE+GETTUPLE
   (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),

tinygrad/tensor.py

@@ -222,7 +222,7 @@ class Tensor(OpMixin):
       param = UOp.param(slot, self.dtype, self.shape, self.device)
     return Tensor(param)
   def call(self, *lst:Tensor, fxn:Tensor|UOp, grad_fxn:Callable|None=None) -> Tensor:
-    fret = (fxn.uop if isinstance(fxn, Tensor) else fxn).call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn)
+    fret = (fxn.uop if isinstance(fxn, Tensor) else fxn).function(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn)
     return Tensor(fret.gettuple(0))
 
   def custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None) -> list[Tensor]:

tinygrad/uop/__init__.py

@@ -26,7 +26,7 @@ class Ops(FastEnum):
 
   # uops that aren't rendered
   NOOP = auto(); REWRITE_ERROR = auto()
-  PARAM = auto(); CALL = auto()
+  PARAM = auto(); CALL = auto(); FUNCTION = auto()
 
   # renderer
   # LINEAR is a list of UOps, SOURCE has a str arg that's human readable, BINARY has bytes arg that's compiled

tinygrad/uop/ops.py

@@ -215,17 +215,17 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       # late ops don't have shape
       case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.STORE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
            Ops.VECTORIZE | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | \
-           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE:
+           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION:
         return None
 
       case Ops.GETTUPLE:
-        # GETTUPLE extracts from a TUPLE (possibly through a CALL)
-        in_tuple = self.src[0].src[0] if self.src[0].op is Ops.CALL else self.src[0]
+        # GETTUPLE extracts from a TUPLE (possibly through a FUNCTION)
+        in_tuple = self.src[0].src[0] if self.src[0].op is Ops.FUNCTION else self.src[0]
         assert in_tuple.op is Ops.TUPLE
         inner_shape = in_tuple.src[self.arg]._shape
         if inner_shape is None: return None
-        # if through a CALL, substitute internal PARAMs in the shape with corresponding args
-        if self.src[0].op is Ops.CALL:
+        # if through a FUNCTION, substitute internal PARAMs in the shape with corresponding args
+        if self.src[0].op is Ops.FUNCTION:
           return tuple(graph_rewrite(s, _pm_resolve_params, self.src[0].src[1:], walk=True) if isinstance(s, UOp) else s for s in inner_shape)
         return inner_shape
 
@@ -262,8 +262,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END:
         return self.src[0]._shape
 
-      case Ops.CALL: return None
-
       # TODO: disallow shape changing bitcast
       case Ops.BITCAST:
         ps = self.src[0]._shape
@@ -421,8 +419,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def maketuple(*srcs:UOp):  # pylint: disable=no-self-argument
     return UOp(Ops.TUPLE, dtypes.void, srcs)
   def gettuple(self, idx:int) -> UOp:
-    in_tuple = self.src[0] if self.op is Ops.CALL else self
-    assert in_tuple.op is Ops.TUPLE, f"gettuple requires CALL or TUPLE source, got {self.op}"
+    in_tuple = self.src[0] if self.op is Ops.FUNCTION else self
+    assert in_tuple.op is Ops.TUPLE, f"gettuple requires FUNCTION or TUPLE source, got {self.op}"
     return UOp(Ops.GETTUPLE, in_tuple.src[idx].dtype, (self,), idx)
   def group(*srcs:UOp|None):  # pylint: disable=no-self-argument
     if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
@@ -690,9 +688,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     """If movement ops on a BUFFER collapse to a contiguous range, return `offset` in elements. Otherwise None."""
     from tinygrad.schedule.rangeify import pm_mops
     from tinygrad.uop.symbolic import symbolic
-    out = graph_rewrite(self._mop(Ops.RESHAPE, (self.size,)).index(UOp.range(self.size, 0)), pm_mops+symbolic, name="contiguous_view_offset")
+    numel = self.numel()
+    out = graph_rewrite(self.flatten().index(UOp.range(numel, 0)), pm_mops+symbolic, name="contiguous_view_offset")
     if out.op is not Ops.INDEX: return None
-    if out.src[1].op is Ops.CONST and self.size == 1:
+    if out.src[1].op is Ops.CONST and resolve(numel == 1, False):
       if not isinstance(out.src[1].arg, int): return None  # masked/padded regions produce InvalidType
       return out.src[1].arg
     if out.src[1].op is Ops.RANGE: return 0
@@ -731,10 +730,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       buf = self.src[0].buffer
       if isinstance(buf, MultiBuffer):
         mbuf = MultiBuffer.__new__(MultiBuffer)
-        mbuf.bufs = [b.view(self.size, self.dtype, self.arg[1] * self.dtype.itemsize) for b in buf.bufs]
+        mbuf.bufs = [b.view(self.arg[0], self.dtype, self.arg[1] * self.dtype.itemsize) for b in buf.bufs]
         return mbuf
       assert isinstance(buf, Buffer), "must be a Buffer for BUFFER_VIEW"
-      return buf.view(self.size, self.dtype, self.arg[1] * self.dtype.itemsize)
+      return buf.view(self.arg[0], self.dtype, self.arg[1] * self.dtype.itemsize)
     if self.op is Ops.MSELECT:
       ret = self.src[0].buffer
       assert isinstance(ret, MultiBuffer)
@@ -748,8 +747,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     assert self.src[0].op is Ops.UNIQUE, f"buffer src[0] must be UNIQUE, not {self.src[0].op}"
     if (cret:=buffers.get(self)) is not None: return cret
     rdtype = self.dtype if isinstance(self.dtype, ImageDType) else self.dtype.base
-    if isinstance(self.device, tuple): ret = MultiBuffer(self.device, self.size, rdtype).ref(1)
-    else: ret = Buffer(self.device, self.size, rdtype).ref(1)
+    if isinstance(self.device, tuple): ret = MultiBuffer(self.device, self.arg, rdtype).ref(1)
+    else: ret = Buffer(self.device, self.arg, rdtype).ref(1)
     buffers[self] = ret
     return ret
   @property
@@ -940,6 +939,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # value-producing bodies are always wrapped in TUPLE so CALL dtype is always void
     body = self if self.op in UOp._NO_TUPLE_WRAP else UOp.maketuple(self)
     return UOp(Ops.CALL, dtypes.void, (body,)+srcs, CallInfo(grad_fxn, metadata, name, precompile, precompile_backward))
+  def function(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(),
+               name:str|None=None, precompile:bool=False, precompile_backward:bool=False) -> UOp:
+    assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
+    return UOp(Ops.FUNCTION, dtypes.void, (UOp.maketuple(self),)+srcs, CallInfo(grad_fxn, metadata, name, precompile, precompile_backward))
   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)]
@@ -1563,7 +1566,7 @@ pm_pyrender_extra = PatternMatcher([
   (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
   (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
   (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
-    f"UOp.new_buffer({repr(d.arg)}, {x.size}, {x.dtype}, {u.arg})"),
+    f"UOp.new_buffer({repr(d.arg)}, {x.arg}, {x.dtype}, {u.arg})"),
   (UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
   (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda ctx,x: f"UOp(Ops.CUSTOM_FUNCTION, {x.dtype}, src={srcs(ctx, x.src)}, arg={x.arg!r})"),
   (UPat(Ops.REDUCE_AXIS, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {r.arg[1]})"),

tinygrad/uop/spec.py

@@ -134,12 +134,13 @@ _tensor_spec = PatternMatcher([
 
   # allow CALL/PARAM/CUSTOM_FUNCTION — CALL dtype is always void
   (UPat(Ops.CALL, dtypes.void), lambda: True),
+  (UPat(Ops.FUNCTION, dtypes.void), lambda: True),
   (UPat(Ops.PARAM), lambda: True),
   (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),
 
   # TUPLE must have void dtype, GETTUPLE can only appear on CALL or TUPLE
   (UPat(Ops.TUPLE, dtypes.void), lambda: True),
-  (UPat(Ops.GETTUPLE, src=(UPat((Ops.CALL, Ops.TUPLE)),), name="g"), lambda g: isinstance(g.arg, int)),
+  (UPat(Ops.GETTUPLE, src=(UPat((Ops.FUNCTION, Ops.TUPLE)),), name="g"), lambda g: isinstance(g.arg, int)),
 
   # ** for custom kernels **
 

tinygrad/viz/js/worker.js

@@ -54,7 +54,7 @@ const layoutUOp = (g, { graph, change }, opts) => {
       width = Math.max(width, ctx.measureText(line).width);
       height += lineHeight;
     }
-    const callNode = label.startsWith("CALL\n");
+    const callNode = label.startsWith("CALL\n") || label.startsWith("FUNCTION\n");
     if (callNode) callCount++;
     g.setNode(k, {...rectDims(width, height), label, ref, id:k, color, tag, callNode});
     // add edges

tinygrad/viz/serve.py

@@ -50,7 +50,8 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
                Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
                Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
-               Ops.CALL: "#00B7C8", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.LINEAR: "#7DF4FF", Ops.BINARY: "#404040",
+               Ops.FUNCTION: "#C07788", Ops.CALL: "#00B7C8",
+               Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.LINEAR: "#7DF4FF", Ops.BINARY: "#404040",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
                Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e", Ops.AFTER: "#8A7866", Ops.END: "#524C46"}
 
@@ -136,7 +137,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
         label += f"\n({multirange_str(rngs, color=True)})"
       if u._shape is not None:
         label += f"\n{shape_to_str(u.shape)}"
-      if u.op is Ops.CALL:
+      if u.op in {Ops.CALL, Ops.FUNCTION}:
         label += f"\n{u.src[0].key.hex()[:8]}"
       if u.op in {Ops.INDEX, Ops.BUFFERIZE}:
         if len(u.toposort()) < 30: label += f"\n{u.render()}"
@@ -340,7 +341,7 @@ def load_amd_counters(data:VizData, profile:list[ProfileEvent]) -> None:
     run_number[k] += 1
     steps:list[dict] = []
     if (pmc:=v.get(ProfilePMCEvent)):
-      steps.append(create_step("PMC", ("/prg-pmc", len(data.ctxs), len(steps)), pmc))
+      steps.append(create_step("PMC", ("/prg-pmc", len(data.ctxs), len(steps)), pmc[0]))
       all_counters[(name, run_number[k], pname)] = pmc[0]
     # to decode a SQTT trace, we need the raw stream, program binary and device properties
     if (sqtt:=v.get(ProfileSQTTEvent)):
@@ -647,7 +648,7 @@ def get_render(viz_data:VizData, query:str) -> dict:
       ret["rows"].append((name, durations[k][n-1], *[r[1] for r in pmc_table["rows"]]))
     ret["cols"] = ["Kernel", "Duration", *ret["cols"]]
     return ret
-  if fmt == "prg-pmc": return unpack_pmc(data[0])
+  if fmt == "prg-pmc": return unpack_pmc(data)
   if fmt.startswith("prg-pkts"):
     ret = {}
     with soft_err(lambda err:ret.update(err)):