CALL with return value is FUNCTION

* fix symbolic contiguous_view_offset

* flatten

docs/abstractions3.py

@@ -38,7 +38,7 @@ optim.schedule_step()   # this will step the optimizer without running realize
 # The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
 # l1.uop and l2.uop define a computation graph
 
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 schedule: List[ExecItem] = Tensor.schedule(l1, l2)
 
 print(f"The schedule contains {len(schedule)} items.")

docs/developer/developer.md

@@ -17,9 +17,9 @@ The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not al
 
 ## Scheduling
 
-The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
+The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
 
-::: tinygrad.engine.schedule.ExecItem
+::: tinygrad.schedule.ExecItem
 
 ## Lowering
 

examples/llm.c/export.py

@@ -5,7 +5,7 @@ from tinygrad import Device, nn, Tensor, dtypes
 from train_gpt2 import GPT, GPTConfig
 from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name
 from tinygrad.engine.realize import get_kernel
-from tinygrad.engine.memory import memory_planner
+from tinygrad.schedule.memory import memory_planner
 from tinygrad.uop.ops import Ops
 
 DEV.value = "CPU"

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_arange.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.ptx import PTXRenderer
 from test.helpers import needs_second_gpu

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.engine.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/backend/test_opt_gemm.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, Device
 from tinygrad.helpers import get_single_element
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 
 class TestOptGemm(unittest.TestCase):
   @classmethod

test/backend/test_quantize_onnx.py

@@ -6,7 +6,7 @@ from tinygrad import Tensor, Context, Device, dtypes
 from tinygrad.uop.ops import Ops
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 
 N = 512
 

test/backend/test_uops.py

@@ -8,7 +8,7 @@ from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.engine.realize import CompiledRunner, get_program, get_runner
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer

test/external/external_test_amd.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Device, Tensor
-from tinygrad.engine.schedule import create_schedule
+from tinygrad.schedule import create_schedule
 from tinygrad.runtime.ops_amd import AMDDevice
 
 class TestAMD(unittest.TestCase):

test/external/external_test_hip_compile.py

@@ -2,7 +2,7 @@ import time, unittest
 from tinygrad.runtime.support.hip_comgr import compile_hip
 from tinygrad import Tensor
 from tinygrad.device import Device
-from tinygrad.engine.schedule import create_schedule
+from tinygrad.schedule import create_schedule
 from tinygrad.codegen.opt.kernel import Kernel
 
 class TestHIPCompileSpeed(unittest.TestCase):

test/external/external_test_opt.py

@@ -7,7 +7,7 @@ from tinygrad import GlobalCounters, Tensor, Device
 from tinygrad.helpers import getenv
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import capturing, run_schedule
-from tinygrad.engine.schedule import linear_to_schedule
+from tinygrad.schedule import linear_to_schedule
 from tinygrad.tensor import _to_np_dtype
 
 class CLCache:

test/external/external_uop_gc.py

@@ -1,6 +1,6 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
-from tinygrad.engine.schedule import schedule_cache
+from tinygrad.schedule import schedule_cache
 from tinygrad.engine.realize import method_cache, get_program
 from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
 from tinygrad.uop.divandmod import fold_divmod_general

test/external/fuzz_graph.py

@@ -5,7 +5,7 @@ from tinygrad.helpers import Context, getenv, from_mv
 from tinygrad.dtype import dtypes
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import BufferXfer, get_runner
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.jit import apply_graph_to_jit
 

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_memory_planner.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.memory import memory_plan_rewrite
+from tinygrad.schedule.memory import memory_plan_rewrite
 
 global_map = {}
 held_bufs: set[UOp] = set()

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_schedule_cache.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Variable, Context
 from tinygrad.helpers import cpu_events
-from tinygrad.engine.schedule import schedule_cache
+from tinygrad.schedule import schedule_cache
 
 def schedule_one():
   Tensor([1]).schedule()

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/null/test_tensor_metadata.py

@@ -2,7 +2,7 @@ import unittest
 from tinygrad import Tensor, dtypes
 from tinygrad.tensor import _METADATA
 from tinygrad.engine.realize import capturing
-from tinygrad.engine.schedule import linear_to_schedule
+from tinygrad.schedule import linear_to_schedule
 from tinygrad.helpers import Context
 
 @unittest.skip("tensor metadata is no longer supported")

test/testextra/test_tk.py

@@ -3,7 +3,7 @@ import unittest, math, time
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.realize import get_runner
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 from tinygrad.engine.jit import TinyJit
 import numpy as np
 

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"))

test/unit/test_llm_server.py

@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import patch
 from tinygrad import Tensor, UOp
-from tinygrad.engine.schedule import schedule_cache
+from tinygrad.schedule import schedule_cache
 from tinygrad.apps.llm import Transformer, TransformerConfig
 
 TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,

test/unit/test_schedule_cache.py

@@ -2,7 +2,7 @@ import unittest
 import functools
 from tinygrad import Tensor, Variable, UOp
 from tinygrad.uop.ops import KernelInfo
-from tinygrad.engine.schedule import schedule_cache
+from tinygrad.schedule import schedule_cache
 
 def custom_set0_kernel(A:UOp, num:int) -> UOp:
   return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))

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/engine/jit.py

@@ -6,8 +6,8 @@ from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
 from tinygrad.dtype import DType, dtypes
 from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite
 from tinygrad.engine.realize import ExecItem, capturing, BufferCopy, BufferXfer, EncDec, CompiledRunner, Runner, Estimates
-from tinygrad.engine.memory import memory_plan_rewrite, _collect_bufs
-from tinygrad.engine.schedule import linear_to_schedule
+from tinygrad.schedule.memory import memory_plan_rewrite, _collect_bufs
+from tinygrad.schedule import linear_to_schedule
 from tinygrad.nn.state import get_parameters
 from tinygrad.schedule.rangeify import mop_cleanup
 from dataclasses import dataclass

tinygrad/engine/schedule.py

@@ -1,182 +0,0 @@
-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, KernelInfo
-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, flatten, BEAM, partition
-from tinygrad.engine.realize import ExecItem
-
-# **** schedule linearizer
-
-# unwrap VIEW/CAST/etc to find the actual data source (kernel output, buffer, or multi-device op)
-def _unwrap_src(s: UOp) -> UOp:
-  while len(s.src) and s.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND}: s = s.src[0]
-  return s
-
-def _split_after(after: UOp) -> tuple[tuple[UOp, ...], tuple[UOp, ...]]:
-  kernels, remaining = partition(after.src[1:], lambda s: s.op in {Ops.CALL, Ops.END})
-  deps, remaining = partition(remaining, lambda s: s.op is Ops.AFTER)
-  if invalid := [s for s in remaining if s.op is not Ops.STORE]:
-    raise AssertionError(f"AFTER source should be CALL, END, STORE, or AFTER, not {invalid[0].op}")
-  return tuple(kernels), tuple(deps)
-
-def create_schedule(sched_sink:UOp) -> UOp:
-  with cpu_profile(TracingKey("toposort sched_sink")):
-    # build kernel dependency graph: edges from producer kernel to consumer kernels
-    children: dict[UOp, list[UOp]] = {}
-    in_degree: dict[UOp, int] = {}
-    for u in sched_sink.toposort(gate_kernel_sink):
-      if u.op is not Ops.AFTER: continue
-      kernels, after_deps = _split_after(u)
-      for k in kernels:
-        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}"
-        kernel_deps = k.src[0].src[1:] if k.op is Ops.END else k.src[1:]
-        for s in kernel_deps + after_deps:
-          match (s := _unwrap_src(s)).op:
-            case Ops.AFTER:
-              for t in _split_after(s)[0]:
-                children.setdefault(t, []).append(k)
-                in_degree[k] += 1
-            case Ops.MSELECT | Ops.MSTACK:
-              for ss in s.src:
-                if ss.op is Ops.MSELECT: ss = ss.src[0]
-                if ss.op not in {Ops.BUFFER, Ops.PARAM}:
-                  assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
-                  for t in _split_after(ss)[0]:
-                    children.setdefault(t, []).append(k)
-                    in_degree[k] += 1
-            case Ops.BUFFER | Ops.PARAM | Ops.BIND:
-              pass  # BUFFER/PARAM is already realized, BIND is a bound variable (not a buffer dependency)
-            case _:
-              raise RuntimeError(f"input to kernel must be AFTER, BUFFER, PARAM, MSELECT, MSTACK, or BIND, not {s.op}")
-
-  with cpu_profile(TracingKey("linearize schedule")):
-    queue: deque[UOp] = deque(k for k,v in in_degree.items() if v == 0)
-    linearized: list[UOp] = []
-    while len(queue):
-      rk = queue.popleft()
-      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))
-
-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:]
-    # create subbuffers if needed
-    if ast.op is Ops.BUFFER_VIEW:
-      base = buf_uops[1].buffer
-      assert isinstance(base, Buffer), "base can't be MultiBuffer"
-      buffers[buf_uops[0]] = base.view(buf_uops[0].arg, ast.dtype, ast.arg[1]*base.dtype.itemsize)
-    # wrap SINK with BEAM UOp when beam search is enabled
-    if ast.op is Ops.SINK and BEAM >= 1: ast = UOp(Ops.BEAM, src=(ast,), arg=BEAM.value)
-    ubufs = [b.buffer for b in buf_uops if b.op is not Ops.BIND]
-    metadata = si.arg.metadata
-    if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph":
-      schedule.append(ExecItem(ast, flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for b in ubufs]), metadata))
-    elif any(isinstance(x, MultiBuffer) for x in ubufs):
-      assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
-      dnums = [x for x in ast.variables() if x.expr == '_device_num']
-      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, cast(list[Buffer|None], ubufs), metadata))
-  return schedule
-
-from tinygrad.engine.memory import memory_plan_rewrite
-from tinygrad.engine.realize import capturing
-from tinygrad.schedule.rangeify import get_kernel_graph
-from tinygrad.helpers import CAPTURING
-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),
-])
-
-pm_resolve_linear_call = PatternMatcher([
-  # call LINEAR is resolved here
-  (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), lambda linear_call:
-   graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")),
-  # LINEAR on LINEAR
-  (UPat(Ops.LINEAR, custom_early_reject={Ops.LINEAR}, name="x"),
-   lambda x: x.replace(src=tuple(flatten(x.src if x.op is Ops.LINEAR else (x,) for x in x.src)))),
-])
-
-schedule_cache: dict[bytes, UOp] = {}
-# ctx is just for DEBUG on inner
-def lower_sink_to_linear(function:UOp) -> UOp|None:
-  st = time.perf_counter()
-  if isinstance(function.arg, KernelInfo): return None
-  cache_key = function.key
-  if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
-    if SPEC: type_verify(function, tensor_spec)
-    # support recursive CALLs
-    linear = create_schedule(get_kernel_graph(function))
-    if SCACHE: schedule_cache[cache_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(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'} {cache_key.hex()[:8]}"+\
-          f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
-  return linear
-
-pm_schedule = PatternMatcher([
-  (UPat(Ops.SINK, name="function"), lower_sink_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_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True)
-
-  # this recursively resolves the linear_call and allocates buffers
-  linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call")
-
-  # 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
-      if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}")
-      var_vals[nm] = val
-
-  # jit captures this schedule, no need to execute.
-  if len(capturing) and CAPTURING:
-    capturing[0].add_linear(linear, var_vals)
-    return [], var_vals
-
-  held_bufs = ({b for b in linear_call.src[1:] if b.op is Ops.BUFFER} if linear_call.op is Ops.CALL else set())
-  linear = memory_plan_rewrite(linear, held_bufs)
-
-  # convert LINEAR to ExecItems
-  schedule: list[ExecItem] = linear_to_schedule(linear)
-  return schedule, var_vals

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/__init__.py

@@ -0,0 +1,182 @@
+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, KernelInfo
+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, flatten, BEAM, partition
+from tinygrad.engine.realize import ExecItem
+
+# **** schedule linearizer
+
+# unwrap VIEW/CAST/etc to find the actual data source (kernel output, buffer, or multi-device op)
+def _unwrap_src(s: UOp) -> UOp:
+  while len(s.src) and s.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND}: s = s.src[0]
+  return s
+
+def _split_after(after: UOp) -> tuple[tuple[UOp, ...], tuple[UOp, ...]]:
+  kernels, remaining = partition(after.src[1:], lambda s: s.op in {Ops.CALL, Ops.END})
+  deps, remaining = partition(remaining, lambda s: s.op is Ops.AFTER)
+  if invalid := [s for s in remaining if s.op is not Ops.STORE]:
+    raise AssertionError(f"AFTER source should be CALL, END, STORE, or AFTER, not {invalid[0].op}")
+  return tuple(kernels), tuple(deps)
+
+def create_schedule(sched_sink:UOp) -> UOp:
+  with cpu_profile(TracingKey("toposort sched_sink")):
+    # build kernel dependency graph: edges from producer kernel to consumer kernels
+    children: dict[UOp, list[UOp]] = {}
+    in_degree: dict[UOp, int] = {}
+    for u in sched_sink.toposort(gate_kernel_sink):
+      if u.op is not Ops.AFTER: continue
+      kernels, after_deps = _split_after(u)
+      for k in kernels:
+        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}"
+        kernel_deps = k.src[0].src[1:] if k.op is Ops.END else k.src[1:]
+        for s in kernel_deps + after_deps:
+          match (s := _unwrap_src(s)).op:
+            case Ops.AFTER:
+              for t in _split_after(s)[0]:
+                children.setdefault(t, []).append(k)
+                in_degree[k] += 1
+            case Ops.MSELECT | Ops.MSTACK:
+              for ss in s.src:
+                if ss.op is Ops.MSELECT: ss = ss.src[0]
+                if ss.op not in {Ops.BUFFER, Ops.PARAM}:
+                  assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
+                  for t in _split_after(ss)[0]:
+                    children.setdefault(t, []).append(k)
+                    in_degree[k] += 1
+            case Ops.BUFFER | Ops.PARAM | Ops.BIND:
+              pass  # BUFFER/PARAM is already realized, BIND is a bound variable (not a buffer dependency)
+            case _:
+              raise RuntimeError(f"input to kernel must be AFTER, BUFFER, PARAM, MSELECT, MSTACK, or BIND, not {s.op}")
+
+  with cpu_profile(TracingKey("linearize schedule")):
+    queue: deque[UOp] = deque(k for k,v in in_degree.items() if v == 0)
+    linearized: list[UOp] = []
+    while len(queue):
+      rk = queue.popleft()
+      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))
+
+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:]
+    # create subbuffers if needed
+    if ast.op is Ops.BUFFER_VIEW:
+      base = buf_uops[1].buffer
+      assert isinstance(base, Buffer), "base can't be MultiBuffer"
+      buffers[buf_uops[0]] = base.view(buf_uops[0].arg, ast.dtype, ast.arg[1]*base.dtype.itemsize)
+    # wrap SINK with BEAM UOp when beam search is enabled
+    if ast.op is Ops.SINK and BEAM >= 1: ast = UOp(Ops.BEAM, src=(ast,), arg=BEAM.value)
+    ubufs = [b.buffer for b in buf_uops if b.op is not Ops.BIND]
+    metadata = si.arg.metadata
+    if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph":
+      schedule.append(ExecItem(ast, flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for b in ubufs]), metadata))
+    elif any(isinstance(x, MultiBuffer) for x in ubufs):
+      assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
+      dnums = [x for x in ast.variables() if x.expr == '_device_num']
+      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, cast(list[Buffer|None], ubufs), metadata))
+  return schedule
+
+from tinygrad.schedule.memory import memory_plan_rewrite
+from tinygrad.engine.realize import capturing
+from tinygrad.schedule.rangeify import get_kernel_graph
+from tinygrad.helpers import CAPTURING
+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),
+])
+
+pm_resolve_linear_call = PatternMatcher([
+  # call LINEAR is resolved here
+  (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), lambda linear_call:
+   graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")),
+  # LINEAR on LINEAR
+  (UPat(Ops.LINEAR, custom_early_reject={Ops.LINEAR}, name="x"),
+   lambda x: x.replace(src=tuple(flatten(x.src if x.op is Ops.LINEAR else (x,) for x in x.src)))),
+])
+
+schedule_cache: dict[bytes, UOp] = {}
+# ctx is just for DEBUG on inner
+def lower_sink_to_linear(function:UOp) -> UOp|None:
+  st = time.perf_counter()
+  if isinstance(function.arg, KernelInfo): return None
+  cache_key = function.key
+  if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
+    if SPEC: type_verify(function, tensor_spec)
+    # support recursive CALLs
+    linear = create_schedule(get_kernel_graph(function))
+    if SCACHE: schedule_cache[cache_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(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'} {cache_key.hex()[:8]}"+\
+          f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
+  return linear
+
+pm_schedule = PatternMatcher([
+  (UPat(Ops.SINK, name="function"), lower_sink_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_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True)
+
+  # this recursively resolves the linear_call and allocates buffers
+  linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call")
+
+  # 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
+      if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}")
+      var_vals[nm] = val
+
+  # jit captures this schedule, no need to execute.
+  if len(capturing) and CAPTURING:
+    capturing[0].add_linear(linear, var_vals)
+    return [], var_vals
+
+  held_bufs = ({b for b in linear_call.src[1:] if b.op is Ops.BUFFER} if linear_call.op is Ops.CALL else set())
+  linear = memory_plan_rewrite(linear, held_bufs)
+
+  # convert LINEAR to ExecItems
+  schedule: list[ExecItem] = linear_to_schedule(linear)
+  return schedule, var_vals

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

@@ -13,10 +13,10 @@ from tinygrad.gradient import compute_gradient
 from tinygrad.mixin import OpMixin, ReductionStr
 from tinygrad.uop.ops import smax, UOp, Ops, sint, all_metadata, _index_to_concrete_int, sint_to_uop, Variable
 from tinygrad.uop.ops import _broadcast_shape
-from tinygrad.engine.schedule import ExecItem, complete_create_schedule_with_vars
+from tinygrad.schedule import ExecItem, complete_create_schedule_with_vars
 from tinygrad.device import Buffer, canonicalize_device
 from tinygrad.engine.realize import run_schedule
-from tinygrad.engine.callify import transform_to_call
+from tinygrad.callify import transform_to_call
 
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
@@ -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)):