Merge branch 'master' into remove_programspec

* system: use addr space

* am: xgmi

* fix

* ugh

docs/developer/layout.md

@@ -26,9 +26,9 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi
 
 ## tinygrad/codegen
 
-Transform the optimized ast into a linearized list of UOps.
+Transform the optimized ast into a linearized and rendered program.
 
-::: tinygrad.codegen.full_rewrite
+::: tinygrad.codegen.full_rewrite_to_program
     options:
         members: false
         show_labels: false

extra/export_model.py

@@ -1,12 +1,11 @@
 from typing import Tuple, Dict, List, Optional
 from tinygrad.dtype import DType
-from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context, to_mv
+from tinygrad.helpers import Context, to_mv, to_function_name
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import Ops, UOp
 import json
 from collections import OrderedDict
 
@@ -15,8 +14,13 @@ EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
 def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
   functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
   for ji in run.jit_cache:
-    fxn: ProgramSpec = ji.prg.p
-    functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
+    prg: UOp = ji.prg.p
+    name = prg.src[0].arg.name
+    function_name = to_function_name(name)
+    src = prg.src[3].arg
+    global_size, local_size = prg.sizes
+    prg_vars = prg.variables()
+    functions[function_name] = src   # NOTE: this assumes all with the same name are the same
     cargs = []
     for i,arg in enumerate(ji.bufs):
       key = id(arg)
@@ -28,8 +32,8 @@ def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str]
           bufnum += 1
           if i > 0: bufs_to_save[bufs[key][0]] = arg   # if first usage of a buffer is not an output, and it's not a special name
       cargs.append(bufs[key][0])
-    cargs += [var for var in fxn.vars if getattr(var, "op", None) is Ops.DEFINE_VAR] # symbolic vars; is it necessary or sufficient to check for DEFINE_VAR?
-    statements.append((fxn.function_name, cargs, fxn.global_size, fxn.local_size))
+    cargs += [var for var in prg_vars if getattr(var, "op", None) is Ops.DEFINE_VAR] # symbolic vars; is it necessary or sufficient to check for DEFINE_VAR?
+    statements.append((function_name, cargs, list(global_size) if global_size else None, list(local_size) if local_size else None))
 
   return functions, statements, {name:(size, dtype, key) for (name,size,dtype,key) in bufs.values()}, bufs_to_save
 

extra/gemm/triton_nv_matmul.py

@@ -4,7 +4,8 @@ import triton.language as tl
 from triton.compiler import AttrsDescriptor, ASTSource, compile as triton_compile
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.engine.realize import CompiledRunner, ExecItem, ProgramSpec
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.engine.realize import CompiledRunner, ExecItem
 from tinygrad.helpers import getenv
 np.set_printoptions(suppress=True)
 
@@ -85,9 +86,12 @@ if __name__ == "__main__":
   # remove debug sections
   src = src.split("\t.file")[0]
   assert '.extern .shared' not in src
-  prg = ProgramSpec("matmul_kernel", src, device=Device.DEFAULT,
-                global_size=[M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1], local_size=[32*compiled.metadata.num_warps, 1, 1],
-                mem_estimate=A.nbytes() + B.nbytes() + C.nbytes())
+  # Create linearized uops with SPECIAL for global/local sizes
+  global_size = [M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1]
+  local_size = [32*compiled.metadata.num_warps, 1, 1]
+  uops = [UOp(Ops.SPECIAL, arg=('g', i), src=(UOp.const(dtypes.int, global_size[i]),)) for i in range(3)]
+  uops += [UOp(Ops.SPECIAL, arg=('l', i), src=(UOp.const(dtypes.int, local_size[i]),)) for i in range(3)]
+  prg = UOp.new_program("matmul_kernel", src, Device.DEFAULT, si.ast, uops)
   ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
   tflops = []
   for i in range(5):

extra/remu/test/hwtest.py

@@ -1,10 +1,10 @@
 import numpy as np
 import unittest
 import subprocess, struct, math
-from tinygrad import Tensor, dtypes, Device, UOp
+from tinygrad import Tensor, dtypes, Device
+from tinygrad.uop.ops import UOp, Ops
 from tinygrad.helpers import getenv
 from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
-from tinygrad.renderer import ProgramSpec
 from tinygrad.engine.realize import CompiledRunner
 
 def get_output(asm:str, n_threads:int=1):
@@ -22,7 +22,9 @@ def get_output(asm:str, n_threads:int=1):
     *(data0_1+l) = res;
   }}"""
   t = Tensor.zeros(n_threads, dtype=dtypes.uint32).contiguous().realize()
-  prg = ProgramSpec("test", src, Device.DEFAULT, UOp.sink(t), global_size=[1, 1, 1], local_size=[n_threads, 1, 1])
+  # Create linearized uops with SPECIAL for local size
+  uops = [UOp(Ops.SPECIAL, arg=('l', 0), src=(UOp.const(dtypes.int, n_threads),))]
+  prg = UOp.new_program("test", src, Device.DEFAULT, UOp.sink(t.uop), uops)
   car = CompiledRunner(prg)
   if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
   car([t.uop.buffer], {}, wait=True)

opencode.json

@@ -1 +1 @@
-{"$schema": "https://opencode.ai/config.json", "formatter": false}
+{"$schema": "https://opencode.ai/config.json", "formatter": false, "lsp": false}

test/external/external_benchmark_op_conv.py

@@ -1,10 +1,8 @@
 # ruff: noqa: E501 E712 F401
 from tinygrad import dtypes, Device
 from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
-from tinygrad.codegen import full_rewrite
 from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
-from tinygrad.renderer import ProgramSpec
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
 from tinygrad.dtype import ImageDType, Invalid
@@ -86,16 +84,13 @@ def dm_conv_172():
 
 ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()
 
-compiler = Device.default.compiler
 renderer = Device.default.renderer
 allocator = Device.default.allocator
 
-uops = full_rewrite(ast, renderer)
-src = renderer.render(uops)
-
-lib = compiler.compile(src)
-ps = ProgramSpec("conv", src, Device.DEFAULT, ast, uops)
-cr = CompiledRunner(ps, precompiled=lib)
+ps = get_program(ast, renderer)
+# update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
+ps = ps.replace(src=(ps.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *ps.src[2:]))
+cr = CompiledRunner(ps)
 
 gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
 # print(len(gs))

test/external/external_test_am.py

@@ -2,7 +2,7 @@ import unittest
 from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_GMC
 from tinygrad.runtime.support.hcq import MMIOInterface
-from tinygrad.runtime.support.memory import PageTableTraverseContext
+from tinygrad.runtime.support.memory import PageTableTraverseContext, AddrSpace
 from tinygrad.runtime.autogen.am import am
 from tinygrad.helpers import mv_address
 
@@ -70,7 +70,7 @@ class TestAMPageTable(unittest.TestCase):
 
     for va,sz in [(0x10000, 0x3000), (0x11000, 0x300000), (0x10000, 0x2000), (0x11000, 0x5000),
                   (0x2000000, 0x2000), (0x4000000, 0x4000000), (0x38000, 0x303000), (0x8000, 0x1000)]:
-      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)])
+      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, helper_va(va))
       results = list(ctx.next(sz))
@@ -102,8 +102,8 @@ class TestAMPageTable(unittest.TestCase):
     mm0 = self.d[0].mm
 
     for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
-      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)])
-      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)])
+      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)], aspace=AddrSpace.PHYS)
+      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)], aspace=AddrSpace.PHYS)
       mm0.unmap_range(helper_va(va2), sz2)
       mm0.unmap_range(helper_va(va1), sz1)
 
@@ -112,24 +112,24 @@ class TestAMPageTable(unittest.TestCase):
 
     for va,sz in [(0x10000, 0x3000), (0x1000000, 0x1000000), (0x12000, 0x4000)]:
       exteranl_va = helper_va(va)
-      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])
+      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)])
+        mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)])
+        mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)])
+        mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)])
+        mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
 
       mm0.unmap_range(vaddr=exteranl_va, size=sz)
 
       # Finally can map and check paddrs
-      mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)])
+      mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)], aspace=AddrSpace.PHYS)
 
       ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
       for tup in ctx.next(0x100000):
@@ -147,13 +147,13 @@ class TestAMPageTable(unittest.TestCase):
     with self.assertRaises(AssertionError):
       mm0.unmap_range(helper_va(0x10000), 0x3000)
 
-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
     mm0.unmap_range(helper_va(0x10000), 0x3000)
 
     with self.assertRaises(AssertionError):
       mm0.unmap_range(helper_va(0x10000), 0x3000)
 
-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
     mm0.unmap_range(helper_va(0x10000), 0x3000)
 
     with self.assertRaises(AssertionError):
@@ -164,16 +164,16 @@ class TestAMPageTable(unittest.TestCase):
 
     # offset from start
     for off in [0, 0x3000, 0x10000]:
-      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
+      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
       mm0.unmap_range(helper_va(0x1000000) + off, (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
       mm0.unmap_range(helper_va(0x1000000), 2 << 20)
 
     # offset from end
     for off in [0x1000, 0x20000]:
-      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
+      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
       mm0.unmap_range(helper_va(0x1000000), (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
       mm0.unmap_range(helper_va(0x1000000), 2 << 20)
 
   def test_frag_size(self):

test/external/process_replay/process_replay.py

@@ -9,7 +9,7 @@ if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
 
 try:
   from tinygrad.schedule.rangeify import get_rangeify_map
-  from tinygrad.renderer import Renderer, ProgramSpec
+  from tinygrad.renderer import Renderer
   from tinygrad.engine.realize import get_program
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
   from tinygrad.codegen.opt import Opt
@@ -51,15 +51,20 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
     return "\n".join([f"{len(asts)} kernels", *asts])
   return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)
 
-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+def replay_get_program(p:UOp, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
   # the ast.arg is non None if we are inside of search.py
-  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
-  input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
+  # p is a PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
+  p_name = p.src[0].arg.name
+  p_applied_opts = p.src[0].arg.applied_opts
+  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p_applied_opts if BEAM>=1 else None)
+  input_ast = ast.replace(arg=replace(sink_arg, name=p_name))
   p2 = get_program(input_ast, renderer=renderer)
-  def to_str(ret:ProgramSpec) -> str:
+  def to_str(ret:UOp) -> str:
     # PYTHON renderer pickles UOps, first unpickle and decode here
-    if p.device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret.src))])
-    return ret.src
+    ret_src = ret.src[3].arg
+    ret_device = ret.device
+    if ret_device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret_src))])
+    return ret_src
   # properly color the name arg
   ast_repr = codecs.decode(str(input_ast), "unicode_escape")
   return to_str(p2), to_str(p), (ast_repr, renderer)

test/helpers.py

@@ -2,18 +2,27 @@ import os, time, struct, functools, unittest
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.engine.realize import Runner
+from tinygrad.engine.realize import Runner, get_program
 from tinygrad.dtype import DType
 from tinygrad.nn.state import get_parameters
 from tinygrad.helpers import T, CI
-from tinygrad.codegen import full_rewrite
+from tinygrad.renderer import Renderer
+from tinygrad.codegen import full_rewrite_to_sink, line_rewrite, pm_linearize_cleanups
+from tinygrad.codegen.late.linearizer import linearize
 
 # decorator to skip slow tests by default, run with RUN_SLOW=1 to include them
 slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run")
 from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
 
+def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
+  """Extract linearized UOps from a sink. Test helper that only does linearization (no render)."""
+  if ren is None: ren = Renderer()
+  if sink.arg is None: sink = sink.replace(arg=KernelInfo())
+  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
+  return line_rewrite(linearize(full_sink), pm_linearize_cleanups)
+
 def derandomize_model(model):
   for p in get_parameters(model):
     p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype))
@@ -51,13 +60,12 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
   bufs = []
   for buf_dt, data in inputs or []:
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
-    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + buf_dt.fmt, *data)))
+    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp(Ops.DEFINE_GLOBAL, uop.dtype.ptr(), arg=0, src=())
-  opts = PythonRenderer()
-  lst = full_rewrite(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), opts)
-  prog = PythonProgram("run", PythonCompiler().compile(opts.render(lst)))
+  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer())
+  prog = PythonProgram("run", PythonCompiler().compile(prg.src[3].arg))  # source code is in src[3].arg
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
-  return out_buf.cast(uop.dtype.fmt).tolist()[0]
+  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 
 def not_support_multi_device():
   # CL and CUDA don't support multi device if in CI

test/opt/test_tensor_cores.py

@@ -1,10 +1,9 @@
 import numpy as np
 import unittest
-from dataclasses import replace
 
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import DType
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import AMX, AMD_LLVM, CPU_LLVM, Context
@@ -26,14 +25,14 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d
   opts_to_apply = [Opt(OptOps.TC, axis, (tc_select, tc_opt, 1))]
 
   if ensure_triggered:
-    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, Device.DEFAULT, opts=opts_to_apply)
+    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
     wmmas = len([uop for uop in program.uops if uop.op is Ops.WMMA])
     tcs = len([x for x in program.applied_opts if x.op is OptOps.TC])
     assert wmmas > 0, "tensor core not triggered"
     assert tcs == 1, "tensor core opt not included"
   else:
     try:
-      program = get_program(realized_ast, Device[Device.DEFAULT].renderer, Device.DEFAULT, opts=opts_to_apply)
+      program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
       assert False, "OptOps.TC triggered, expected KernelOptError"
     except KernelOptError: pass
 
@@ -44,7 +43,10 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
   if dtype_in == dtypes.bfloat16: r = r.float()
   realized_ast, bufs = helper_realized_ast(r)
   opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
-  prg = CompiledRunner(replace(get_program(realized_ast, Device[Device.DEFAULT].renderer, Device.DEFAULT, opts=opts), device=Device.DEFAULT))
+  p = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts)
+  # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
+  p = p.replace(src=(p.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *p.src[2:]))
+  prg = CompiledRunner(p)
   if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
   assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
   prg.exec(bufs)

test/test_fusion_op.py

@@ -28,7 +28,7 @@ class TestFusionOp(unittest.TestCase):
     sched = a.schedule()
     sched[-1].lower()
     self.assertLess(time.perf_counter()-st, 2.0)
-    assert len(sched[-1].prg.p.src.splitlines()) < 250
+    assert len(sched[-1].prg.p.src[3].arg.splitlines()) < 250  # source code is in src[3].arg
 
   def test_recursive_add_cmp(self):
     st = time.perf_counter()

test/test_linearizer.py

@@ -1,6 +1,5 @@
 import numpy as np
 import unittest
-from dataclasses import replace
 
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.gpudims import get_grouped_dims
@@ -168,8 +167,8 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
   def test_upcast_with_locals_cpu(self):
     out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)]).uops
-    self.assertEqual(len(prg.src.split("for")), 5)
+    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)])
+    self.assertEqual(len(prg.src[3].arg.split("for")), 5)  # source code is in src[3].arg
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -517,7 +516,11 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
   device = real_bufs[0].device
   wanna_output = [np.array(x).flatten() for x in wanna_output]
 
-  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts), device=device))
+  def get_prg(opts):
+    prg = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts)
+    # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
+    prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=device), *prg.src[2:]))
+    return CompiledRunner(prg)
 
   def check_opt(opts):
     prg = get_prg(opts=opts)

test/test_opts.py

@@ -14,7 +14,7 @@ class TestOpts(unittest.TestCase):
     self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
     if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP:
       prg = get_program(s[-1].ast, renderer=Device[Device.DEFAULT].renderer)
-      self.assertIn('float4', prg.src)
+      self.assertIn('float4', prg.src[3].arg)  # source code is in src[3].arg
 
 if __name__ == '__main__':
   unittest.main()

test/test_renderer_failures.py

@@ -2,30 +2,24 @@ import unittest
 import numpy as np
 from tinygrad.device import Buffer, Device, is_dtype_supported
 from tinygrad.dtype import dtypes, ConstType
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import dedup, flatten, prod
+from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.helpers import prod
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.runtime.ops_python import PythonRenderer
 from tinygrad.uop.ops import UOp, Ops, python_alu
-from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.codegen import full_rewrite
 
-def _test_uop_result(inputs:list[Tensor], stores:list[UOp], local_size=None):
+def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
   for x in inputs: x.realize()
-  # NOTE: we only toposort the stores
-  uops: list[UOp] = []
-  def _recursive_add(uop:UOp) -> list[UOp]: return flatten([_recursive_add(x) for x in uop.src])+[uop]
-  uops = dedup(flatten(_recursive_add(st) for st in stores))
+  uops = prg.uops
   outbufs = [Buffer(Device.DEFAULT, sz:=(1 if local_size is None else prod(local_size)), (dtype:=u.src[1].dtype), \
       initial_value=np.zeros(sz, dtype=_to_np_dtype(dtype)).data) for u in uops if u.op is Ops.STORE]
   inbufs = [x.uop.base.buffer for x in inputs]
-  src = Device[Device.DEFAULT].renderer.render(uops)
-  lib = Device[Device.DEFAULT].compiler.compile_cached(src)
-  ei = CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test",
-                                  src, Device.DEFAULT, uops[-1], uops=uops, lib=lib, local_size=local_size))
+  # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
+  prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *prg.src[2:]))
+  ei = CompiledRunner(prg)
   ei.exec(outbufs+inbufs)
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
 
@@ -38,8 +32,8 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   alu = ld.alu(alu_op, *alu_src_uops)
   store = UOp.store(a.index(idx), alu)
   sink = UOp(Ops.SINK, dtypes.void, (store,))
-  uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
-  return _test_uop_result([Tensor([input_val])], uops)[0]
+  prg = get_program(sink, Device[Device.DEFAULT].renderer)
+  return _test_uop_result([Tensor([input_val])], prg)[0]
 
 class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
@@ -48,8 +42,8 @@ class TestRendererFailures(unittest.TestCase):
     gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
-    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
+    prg = get_program(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
 
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
@@ -59,8 +53,8 @@ class TestRendererFailures(unittest.TestCase):
     gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
-    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], uops, local_size=[4, 2, 1])[0]
+    prg = get_program(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], prg, local_size=[4, 2, 1])[0]
     np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
 
 @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage), "uops are for cstyle")
@@ -77,7 +71,7 @@ class TestCStyleFailures(unittest.TestCase):
     schedule = ret.schedule()
     assert len(schedule) == 1
     schedule[0].lower()
-    src = schedule[0].prg.p.src
+    src = schedule[0].prg.p.src[3].arg  # source code is in src[3].arg
     self.assertEqual("("*5 not in src, should_strip_paren)
 
   def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
@@ -105,8 +99,8 @@ class TestPTXFailures(unittest.TestCase):
     if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, if_uop), val))
     sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
-    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
+    prg = get_program(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")

test/test_tensor.py

@@ -10,7 +10,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.codegen import full_rewrite
+from tinygrad.engine.realize import get_program
 from tinygrad.dtype import DType
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
@@ -869,9 +869,8 @@ class TestIdxUpcast(unittest.TestCase):
     for s in schedule:
       if s.ast.op is Ops.SINK:
         renderer = Device[s.bufs[0].device].renderer
-        uops = full_rewrite(s.ast, renderer)
-        renderer.render(uops)
-        return uops
+        prg = get_program(s.ast, renderer)
+        return prg.uops
 
   def _assert(self, dtype: DType, a: Tensor):
     uops = self._schedule_render(a)

test/test_uops.py

@@ -10,22 +10,25 @@ from tinygrad.uop.spec import shared_spec
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.engine.realize import CompiledRunner, get_program, get_runner
 from tinygrad.engine.schedule import ExecItem
-from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
+from test.helpers import get_uops
 
 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   sink = UOp.group(*u)
   for r in sink.ranges: sink = sink.end(r)
   # we strip the SINK here for legacy reasons
-  ret = full_rewrite(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
+  ret = get_uops(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
   assert ret[-1].op is Ops.SINK
   return ret[:-1]
 
 def _uops_to_prg(uops_list):
-  return CompiledRunner(get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer, Device.DEFAULT))
+  prg = get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer)
+  # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
+  prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *prg.src[2:]))
+  return CompiledRunner(prg)
 
 def uop(uops:list[UOp], uop:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
   uops.append(UOp(uop, dtype, tuple(src), arg))

test/test_uops_stats.py

@@ -2,9 +2,7 @@ import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import getenv, GlobalCounters, EMULATE
 from tinygrad.engine.realize import get_program
-from tinygrad.renderer import ProgramSpec
 from tinygrad.renderer import Estimates
-from tinygrad.codegen import full_rewrite
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
@@ -146,7 +144,7 @@ class TestUOpsStats(unittest.TestCase):
     u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
     u4 = UOp(Ops.MUL, dtypes.int, (u1,u2))
     u5 = UOp(Ops.ADD, dtypes.int, (u4,u3))
-    uops = full_rewrite(u5.sink())
+    uops = list(u5.toposort())
 
     globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
     o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
@@ -155,7 +153,7 @@ class TestUOpsStats(unittest.TestCase):
     u2 = globl.index(o2)
     u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
     u4 = UOp(Ops.MULACC, dtypes.int, (u1,u2,u3))
-    uops_fma = full_rewrite(u4.sink())
+    uops_fma = list(u4.toposort())
 
     self.assertEqual(flops_mem(uops), flops_mem(uops_fma))
 
@@ -167,24 +165,26 @@ class TestStatsOptimized(unittest.TestCase):
     cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1].ast
     cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule()[-1].ast
 
-  def check_gemm(self, p:ProgramSpec, extra_flops=0):
+  def check_gemm(self, p:UOp, extra_flops=0):
     #p.uops.print()
-    #print(p.src)
-    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
-    self.assertEqual(p.estimates.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
-    self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats
+    #print(p.src[3].arg)
+    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+    print(p.src[0].arg.name, estimates.ops, estimates.mem, estimates.lds)
+    self.assertEqual(estimates.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
+    self.assertEqual(estimates.mem, 3*N*N*4) # 3 NxN mats with floats
 
   def test_gemm(self):
     p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[])
     self.check_gemm(p)
-    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)
+    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+    self.assertEqual(estimates.lds, 2*N*N*N*4 + 4*N*N)
 
   def test_gemm_tc_unroll(self):
     try:
       p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
     except KernelOptError:
       raise unittest.SkipTest("no tensor cores")
-    print(p.src)
+    print(p.src[3].arg)
     self.check_gemm(p)
 
   # this is a good lesson about why UPCASTing is a good idea
@@ -192,13 +192,15 @@ class TestStatsOptimized(unittest.TestCase):
   def test_gemm_one_upcasted(self):
     p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4)])
     self.check_gemm(p)
-    self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
+    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+    self.assertEqual(estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted(self):
     p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer,
                     opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
     self.check_gemm(p)
-    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
+    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+    self.assertEqual(estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted_locals(self):
     try:
@@ -207,7 +209,8 @@ class TestStatsOptimized(unittest.TestCase):
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     self.check_gemm(p)
-    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
+    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+    self.assertEqual(estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_group(self):
     try:
@@ -217,13 +220,15 @@ class TestStatsOptimized(unittest.TestCase):
     SZ = N*N*4
     # NOTE: these are sort of wrong. they aren't honoring the IF statement
     self.check_gemm(p, extra_flops=SZ*4)
-    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
+    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+    self.assertEqual(estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
 
   def test_reduce(self):
     p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[])
-    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
-    self.assertEqual(p.estimates.ops, N*N)
-    self.assertEqual(p.estimates.mem, N*N*4 + 4)
+    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+    print(p.src[0].arg.name, estimates.ops, estimates.mem, estimates.lds)
+    self.assertEqual(estimates.ops, N*N)
+    self.assertEqual(estimates.mem, N*N*4 + 4)
 
   def test_reduce_group(self):
     try:
@@ -231,7 +236,8 @@ class TestStatsOptimized(unittest.TestCase):
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     # NOTE: these are wrong, they don't respect the if statement
-    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
+    estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+    print(p.src[0].arg.name, estimates.ops, estimates.mem, estimates.lds)
 
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/testextra/test_cfg_viz.py

@@ -3,7 +3,6 @@ import textwrap
 
 from tinygrad import Device, Tensor
 from tinygrad.uop.ops import UOp, Ops, track_rewrites
-from tinygrad.renderer import ProgramSpec
 from tinygrad.helpers import TracingKey
 from tinygrad.engine.realize import ExecItem, CompiledRunner
 
@@ -51,9 +50,9 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
 
-@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, ret=ret))
-def run_asm(name:str, src:str) -> ProgramSpec:
-  prg = ProgramSpec(name, template.replace("fn_name", name).replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK))
+@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.src[0].arg.name, ret=ret))
+def run_asm(name:str, src:str) -> UOp:
+  prg = UOp.new_program(name, template.replace("fn_name", name).replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK), [])
   ei = ExecItem(UOp(Ops.SINK), [Tensor.empty(1).uop.buffer.ensure_allocated()], prg=CompiledRunner(prg))
   ei.run()
   return prg

test/unit/test_dtype_spec.py

@@ -8,7 +8,7 @@ from hypothesis import given, settings, strategies as strat
 import numpy as np
 import torch
 
-settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.register_profile("my_profile", max_examples=50, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
 
 core_dtypes = list(DTYPES_DICT.values())

test/unit/test_llm_server.py

@@ -27,8 +27,8 @@ class TestLLMServer(unittest.TestCase):
     from tinygrad.apps.llm import Handler
     from tinygrad.helpers import TCPServerWithReuse
 
-    cls.port = 11435
-    cls.server = TCPServerWithReuse(('127.0.0.1', cls.port), Handler)
+    cls.server = TCPServerWithReuse(('127.0.0.1', 0), Handler)
+    cls.port = cls.server.server_address[1]
     cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
     cls.server_thread.start()
     time.sleep(0.1)

test/unit/test_uop_symbolic.py

@@ -3,8 +3,8 @@ import unittest, pickle, functools, math
 import z3
 
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
-from tinygrad.codegen import full_rewrite
 from tinygrad.helpers import Context
+from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid
 from tinygrad.uop.validate import uops_to_z3
@@ -747,7 +747,7 @@ class TestSymbolic(unittest.TestCase):
 
     # TODO: copied from render, render does not support cast
     glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
-    uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
+    uops = get_uops(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
     rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
 
     self.assertEqual(rewritten_uop, cond.where(a.cast(dtypes.half), b.cast(dtypes.half)))

tinygrad/apps/llm.py

@@ -11,8 +11,8 @@ class SimpleTokenizer:
     self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
 
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
-    # TODO: ucat_range is slow
-    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
+    # 0x323b0 is one past the max codepoint in unicode categories L/N/Z (0x323af is max L)
+    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
     r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
     self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
       f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")

tinygrad/codegen/__init__.py

@@ -18,7 +18,6 @@ from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_split_store
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
-from tinygrad.device import Compiler
 
 pm_syntactic_sugar = PatternMatcher([
   # INDEX on ptr INDEX concats them
@@ -129,39 +128,16 @@ def do_linearize(prg:UOp, sink:UOp) -> UOp:
   if SPEC: type_verify(lst, program_spec)
   return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),))
 
-def do_render(ctx:tuple[Renderer, Compiler], prg:UOp, lin:UOp) -> UOp:
-  src = ctx[0].render(list(lin.src))
+def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
+  src = ctx.render(list(lin.src))
   return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),))
 
-def do_compile(ctx:tuple[Renderer, Compiler], prg:UOp, src:UOp) -> UOp:
-  lib = ctx[1].compile_cached(src.arg)
-  return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),))
-
 pm_to_program = PatternMatcher([
-  (UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"),), name="prg"), do_linearize),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR, name="lin")), name="prg"), do_render),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(Ops.SOURCE, name="src")), name="prg"), do_compile),
+  (UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"), UPat(Ops.DEVICE)), name="prg"), do_linearize),
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR, name="lin")), name="prg"), do_render),
 ])
 
-def full_rewrite_to_program(sink:UOp, ren:Renderer, compiler:Compiler) -> UOp:
+def full_rewrite_to_program(sink:UOp, ren:Renderer) -> UOp:
   full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  sink = UOp(Ops.PROGRAM, src=(full_sink,))
-  return graph_rewrite(sink, pm_to_program, ctx=(ren, compiler), name="linearize/render/compile")
-
-def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
-  """
-  Function to transform the Kernel UOp graph into a linearized program.
-
-  Args:
-    sink: The Ops.SINK rooting the Kernel graph.
-    ren: The Renderer (can change how things are processed, fix this).
-
-  Returns:
-    Linear program in UOps.
-  """
-
-  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  assert len(full_sink.ranges) == 0, f"all ranges must end by the sink, {full_sink.ranges}"
-  lst = line_rewrite(linearize(full_sink), pm_linearize_cleanups)
-  if SPEC: type_verify(lst, program_spec)
-  return lst
+  sink = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=ren.device)))
+  return graph_rewrite(sink, pm_to_program, ctx=ren, name="linearize/render")

tinygrad/codegen/opt/search.py

@@ -1,13 +1,11 @@
 import functools, math, time, multiprocessing, traceback, signal, atexit
-from dataclasses import replace
-from tinygrad.uop.ops import sym_infer, AxisType, pyrender
+from tinygrad.uop.ops import sym_infer, AxisType, pyrender, UOp
 from tinygrad.device import Device, Buffer, Compiler
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str, unwrap
 from tinygrad.helpers import IGNORE_BEAM_CACHE
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.tensor import Tensor
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.renderer import ProgramSpec
+from tinygrad.engine.realize import CompiledRunner, get_program, Estimates
 from tinygrad.codegen.opt.postrange import Scheduler
 
 actions = [Opt(op=OptOps.UPCAST, axis=axis, arg=amt) for amt in [0,2,3,4,5,7] for axis in range(8)]
@@ -34,19 +32,22 @@ def get_test_global_size(global_size, max_global_size, var_vals):
         break
   return test_global_size, input_size / prod(test_global_size)
 
-def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
+def _time_program(p:UOp, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
                   allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test") -> list[float]:
   factor = 1
-  if allow_test_size and p.global_size is not None and max_global_size is not None:
-    global_size, factor = get_test_global_size(p.global_size, max_global_size, var_vals)
-    p = replace(p, global_size=global_size)
-  try: car = CompiledRunner(replace(p, lib=lib))
+  global_size, local_size = p.sizes
+  if allow_test_size and global_size is not None and max_global_size is not None:
+    test_global_size, factor = get_test_global_size(global_size, max_global_size, var_vals)
+    # NOTE: we can't modify p.sizes directly, but optimize_local_size doesn't run in BEAM so this is ok
+  try: car = CompiledRunner(p, precompiled=lib)
   except AssertionError: return [math.inf] * cnt
   tms = []
   input_bufs = [rawbufs[i] for i in car.p.globals]
   for _ in range(cnt):
     if clear_l2:
-      if hasattr(dev:=Device[p.device], 'invalidate_caches'): dev.invalidate_caches()
+      p_device = p.device
+      assert isinstance(p_device, str), f"PROGRAM device must be a string, not {type(p_device)}"
+      if hasattr(dev:=Device[p_device], 'invalidate_caches'): dev.invalidate_caches()
       else:
         with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False)
     tms.append(unwrap(car(input_bufs, var_vals, wait=True))*factor)
@@ -58,7 +59,7 @@ def timeout_handler(signum, frame):
   if DEBUG >= 2: print("*** BEAM COMPILE TIMEOUT")
   raise TimeoutException()
 
-def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[ProgramSpec, bytes, float]|None]:
+def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[UOp, bytes, float]|None]:
   if hasattr(signal, "alarm"):
     signal.signal(getattr(signal, 'SIGALRM'), timeout_handler)
     # set timeout
@@ -66,12 +67,12 @@ def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[
   ret = None
   try:
     p = get_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].ren)
-    assert p.uops is not None, "uop list wasn't generated?"
-    if len(p.uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
-      if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(p.uops)=}, {uops_max=}")
+    uops = list(p.src[2].src)  # LINEAR is src[2]
+    if len(uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
+      if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(uops)=}, {uops_max=}")
       raise RuntimeError("too many uops")
     st = time.perf_counter()
-    prog = compiler.compile(p.src)
+    prog = compiler.compile(p.src[3].arg)  # SOURCE is src[3]
     et = time.perf_counter() - st
     ret = (p, prog, et)
   except RuntimeError:
@@ -154,7 +155,8 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True
         p, lib, compile_et = proc
         if lib in seen_libs: continue
         # filter out kernels that use 1000x more compute than the smallest
-        least_compute_ops = min(this_compute_ops:=sym_infer(p.estimates.ops, var_vals), least_compute_ops)
+        estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+        least_compute_ops = min(this_compute_ops:=sym_infer(estimates.ops, var_vals), least_compute_ops)
         if least_compute_ops*1000 < this_compute_ops:
           if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too much compute. {this_compute_ops} when least is {least_compute_ops}")
           continue
@@ -167,7 +169,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True
           raise
         timed.append((candidates[i], min(tms)))
         if BEAM_DEBUG > 1:
-          print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(unwrap(p.uops)):5d} uops",
+          print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(p.src[2].src):5d} uops",
                 f"{time_to_str(compile_et, w=12)} compile/{time_to_str(timed[-1][1], w=12)} run",
                 f"      {len(timed):4d}/{len(candidates):4d}         {timed[-1][0].colored_shape()}")
         elif DEBUG >= 2:

tinygrad/engine/jit.py

@@ -4,7 +4,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, colored, JIT, JIT_BATCH_SIZE, dedup, partition, unwrap
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
 from tinygrad.dtype import DType
-from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops
+from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops, sint
 from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, EncDec, CompiledRunner, Runner, Estimates
 from tinygrad.engine.memory import _internal_memory_planner
 from tinygrad.nn.state import get_parameters
@@ -78,13 +78,18 @@ class GraphRunner(Runner):
     self.input_replace:dict[tuple[int, int], int] = get_input_replace(jit_cache, input_rawbuffers)
     self.var_vals_replace:dict[int, list[tuple[int, int]]] = {}
     self.launch_dims_replace:dict[int, tuple[int|None, int|None]] = {}
-    self.launch_dims_base:dict[int, tuple[tuple[int, ...], tuple[int, ...]]] = {}
+    self.launch_dims_base:dict[int, tuple[tuple[sint, ...], tuple[sint, ...]]] = {}
 
-    def is_sym_dim(dim) -> bool: return not all(isinstance(d, (int, float)) for d in dim)
+    def is_sym_dim(dim) -> bool: return dim is not None and not all(isinstance(d, (int, float)) for d in dim)
 
     self.vars = sorted(var_vals.keys())
-    self.symbolic_dims = dedup([tuple(d) for ji in jit_cache if isinstance(ji.prg, CompiledRunner) and (d:=ji.prg.p.local_size) and is_sym_dim(d)] +
-                               [tuple(d) for ji in jit_cache if isinstance(ji.prg, CompiledRunner) and (d:=ji.prg.p.global_size) and is_sym_dim(d)])
+    sym_dims: list[tuple[sint, ...]] = []
+    for ji in jit_cache:
+      if isinstance(ji.prg, CompiledRunner):
+        global_size, local_size = ji.prg.p.sizes
+        if local_size is not None and is_sym_dim(local_size): sym_dims.append(tuple(local_size))
+        if global_size is not None and is_sym_dim(global_size): sym_dims.append(tuple(global_size))
+    self.symbolic_dims = dedup(sym_dims)
     def find_symbolic_dim(dim): return self.symbolic_dims.index(tuple(dim)) if dim is not None and tuple(dim) in self.symbolic_dims else None
 
     estimates = Estimates()
@@ -92,13 +97,15 @@ class GraphRunner(Runner):
       assert ji.prg is not None
       estimates += ji.prg.estimates
       if isinstance(ji.prg, CompiledRunner):
-        if ji.prg.p.vars: self.var_vals_replace[j] = [(i, self.vars.index(v.expr)) for i, v in enumerate(ji.prg.p.vars) if v.expr not in ji.fixedvars]
+        prg_vars = ji.prg.p.variables()
+        if prg_vars: self.var_vals_replace[j] = [(i, self.vars.index(v.expr)) for i, v in enumerate(prg_vars) if v.expr not in ji.fixedvars]
 
-        global_dim_idx, local_dim_idx = find_symbolic_dim(ji.prg.p.global_size), find_symbolic_dim(ji.prg.p.local_size)
+        global_size, local_size = ji.prg.p.sizes
+        global_dim_idx, local_dim_idx = find_symbolic_dim(global_size), find_symbolic_dim(local_size)
         if global_dim_idx is not None or local_dim_idx is not None:
           self.launch_dims_replace[j] = (global_dim_idx, local_dim_idx)
-          assert ji.prg.p.global_size is not None and ji.prg.p.local_size is not None
-          self.launch_dims_base[j] = (tuple(ji.prg.p.global_size), tuple(ji.prg.p.local_size))
+          assert global_size is not None and local_size is not None
+          self.launch_dims_base[j] = (tuple(global_size), tuple(local_size))
 
     # used in MultiGraphRunner. the ints are id() of _bufs
     self.w_dependency_map: dict[int, Any] = {}

tinygrad/engine/realize.py

@@ -1,31 +1,30 @@
 from typing import cast, Callable
 import time, pprint, random, itertools, math
-from dataclasses import dataclass, replace, field
+from dataclasses import dataclass, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
-from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, PROFILE, ProfilePointEvent, cpu_events, prod, Context
-from tinygrad.helpers import unwrap
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, track_rewrites, KernelInfo, pyrender
+from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
+from tinygrad.helpers import unwrap, to_function_name
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo, pyrender
 from tinygrad.device import Device, Buffer
-from tinygrad.renderer import Renderer, ProgramSpec, Estimates
+from tinygrad.renderer import Renderer, Estimates
 from tinygrad.codegen import full_rewrite_to_program
 from tinygrad.codegen.opt import Opt
 
 # **************** Program Creation ****************
 
-@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret), replay=True)
-def get_program(ast:UOp, renderer:Renderer, device:str|None=None, opts:list[Opt]|None=None) -> ProgramSpec:
+@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.src[0].arg.name, (to_function_name(ret.src[0].arg.name), ret.src[0]), ret=ret),
+                replay=True)
+def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> UOp:
   """
-  Transform an AST into a ProgramSpec. May trigger BEAM search.
+  Transform an AST into a PROGRAM UOp. May trigger BEAM search.
 
   Args:
     ast: The Ops.SINK rooted AST
     renderer: The renderer used to generate the code
-    device: The device to compile for (defaults to renderer.device)
 
   Returns:
-    The ProgramSpec of the program.
+    The PROGRAM UOp with structure (SINK, DEVICE, LINEAR, SOURCE).
   """
-  device = device or renderer.device
 
   if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(ast))
@@ -36,14 +35,12 @@ def get_program(ast:UOp, renderer:Renderer, device:str|None=None, opts:list[Opt]
     ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
   if ast.arg is None: ast = ast.replace(arg=KernelInfo())
 
-  prg = full_rewrite_to_program(ast, renderer, Device[device].compiler)
-  # SINK/LINEAR/SOURCE/BINARY
-  sink, linear, source, binary = prg.src
+  prg = full_rewrite_to_program(ast, renderer)
 
-  # legacy
-  return ProgramSpec(sink.arg.name, source.arg, device, sink, list(linear.src), binary.arg,
-                     global_size=[1,1,1] if renderer.has_local or renderer.has_threads else None,
-                     local_size=[1,1,1] if renderer.has_local else None)
+  # print
+  if DEBUG >= 6: print_uops(list(prg.src[2].src))  # LINEAR is src[2]
+
+  return prg
 
 # **************** Runners ****************
 
@@ -72,25 +69,36 @@ def optimize_local_size(_prg:Callable, global_size:list[int], rawbufs:list[Buffe
   return ret[1]
 
 class CompiledRunner(Runner):
-  def __init__(self, p:ProgramSpec, prg=None):
-    if DEBUG >= 3: print(p.applied_opts)
-    if DEBUG >= 4: print(p.src)
-    assert p.lib is not None, "lib must be provided"
-    self.p:ProgramSpec = p
-    if DEBUG >= 7: Device[p.device].compiler.disassemble(p.lib)
-    self._prg = Device[p.device].runtime(p.function_name, p.lib) if prg is None else prg
-    super().__init__(p.name, p.device, p.estimates)
+  def __init__(self, p:UOp, precompiled:bytes|None=None, prg=None):
+    assert p.op is Ops.PROGRAM, f"CompiledRunner requires PROGRAM UOp, not {p.op}"
+    self.p:UOp = p
+    dev = p.device
+    assert isinstance(dev, str), f"PROGRAM device must be a string, not {type(dev)}"
+    name = p.src[0].arg.name
+    src = p.src[3].arg
+    function_name = to_function_name(name)
+    if DEBUG >= 3: print(p.src[0].arg.applied_opts)
+    if DEBUG >= 4: print(src)
+    if precompiled is not None: self.lib = precompiled
+    else:
+      with cpu_profile(TracingKey(f"compile {name}", (function_name,)), "TINY"):
+        self.lib = Device[dev].compiler.compile_cached(src)
+    if DEBUG >= 7: Device[dev].compiler.disassemble(self.lib)
+    self._prg = Device[dev].runtime(function_name, self.lib) if prg is None else prg
+    super().__init__(name, dev, Estimates.from_uops(list(p.src[2].src), ignore_indexing=True))
 
-  def __reduce__(self): return self.__class__, (self.p,)
+  def __reduce__(self): return self.__class__, (self.p, self.lib)
 
   def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False) -> float|None:
     if var_vals is None: var_vals = {}
-    has_local = Device[self.p.device].renderer.has_local
+    dev = self.p.device
+    assert isinstance(dev, str), f"PROGRAM device must be a string, not {type(dev)}"
+    has_local = Device[dev].renderer.has_local
     global_size, local_size = self.p.launch_dims(var_vals)
-    if has_local and global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type]
+    sym_global_size, sym_local_size = self.p.sizes
+    if has_local and global_size is not None and local_size is None and sym_global_size is not None and all_int(sym_global_size):
       local_size = optimize_local_size(self._prg, global_size, rawbufs)
-      global_size = [g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)]
-      self.p = replace(self.p, global_size=global_size, local_size=local_size)
+      global_size = [g//l if g%l == 0 else int(g/l) for g,l in zip(global_size, local_size)]
     lra = {}
     if global_size:
       lra['global_size'] = tuple(global_size)
@@ -98,7 +106,7 @@ class CompiledRunner(Runner):
     if local_size:
       lra['local_size'] = tuple(local_size)
       assert len(local_size) == 3, "local size must have len 3"
-    return self._prg(*[x._buf for x in rawbufs], **lra, vals=tuple(var_vals[k.expr] for k in self.p.vars), wait=wait)
+    return self._prg(*[x._buf for x in rawbufs], **lra, vals=tuple(var_vals[k.expr] for k in self.p.variables()), wait=wait)
 
 class ViewOp(Runner):
   def __init__(self, buf:Buffer): super().__init__(colored(f"view {buf.nbytes:8d} @ {buf.offset:<10d}", "yellow"), buf.device)
@@ -155,10 +163,14 @@ def get_runner(device:str, ast:UOp) -> CompiledRunner:
   if cret:=method_cache.get(ckey): return cret
   bkey = (device.split(":")[0], type(Device[device].compiler), ast.key, context, True)
   if bret:=method_cache.get(bkey):
-    method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device))
+    # update device in PROGRAM UOp
+    new_p = bret.p.replace(src=(bret.p.src[0], UOp(Ops.DEVICE, arg=device), *bret.p.src[2:]))
+    method_cache[ckey] = ret = CompiledRunner(new_p, bret.lib)
   else:
-    prg: ProgramSpec = get_program(ast, Device[device].renderer, device)
-    method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, device=device))
+    prg = get_program(ast, Device[device].renderer)
+    # update device in PROGRAM UOp to match the actual device
+    prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=device), *prg.src[2:]))
+    method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(prg)
   return ret
 
 # **************** lowering functions ****************

tinygrad/engine/schedule.py

@@ -90,23 +90,29 @@ from tinygrad.engine.memory import memory_planner
 from tinygrad.schedule.rangeify import get_rangeify_map
 from tinygrad.schedule.multi import get_multi_map
 
-def replace_input_buffer(ctx:dict[UOp, UOp], b:UOp):
-  if (ret:=ctx.get(b, None)) is None:
+def replace_input_buffer(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
+  if (ret:=ctx[0].get(b, None)) is None:
     if b.op is Ops.BUFFER:
-      ctx[b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx)), b.src[1]))
+      ctx[0][b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx[0])), b.src[1]))
     else:
       # TODO: flip args in CONST
       assert b.op is Ops.CONST
-      ctx[b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx))))
+      ctx[0][b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx[0]))))
   return ret
 
+def strip_bind(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
+  var, val = b.src[0], b.src[1].arg
+  assert var.expr not in ctx[1] or ctx[1][var.expr] == val, f"bind mismatch on {var}, {ctx[1][var.expr]} != {val}"
+  ctx[1][var.expr] = val
+  return ctx[0].setdefault(b, b.replace(src=(b.src[0],)))
+
 pm_pre_sched_cache = PatternMatcher([
   # replace input buffers
   (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
   # remove unique consts
   (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.UNIQUE)), name="b"), replace_input_buffer),
   # strip value from BIND for cache key normalization, so different values hit same cache
-  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), lambda ctx,b: ctx.setdefault(b, b.replace(src=(b.src[0],)))),
+  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), strip_bind),
 ])
 
 def replace_input_buffer_back(ctx:dict[UOp, UOp], b:UOp):
@@ -129,9 +135,10 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
   # big_sink srcs are all the Tensors
   st = time.perf_counter()
 
-  # replace all UNIQUE buffers with LUNIQUE, strip BIND values for cache key
+  # replace all UNIQUE buffers with LUNIQUE, strip BIND values for cache key, extract var_vals
   input_buffers: dict[UOp, UOp] = {}
-  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=input_buffers, name="rewrite for sched cache")
+  var_vals: dict[str, int] = {}
+  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=(input_buffers, var_vals), name="rewrite for sched cache")
   sched_cache_key = big_sink_cache.key
 
   if (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None:
@@ -139,7 +146,7 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
     if SPEC: type_verify(big_sink, tensor_spec)
 
     # hack to preserve metadata
-    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx={}, name="preserve metadata")
+    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx=({}, {}), name="preserve metadata")
 
     # tensor map is what we return
     tensor_map: dict[UOp, UOp] = {}
@@ -191,17 +198,8 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
       schedule.append(ExecItem(si.ast, list(ubufs), si.metadata, si.fixedvars))
   with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
 
-  # extract var_vals from BINDs that were stripped (only if there are kernels)
-  var_vals: dict[str, int] = {}
-  if schedule:
-    for u in input_buffers:
-      if u.op is Ops.BIND:
-        var, val = u.unbind()
-        assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
-        var_vals[var.expr] = val
-
   if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
     print(f"scheduled {len(schedule):4d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
           f" | {' cache hit' if sc_ret is not None else 'CACHE MISS'} {sched_cache_key.hex()[:8]}"+\
           f" | {len(UOpMetaClass.ucache)} uops in cache")
-  return tensor_map, schedule, var_vals
+  return tensor_map, schedule, var_vals if schedule else {}

tinygrad/renderer/__init__.py

@@ -1,12 +1,10 @@
 from __future__ import annotations
 from typing import Callable, cast
-import functools
-from dataclasses import dataclass, field
-from tinygrad.helpers import to_function_name, dedup, prod
-from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher
+from dataclasses import dataclass
+from tinygrad.helpers import prod
+from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, GroupOp, PatternMatcher
 from tinygrad.dtype import AddrSpace, PtrDType
 from tinygrad.codegen.opt.tc import TensorCore
-from tinygrad.codegen.opt import Opt
 
 @dataclass(frozen=True)
 class Estimates:
@@ -57,63 +55,6 @@ class Estimates:
       elif u.op is Ops.WMMA and u not in dont_count: flops += 2 * prod(u.arg[1]) // u.arg[5] * mults
     return Estimates(flops, lds, sum(mem.values()))
 
-@dataclass
-class ProgramSpec:
-  name:str
-  src:str
-  device:str
-  ast:UOp  # save the base ast (this is method cache key)
-  uops:list[UOp]|None=None
-  lib:bytes|None=None  # compiled binary
-
-  # filled in from uops (if we have uops)
-  global_size:list[int]|None=None
-  local_size:list[int]|None=None
-  vars:list[Variable]=field(default_factory=list)
-  globals:list[int]=field(default_factory=list)
-  outs:list[int]=field(default_factory=list)
-  ins:list[int]=field(default_factory=list)
-  _ran_post_init:bool=False  # NOTE: this is needed if you call replace on the Program
-
-  def __post_init__(self):
-    if not self._ran_post_init and self.uops is not None:
-      # single pass through the uops
-      for u in self.uops:
-        if u.op is Ops.DEFINE_VAR: self.vars.append(u)
-        if u.op is Ops.DEFINE_GLOBAL: self.globals.append(u.arg)
-        if u.op in (Ops.STORE, Ops.LOAD):
-          if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX):
-            if (buf:=idx.src[0]).op is Ops.DEFINE_GLOBAL: (self.outs if u.op is Ops.STORE else self.ins).append(buf.arg)
-          # TODO: can else happen?
-        if u.op is Ops.SPECIAL:
-          # NOTE: you have to set local_size and global_size to the base [1,1,1] outside this
-          if u.arg[0] == 'i': self.local_size = None
-          special_size = self.local_size if u.arg[0] == 'l' else self.global_size
-          # TODO: this cast is wrong, u.src[0].ssimplify() can be sint
-          if special_size is not None: special_size[int(u.arg[-1])] = cast(int, u.src[0].ssimplify())
-      self.vars = sorted(self.vars, key=lambda v: v.arg)
-      self.outs = sorted(dedup(self.outs))
-      self.ins = sorted(dedup(self.ins))
-      self._ran_post_init = True
-
-  @functools.cached_property
-  def estimates(self) -> Estimates:
-    return Estimates() if self.uops is None else Estimates.from_uops(self.uops, ignore_indexing=True)
-
-  @functools.cached_property
-  def function_name(self) -> str: return to_function_name(self.name)
-
-  @property
-  def applied_opts(self) -> tuple[Opt, ...]|None:
-    if self.uops is None: return None
-    assert self.uops[-1].op is Ops.SINK, self.uops[-1].op
-    return self.uops[-1].arg.applied_opts
-
-  def launch_dims(self, var_vals:dict[str, int]):
-    global_size = [sym_infer(sz, var_vals) for sz in self.global_size] if self.global_size is not None else None
-    local_size = [sym_infer(sz, var_vals) for sz in self.local_size] if self.local_size is not None else None
-    return global_size, local_size
-
 class Renderer:
   device: str = ""
   suffix: str = ""

tinygrad/runtime/graph/cuda.py

@@ -22,12 +22,14 @@ class CUDAGraph(MultiGraphRunner):
     for j,ji in enumerate(jit_cache):
       if isinstance(ji.prg, CompiledRunner):
         global_size, local_size = ji.prg.p.launch_dims(var_vals)
+        assert global_size is not None and local_size is not None
 
         new_node = cuda.CUgraphNode()
         deps = self._access_resources([x.base for x in ji.bufs if x is not None], ji.prg.p.outs, new_dependency=new_node)
         c_deps = (cuda.CUgraphNode*len(deps))(*deps) if deps else None
 
-        c_args, vargs = encode_args([cast(Buffer, x)._buf for x in ji.bufs], [var_vals.get(x.expr, ji.fixedvars.get(x.expr)) for x in ji.prg.p.vars])
+        prg_vars = ji.prg.p.variables()
+        c_args, vargs = encode_args([cast(Buffer, x)._buf for x in ji.bufs], [var_vals.get(x.expr, ji.fixedvars.get(x.expr)) for x in prg_vars])
         kern_params = cuda.CUDA_KERNEL_NODE_PARAMS_v1(ji.prg._prg.prg, *global_size, *local_size, 0, None, vargs)
         check(cuda.cuGraphAddKernelNode(ctypes.byref(new_node), self.graph, c_deps, len(deps), ctypes.byref(kern_params)))
 

tinygrad/runtime/graph/hcq.py

@@ -39,7 +39,7 @@ class HCQGraph(MultiGraphRunner):
       if not isinstance(ji.prg, CompiledRunner): continue
 
       argsbuf = self.kernargs_bufs[ji.prg.dev].offset(kargs_alloc[ji.prg.dev].alloc(ji.prg._prg.kernargs_alloc_size, 16))
-      self.ji_args[j] = ji.prg._prg.fill_kernargs(self.hcq_bufs[j], ji.prg.p.vars, argsbuf)
+      self.ji_args[j] = ji.prg._prg.fill_kernargs(self.hcq_bufs[j], ji.prg.p.prog_vars(), argsbuf)
 
     # Schedule Dependencies.
     # There are two types of queues on each device: copy and compute. Both must synchronize with all external operations before launching any
@@ -159,7 +159,8 @@ class HCQGraph(MultiGraphRunner):
 
       # Encode main commands based on ji type.
       if isinstance(ji.prg, CompiledRunner):
-        enqueue_queue.exec(ji.prg._prg, self.ji_args[j], tuple(ji.prg.p.global_size or (1,1,1)), tuple(ji.prg.p.local_size or (1,1,1)))
+        global_size, local_size = ji.prg.p.sizes
+        enqueue_queue.exec(ji.prg._prg, self.ji_args[j], tuple(global_size or (1,1,1)), tuple(local_size or (1,1,1)))
       elif isinstance(ji.prg, (BufferXfer, BufferCopy)):
         dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
         for bufid, src in enumerate(cast(list[Buffer], ji.bufs)):

tinygrad/runtime/graph/metal.py

@@ -42,9 +42,11 @@ class MetalGraph(GraphRunner):
         if b is not None and b not in input_rawbuffers:
           icb_command.setKernelBuffer_offset_atIndex(b._buf.buf, b._buf.offset, i)
           all_resources.append(b._buf.buf)
-      for i,v in enumerate(prg.p.vars): icb_command.setKernelBuffer_offset_atIndex(self.int_buf.buf, self.varlist.index(v.expr)*4, len(ji.bufs)+i)
+      for i,v in enumerate(prg.p.variables()):
+        icb_command.setKernelBuffer_offset_atIndex(self.int_buf.buf, self.varlist.index(v.expr)*4, len(ji.bufs)+i)
 
       global_size, local_size = prg.p.launch_dims(var_vals)
+      assert global_size is not None and local_size is not None
       icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size))
       icb_command.setBarrier()
 

tinygrad/runtime/ops_amd.py

@@ -18,6 +18,7 @@ from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets, import_pmc
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, PCIDevice, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
+from tinygrad.runtime.support.memory import AddrSpace
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
 SQTT = ContextVar("SQTT", abs(VIZ.value)>=2)
@@ -859,7 +860,7 @@ class USBIface(PCIIface):
     self.sys_buf, self.sys_next_off = self._dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0x800
 
   def _dma_region(self, ctrl_addr, sys_addr, size):
-    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], system=True, uncached=True)
+    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
     return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
 
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:

tinygrad/runtime/support/am/amdev.py

@@ -4,7 +4,7 @@ from tinygrad.helpers import mv_address, getenv, DEBUG, fetch
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.amd import AMDReg, import_module, import_asic_regs
-from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager
+from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
 from tinygrad.runtime.support.system import PCIDevice, PCIDevImplBase
 from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
 
@@ -120,10 +120,11 @@ class AMFirmware:
 class AMPageTableEntry:
   def __init__(self, adev, paddr, lv): self.adev, self.paddr, self.lv, self.entries = adev, paddr, lv, adev.vram.view(paddr, 0x1000, fmt='Q')
 
-  def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, system=False, snooped=False, frag=0, valid=True):
-    if not system: paddr = self.adev.paddr2xgmi(paddr)
+  def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, aspace=AddrSpace.PHYS, snooped=False, frag=0, valid=True):
+    is_sys = aspace is AddrSpace.SYS
+    if aspace is AddrSpace.PHYS: paddr = self.adev.paddr2xgmi(paddr)
     assert paddr & self.adev.gmc.address_space_mask == paddr, f"Invalid physical address {paddr:#x}"
-    self.entries[entry_id] = self.adev.gmc.get_pte_flags(self.lv, table, frag, uncached, system, snooped, valid) | (paddr & 0x0000FFFFFFFFF000)
+    self.entries[entry_id] = self.adev.gmc.get_pte_flags(self.lv, table, frag, uncached, is_sys, snooped, valid) | (paddr & 0x0000FFFFFFFFF000)
 
   def entry(self, entry_id:int) -> int: return self.entries[entry_id]
   def valid(self, entry_id:int) -> bool: return (self.entries[entry_id] & am.AMDGPU_PTE_VALID) != 0

tinygrad/runtime/support/am/ip.py

@@ -3,6 +3,7 @@ from typing import Literal
 from tinygrad.helpers import to_mv, data64, lo32, hi32, DEBUG, wait_cond, pad_bytes
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.amd import import_soc
+from tinygrad.runtime.support.memory import AddrSpace
 
 class AM_IP:
   def __init__(self, adev): self.adev = adev
@@ -468,7 +469,7 @@ class AM_PSP(AM_IP):
     msg1_region = next((reg for reg in self.adev.dma_regions or [] if reg[1].nbytes >= (512 << 10)), None)
     if msg1_region is not None:
       self.msg1_addr, self.msg1_view = self.adev.mm.alloc_vaddr(size=msg1_region[1].nbytes, align=am.PSP_1_MEG), msg1_region[1]
-      self.adev.mm.map_range(self.msg1_addr, msg1_region[1].nbytes, [(msg1_region[0], msg1_region[1].nbytes)], system=True, uncached=True, boot=True)
+      self.adev.mm.map_range(self.msg1_addr, msg1_region[1].nbytes, [(msg1_region[0],msg1_region[1].nbytes)], AddrSpace.SYS, uncached=True, boot=True)
     else:
       self.msg1_paddr = self.adev.mm.palloc(am.PSP_1_MEG, align=am.PSP_1_MEG, zero=False, boot=True)
       self.msg1_addr, self.msg1_view = self.adev.paddr2mc(self.msg1_paddr), self.adev.vram.view(self.msg1_paddr, am.PSP_1_MEG, 'B')

tinygrad/runtime/support/memory.py

@@ -1,4 +1,4 @@
-import collections, functools, dataclasses
+import collections, functools, dataclasses, enum
 from typing import Any, ClassVar
 from tinygrad.helpers import round_up, getenv
 
@@ -107,8 +107,10 @@ class TLSFAllocator:
 
 # Memory Managment
 
+class AddrSpace(enum.Enum): PHYS = enum.auto(); SYS = enum.auto(); PEER = enum.auto() # noqa: E702
+
 @dataclasses.dataclass(frozen=True)
-class VirtMapping: va_addr:int; size:int; paddrs:list[tuple[int, int]]; uncached:bool=False; system:bool=False; snooped:bool=False # noqa: E702
+class VirtMapping: va_addr:int; size:int; paddrs:list[tuple[int, int]]; aspace:AddrSpace; uncached:bool=False; snooped:bool=False # noqa: E702
 
 class PageTableTraverseContext:
   def __init__(self, dev, pt, vaddr, create_pts=False, free_pts=False, boot=False):
@@ -190,7 +192,7 @@ class MemoryManager:
     ctx = PageTableTraverseContext(self.dev, self.root_page_table, vaddr, create_pts=True)
     for _ in ctx.next(size, paddr=0): return [pt for pt, _, _ in ctx.pt_stack]
 
-  def map_range(self, vaddr:int, size:int, paddrs:list[tuple[int, int]], uncached=False, system=False, snooped=False, boot=False) -> VirtMapping:
+  def map_range(self, vaddr:int, size:int, paddrs:list[tuple[int, int]], aspace:AddrSpace, uncached=False, snooped=False, boot=False) -> VirtMapping:
     if getenv("MM_DEBUG", 0): print(f"mm {self.dev.devfmt}: mapping {vaddr=:#x} ({size=:#x})")
 
     assert size == sum(p[1] for p in paddrs), f"Size mismatch {size=} {sum(p[1] for p in paddrs)=}"
@@ -200,11 +202,11 @@ class MemoryManager:
       for off, pt, pte_idx, pte_cnt, pte_covers in ctx.next(psize, paddr=paddr):
         for pte_off in range(pte_cnt):
           assert not pt.valid(pte_idx + pte_off), f"PTE already mapped: {pt.entry(pte_idx + pte_off):#x}"
-          pt.set_entry(pte_idx + pte_off, paddr + off + pte_off * pte_covers, uncached=uncached, system=system, snooped=snooped,
+          pt.set_entry(pte_idx + pte_off, paddr + off + pte_off * pte_covers, uncached=uncached, aspace=aspace, snooped=snooped,
                        frag=self._frag_size(ctx.vaddr+off, pte_cnt * pte_covers), valid=True)
 
     self.on_range_mapped()
-    return VirtMapping(vaddr, size, paddrs, uncached=uncached, system=system, snooped=snooped)
+    return VirtMapping(vaddr, size, paddrs, aspace=aspace, uncached=uncached, snooped=snooped)
 
   def unmap_range(self, vaddr:int, size:int):
     if getenv("MM_DEBUG", 0): print(f"mm {self.dev.devfmt}: unmapping {vaddr=:#x} ({size=:#x})")
@@ -243,7 +245,7 @@ class MemoryManager:
           continue
         rem_size -= self.palloc_ranges[nxt_range][0]
 
-    return self.map_range(va, size, paddrs, uncached=uncached)
+    return self.map_range(va, size, paddrs, aspace=AddrSpace.PHYS, uncached=uncached)
 
   def vfree(self, vm:VirtMapping):
     assert self.va_allocator is not None, "must be set it"

tinygrad/runtime/support/nv/nvdev.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import ctypes, time, functools, re, gzip, struct
 from tinygrad.helpers import getenv, DEBUG, fetch, getbits
-from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager
+from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
 from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP
 from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
 
@@ -33,9 +33,9 @@ class NVPageTableEntry:
 
   def _is_dual_pde(self) -> bool: return self.lv == self.nvdev.mm.level_cnt - 2
 
-  def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, system=False, snooped=False, frag=0, valid=True):
+  def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, aspace=AddrSpace.PHYS, snooped=False, frag=0, valid=True):
     if not table:
-      x = self.nvdev.pte_t.encode(valid=valid, address_sys=paddr >> 12, aperture=2 if system else 0, kind=6,
+      x = self.nvdev.pte_t.encode(valid=valid, address_sys=paddr >> 12, aperture=2 if aspace is AddrSpace.SYS else 0, kind=6,
         **({'pcf': int(uncached)} if self.nvdev.mmu_ver == 3 else {'vol': uncached}))
     else:
       pde = self.nvdev.dual_pde_t if self._is_dual_pde() else self.nvdev.pde_t

tinygrad/runtime/support/system.py

@@ -3,7 +3,7 @@ from typing import cast, ClassVar
 from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv
 from tinygrad.runtime.autogen import libc, vfio, pci
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
-from tinygrad.runtime.support.memory import MemoryManager, VirtMapping
+from tinygrad.runtime.support.memory import MemoryManager, VirtMapping, AddrSpace
 from tinygrad.runtime.support.usb import ASM24Controller, USBMMIOInterface
 
 MAP_FIXED, MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0x10, 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000), 0x400
@@ -262,7 +262,7 @@ class LNXPCIIfaceBase:
     if should_use_sysmem:
       vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
       memview, paddrs = System.alloc_sysmem(size, vaddr=vaddr, contiguous=contiguous)
-      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], system=True, snooped=True, uncached=True)
+      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], aspace=AddrSpace.SYS, snooped=True, uncached=True)
       return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), view=memview, owner=self.dev)
 
     mapping = self.dev_impl.mm.valloc(size:=round_up(size, 0x1000), uncached=uncached, contiguous=cpu_access)
@@ -271,19 +271,23 @@ class LNXPCIIfaceBase:
 
   def free(self, b:HCQBuffer):
     for dev in b.mapped_devs[1:]: dev.iface.dev_impl.mm.unmap_range(b.va_addr, b.size)
-    if not b.meta.mapping.system: self.dev_impl.mm.vfree(b.meta.mapping)
+    if b.meta.mapping.aspace is AddrSpace.PHYS: self.dev_impl.mm.vfree(b.meta.mapping)
     if b.owner == self.dev and b.meta.has_cpu_mapping and not OSX: FileIOInterface.munmap(b.va_addr, b.size)
 
   def map(self, b:HCQBuffer):
     if b.owner is not None and b.owner._is_cpu():
       System.lock_memory(cast(int, b.va_addr), b.size)
-      paddrs, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, True
+      paddrs, aspace = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], AddrSpace.SYS
+      snooped, uncached = True, True
     elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, LNXPCIIfaceBase):
-      paddrs = [(paddr if b.meta.mapping.system else (paddr + ifa.p2p_base_addr), size) for paddr,size in b.meta.mapping.paddrs]
-      snooped, uncached = b.meta.mapping.snooped, b.meta.mapping.uncached
+      snooped, uncached = True, b.meta.mapping.uncached
+      if b.meta.mapping.aspace is AddrSpace.SYS: paddrs, aspace = b.meta.mapping.paddrs, AddrSpace.SYS
+      elif hasattr(ifa.dev_impl, 'paddr2xgmi') and ifa.dev_impl.gmc.xgmi_seg_sz > 0:
+        paddrs, aspace = [(ifa.dev_impl.paddr2xgmi(p), sz) for p, sz in b.meta.mapping.paddrs], AddrSpace.PEER
+      else: paddrs, aspace = [(p + ifa.p2p_base_addr, sz) for p, sz in b.meta.mapping.paddrs], AddrSpace.SYS
     else: raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
 
-    self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, system=True, snooped=snooped, uncached=uncached)
+    self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
 
 class APLPCIIfaceBase(LNXPCIIfaceBase):
   def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):

tinygrad/uop/__init__.py

@@ -27,9 +27,9 @@ class Ops(FastEnum):
   # uops that aren't rendered
   NOOP = auto(); REWRITE_ERROR = auto()
 
-  # renderer/compiler
-  # LINEAR is a list of UOps, SOURCE has a str arg that's human readable, BINARY has a bytes arg that's not
-  PROGRAM = auto(); LINEAR = auto(); SOURCE = auto(); BINARY = auto()
+  # renderer
+  # LINEAR is a list of UOps, SOURCE has a str arg that's human readable
+  PROGRAM = auto(); LINEAR = auto(); SOURCE = auto()
 
   # AFTER passes src[0] through and promises in the toposort that any consumers of the AFTER run after src[1:]
   # GROUP is a NOOP that just merges things together

tinygrad/uop/ops.py

@@ -219,7 +219,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       # late ops don't have shape
       case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
            Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL | \
-           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY:
+           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE:
         return None
 
       case Ops.INDEX:
@@ -609,11 +609,17 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   @staticmethod
   def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
     return UOp(Ops.BUFFER, dtype, (UOp.unique(num), UOp(Ops.DEVICE, arg=device)), size)
+  @staticmethod
+  def new_program(name:str, src:str, device:str, ast:UOp, uops:list[UOp]):
+    """Create a PROGRAM UOp from raw components."""
+    sink = ast.replace(arg=KernelInfo(name=name)) if ast.arg is None else ast.replace(arg=ast.arg.replace(name=name))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=tuple(uops)), UOp(Ops.SOURCE, arg=src)))
   @property
   def device(self) -> str|tuple[str, ...]: return unwrap(self._device)
   @recursive_property
   def _device(self) -> str|tuple[str, ...]|None:
     if self.op is Ops.DEVICE: return self.arg
+    if self.op is Ops.PROGRAM: return self.src[1].arg  # PROGRAM src[1] is DEVICE
     if self.op is Ops.BUFFERIZE: return self.arg.device
     if self.op is Ops.AFTER: return self.src[0]._device
     if self.op is Ops.MSELECT:
@@ -624,6 +630,104 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     for x in self.src:
       if x._device is not None: return x._device
     return None
+
+  # *** PROGRAM UOp properties ***
+
+  @property
+  def uops(self) -> list[UOp]:
+    """Linearized uops list. Only valid for PROGRAM."""
+    assert self.op is Ops.PROGRAM, f"uops only valid for PROGRAM, not {self.op}"
+    return list(self.src[2].src)
+
+  @property
+  def name(self) -> str:
+    """Kernel name. Only valid for PROGRAM."""
+    assert self.op is Ops.PROGRAM, f"name only valid for PROGRAM, not {self.op}"
+    return self.src[0].arg.name
+
+  @property
+  def applied_opts(self):
+    """Applied optimizations. Only valid for PROGRAM."""
+    assert self.op is Ops.PROGRAM, f"applied_opts only valid for PROGRAM, not {self.op}"
+    return self.src[0].arg.applied_opts
+
+  @functools.cached_property
+  def estimates(self):
+    """Estimates for this program. Only valid for PROGRAM."""
+    assert self.op is Ops.PROGRAM, f"estimates only valid for PROGRAM, not {self.op}"
+    from tinygrad.renderer import Estimates
+    return Estimates.from_uops(self.uops, ignore_indexing=True)
+
+  @property
+  def globals(self) -> list[int]:
+    """DEFINE_GLOBAL arg indices from linearized uops. Only valid for PROGRAM."""
+    assert self.op is Ops.PROGRAM, f"globals only valid for PROGRAM, not {self.op}"
+    return [u.arg for u in self.src[2].src if u.op is Ops.DEFINE_GLOBAL]
+
+  @property
+  def outs(self) -> list[int]:
+    """Buffer indices written to (STORE). Only valid for PROGRAM."""
+    assert self.op is Ops.PROGRAM, f"outs only valid for PROGRAM, not {self.op}"
+    ret = []
+    for u in self.src[2].src:
+      if u.op is Ops.STORE:
+        idx = u.src[0]
+        if idx.op is Ops.CAST: idx = idx.src[0]
+        if idx.op is Ops.INDEX and idx.src[0].op is Ops.DEFINE_GLOBAL: ret.append(idx.src[0].arg)
+    return sorted(set(ret))
+
+  @property
+  def ins(self) -> list[int]:
+    """Buffer indices read from (LOAD). Only valid for PROGRAM."""
+    assert self.op is Ops.PROGRAM, f"ins only valid for PROGRAM, not {self.op}"
+    ret = []
+    for u in self.src[2].src:
+      if u.op is Ops.LOAD:
+        idx = u.src[0]
+        if idx.op is Ops.CAST: idx = idx.src[0]
+        if idx.op is Ops.INDEX and idx.src[0].op is Ops.DEFINE_GLOBAL: ret.append(idx.src[0].arg)
+    return sorted(set(ret))
+
+  @functools.cached_property
+  def sizes(self) -> tuple[list[sint]|None, list[sint]|None]:
+    """Get (global_size, local_size) which may contain symbolic values. Only valid for PROGRAM."""
+    assert self.op is Ops.PROGRAM, f"sizes only valid for PROGRAM, not {self.op}"
+    from tinygrad.device import Device
+    dev = self.device
+    assert isinstance(dev, str), f"PROGRAM device must be a string, not {type(dev)}"
+    ren = Device[dev].renderer
+    global_size:list[sint]|None = [1,1,1] if ren.has_local or ren.has_threads else None
+    local_size:list[sint]|None = [1,1,1] if ren.has_local else None
+    for u in self.src[2].src:
+      if u.op is Ops.SPECIAL:
+        if u.arg[0] == 'i': local_size = None
+        special_size = local_size if u.arg[0] == 'l' else global_size
+        if special_size is not None: special_size[int(u.arg[-1])] = cast(sint, u.src[0].ssimplify())
+    return global_size, local_size
+
+  def launch_dims(self, var_vals:dict[str, int]) -> tuple[list[int]|None, list[int]|None]:
+    """Resolve global/local sizes to concrete ints. Only valid for PROGRAM."""
+    global_size, local_size = self.sizes
+    global_ret = [sym_infer(sz, var_vals) for sz in global_size] if global_size is not None else None
+    local_ret = [sym_infer(sz, var_vals) for sz in local_size] if local_size is not None else None
+    return global_ret, local_ret
+
+  @property
+  def global_size(self) -> list[sint]|None:
+    """Global size (may be symbolic). Only valid for PROGRAM."""
+    return self.sizes[0]
+
+  @property
+  def local_size(self) -> list[sint]|None:
+    """Local size (may be symbolic). Only valid for PROGRAM."""
+    return self.sizes[1]
+
+  def prog_vars(self) -> list:
+    """Variables list for this program. Only valid for PROGRAM."""
+    assert self.op is Ops.PROGRAM, f"prog_vars only valid for PROGRAM, not {self.op}"
+    # Get variables from the linearized uops
+    linear_uops = self.src[2]
+    return linear_uops.variables()
   @property
   def buf_uop(self) -> UOp:
     if self.op is Ops.BUFFER: return self

tinygrad/uop/spec.py

@@ -249,15 +249,13 @@ full_spec = PatternMatcher([
   # in progress MSTACK may lose device
   (UPat((Ops.MSELECT, Ops.MSTACK), name="x"), lambda x: True),
 
-  # codegen: PROGRAM with progressive sources through the pipeline
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK),)), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.LINEAR))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.LINEAR), UPat(Ops.SOURCE), UPat(Ops.BINARY))), lambda: True),
-  # codegen: standalone LINEAR/SOURCE/BINARY
+  # codegen: PROGRAM with progressive sources through the pipeline (SINK, DEVICE, LINEAR, SOURCE)
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE))), lambda: True),
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR))), lambda: True),
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
+  # codegen: standalone LINEAR/SOURCE
   (UPat(Ops.LINEAR, dtypes.void), lambda: True),
   (UPat(Ops.SOURCE, dtypes.void, src=()), lambda: True),
-  (UPat(Ops.BINARY, dtypes.void, src=()), lambda: True),
 
   # temp VECTORIZE/INDEX during rewrite have the wrong dtype
   (UPat(Ops.VECTORIZE), lambda: True),

tinygrad/viz/serve.py

@@ -10,7 +10,7 @@ from tinygrad.helpers import printable, TCPServerWithReuse, HTTPRequestHandler
 from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupOp, srender, sint, sym_infer, range_str, pyrender
 from tinygrad.uop.ops import print_uops, range_start, multirange_str
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device, ProfileProgramEvent
-from tinygrad.renderer import ProgramSpec
+from tinygrad.renderer import Estimates
 from tinygrad.dtype import dtypes
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
@@ -39,7 +39,7 @@ def get_rewrites(t:RewriteTrace) -> list[dict]:
   for i,(k,v) in enumerate(zip(t.keys, t.rewrites)):
     steps = [create_step(s.name, ("/graph-rewrites", i, j), loc=s.loc, match_count=len(s.matches), code_line=printable(s.loc),
                          trace=k.tb if j==0 else None, depth=s.depth) for j,s in enumerate(v)]
-    if isinstance(k.ret, ProgramSpec):
+    if isinstance(k.ret, UOp) and k.ret.op is Ops.PROGRAM:
       steps.append(create_step("View UOp List", ("/uops", i, len(steps)), k.ret))
       steps.append(create_step("View Program", ("/code", i, len(steps)), k.ret))
       steps.append(create_step("View Disassembly", ("/asm", i, len(steps)), k.ret))
@@ -161,9 +161,10 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:
     name, fmt, key = e.name, [], None
     if (ref:=ref_map.get(name)) is not None:
       name = ctxs[ref]["name"]
-      if isinstance(p:=trace.keys[ref].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None:
-        flops = sym_infer(p.estimates.ops, var_vals:=ei.arg['var_vals'])/(t:=dur*1e-6)
-        membw, ldsbw = sym_infer(p.estimates.mem, var_vals)/t, sym_infer(p.estimates.lds, var_vals)/t
+      if isinstance(p:=trace.keys[ref].ret, UOp) and p.op is Ops.PROGRAM and (ei:=exec_points.get(p.src[0].arg.name)) is not None:
+        estimates = Estimates.from_uops(list(p.src[2].src), ignore_indexing=True)
+        flops = sym_infer(estimates.ops, var_vals:=ei.arg['var_vals'])/(t:=dur*1e-6)
+        membw, ldsbw = sym_infer(estimates.mem, var_vals)/t, sym_infer(estimates.lds, var_vals)/t
         fmt = [f"{flops*1e-9:.0f} GFLOPS" if flops < 1e14 else f"{flops*1e-12:.0f} TFLOPS",
               (f"{membw*1e-9:.0f} GB/s" if membw < 1e13 else f"{membw*1e-12:.0f} TB/s")+" mem",
               (f"{ldsbw*1e-9:.0f} GB/s" if ldsbw < 1e15 else f"{ldsbw*1e-12:.0f} TB/s")+" lds"]
@@ -425,10 +426,12 @@ def get_render(i:int, j:int, fmt:str) -> dict:
   data = ctxs[i]["steps"][j]["data"]
   if fmt == "graph-rewrites": return {"value":get_full_rewrite(trace.rewrites[i][j]), "content_type":"text/event-stream"}
   if fmt == "uops": return {"src":get_stdout(lambda: print_uops(data.uops or [])), "lang":"txt"}
-  if fmt == "code": return {"src":data.src, "lang":"cpp"}
+  # PROGRAM UOp: src[3].arg is source code
+  source_code = data.src[3].arg
+  if fmt == "code": return {"src":source_code, "lang":"cpp"}
   if fmt == "asm":
     compiler = Device[data.device].compiler
-    disasm_str = get_stdout(lambda: compiler.disassemble(compiler.compile(data.src)))
+    disasm_str = get_stdout(lambda: compiler.disassemble(compiler.compile(source_code)))
     ret:dict = {"src":disasm_str}
     if data.device.startswith("AMD"):
       with soft_err(lambda err: ret.update(err)):