Merge branch 'master' into remove_programspec

docs/developer/layout.md

@@ -28,7 +28,7 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi
 
 Transform the optimized ast into a linearized and rendered program.
 
-::: tinygrad.codegen.get_program
+::: 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)

test/external/external_benchmark_op_conv.py

@@ -1,5 +1,4 @@
 # ruff: noqa: E501 E712 F401
-from dataclasses import replace
 from tinygrad import dtypes, Device
 from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
 from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
@@ -89,7 +88,9 @@ renderer = Device.default.renderer
 allocator = Device.default.allocator
 
 ps = get_program(ast, renderer)
-cr = CompiledRunner(replace(ps, device=Device.DEFAULT))
+# 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_hip_compiler_bug.py

@@ -0,0 +1,225 @@
+# [<buf device:HIP size:1605632 dtype:dtypes.float>, <buf device:HIP size:301506 dtype:dtypes.float>, <buf device:HIP size:9408 dtype:dtypes.float>]
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer, CompiledRunner
+
+import ctypes
+import gpuctypes.hip as hip
+from tinygrad.helpers import to_char_p_p, init_c_var
+def get_bytes(arg, get_sz, get_str, check) -> bytes: return (sz := init_c_var(ctypes.c_size_t(), lambda x: check(get_sz(arg, ctypes.byref(x)))), ctypes.string_at(init_c_var(ctypes.create_string_buffer(sz.value), lambda x: check(get_str(arg, x))), size=sz.value))[1]  # noqa: E501
+def check(status):
+  if status != 0: raise RuntimeError(f"HIP Error {status}, {ctypes.string_at(hip.hipGetErrorString(status)).decode()}")
+def compile_hip(prg:str, arch="gfx1100") -> bytes:
+  check(hip.hiprtcCreateProgram(ctypes.byref(prog := hip.hiprtcProgram()), prg.encode(), "<null>".encode(), 0, None, None))
+  compile_options = [f'--offload-arch={arch}', '-I/opt/rocm/include']
+  status = hip.hiprtcCompileProgram(prog, len(compile_options), to_char_p_p([o.encode() for o in compile_options]))
+  if status != 0: raise RuntimeError(f"compile failed: {get_bytes(prog, hip.hiprtcGetProgramLogSize, hip.hiprtcGetProgramLog, check).decode()}")
+  return get_bytes(prog, hip.hiprtcGetCodeSize, hip.hiprtcGetCode, check)
+
+prefix = """
+typedef long unsigned int size_t;
+extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_id(unsigned int);
+extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_group_id(unsigned int);
+extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_size(unsigned int);
+typedef float float2 __attribute__((ext_vector_type(2)));
+static inline __attribute__((device)) float2 make_float2(float x, float y) { return {x, y}; }
+"""
+
+code = """
+extern "C" __attribute__((global))void r_2_8_7_7_4_8_3_7_7_4_4_2_2(float* data0, const float* data1, const float* data2) {
+  int gidx0 = __ockl_get_group_id(2); /* 2 */
+  int gidx1 = __ockl_get_group_id(1); /* 8 */
+  int gidx2 = __ockl_get_group_id(0); /* 49 */
+  int lidx4 = __ockl_get_local_id(1); /* 4 */
+  int lidx5 = __ockl_get_local_id(0); /* 8 */
+  float2 acc0 = make_float2(0.0f,0.0f);
+  float2 acc1 = make_float2(0.0f,0.0f);
+  float2 acc2 = make_float2(0.0f,0.0f);
+  float2 acc3 = make_float2(0.0f,0.0f);
+  float2 acc4 = make_float2(0.0f,0.0f);
+  float2 acc5 = make_float2(0.0f,0.0f);
+  float2 acc6 = make_float2(0.0f,0.0f);
+  float2 acc7 = make_float2(0.0f,0.0f);
+  float2 acc8 = make_float2(0.0f,0.0f);
+  float2 acc9 = make_float2(0.0f,0.0f);
+  float2 acc10 = make_float2(0.0f,0.0f);
+  float2 acc11 = make_float2(0.0f,0.0f);
+  float2 acc12 = make_float2(0.0f,0.0f);
+  float2 acc13 = make_float2(0.0f,0.0f);
+  float2 acc14 = make_float2(0.0f,0.0f);
+  float2 acc15 = make_float2(0.0f,0.0f);
+  float2 acc16 = make_float2(0.0f,0.0f);
+  float2 acc17 = make_float2(0.0f,0.0f);
+  float2 acc18 = make_float2(0.0f,0.0f);
+  float2 acc19 = make_float2(0.0f,0.0f);
+  float2 acc20 = make_float2(0.0f,0.0f);
+  float2 acc21 = make_float2(0.0f,0.0f);
+  float2 acc22 = make_float2(0.0f,0.0f);
+  float2 acc23 = make_float2(0.0f,0.0f);
+  float2 acc24 = make_float2(0.0f,0.0f);
+  float2 acc25 = make_float2(0.0f,0.0f);
+  float2 acc26 = make_float2(0.0f,0.0f);
+  float2 acc27 = make_float2(0.0f,0.0f);
+  float2 acc28 = make_float2(0.0f,0.0f);
+  float2 acc29 = make_float2(0.0f,0.0f);
+  float2 acc30 = make_float2(0.0f,0.0f);
+  float2 acc31 = make_float2(0.0f,0.0f);
+  int alu0 = (gidx2/7);
+  int alu1 = (gidx2%7);
+  int alu2 = (alu1*32);
+  int alu3 = (lidx5*4);
+  int alu4 = ((gidx0*802816)+(gidx1*100352)+(alu0*1792)+(alu1*16)+(lidx4*448)+(lidx5*2));
+  for (int ridx0 = 0; ridx0 < 3; ridx0++) {
+    for (int ridx1 = 0; ridx1 < 7; ridx1++) {
+      int alu5 = ((alu0*(-32))+(lidx4*(-8))+(ridx1*(-1)));
+      bool alu6 = (alu5<(-2));
+      bool alu7 = (alu5<0);
+      bool alu8 = (((alu0*32)+(lidx4*8)+ridx1)<221);
+      for (int ridx2 = 0; ridx2 < 7; ridx2++) {
+        int alu9 = ((gidx0*150528)+(ridx0*50176)+(alu0*7168)+(lidx4*1792)+(ridx1*224)+alu2+alu3+ridx2);
+        int alu10 = ((alu1*(-32))+(lidx5*(-4))+(ridx2*(-1)));
+        bool alu11 = (alu10<(-2));
+        float val0 = 0.0f;
+        if ((alu6*alu11)) { val0 = data1[alu9+(-675)]; }
+        float val1 = 0.0f;
+        if ((alu7*alu11)) { val1 = data1[alu9+(-227)]; }
+        float val2 = 0.0f;
+        if (alu11) { val2 = data1[alu9+221]; }
+        float val3 = 0.0f;
+        if ((alu8*alu11)) { val3 = data1[alu9+669]; }
+        bool alu12 = (alu10<0);
+        bool alu13 = ((alu2+alu3+ridx2)<225);
+        float val4 = 0.0f;
+        if ((alu6*alu12*alu13)) { val4 = data1[alu9+(-673)]; }
+        float val5 = 0.0f;
+        if ((alu7*alu12*alu13)) { val5 = data1[alu9+(-225)]; }
+        float val6 = 0.0f;
+        if ((alu12*alu13)) { val6 = data1[alu9+223]; }
+        float val7 = 0.0f;
+        if ((alu8*alu12*alu13)) { val7 = data1[alu9+671]; }
+        int alu14 = ((gidx1*1176)+(ridx0*49)+(ridx1*7)+ridx2);
+        float val8 = data2[alu14];
+        float val9 = data2[alu14+147];
+        float val10 = data2[alu14+294];
+        float val11 = data2[alu14+441];
+        float val12 = data2[alu14+588];
+        float val13 = data2[alu14+735];
+        float val14 = data2[alu14+882];
+        float val15 = data2[alu14+1029];
+        (acc0).x = ((val0*val8)+(acc0).x);
+        (acc1).x = ((val0*val9)+(acc1).x);
+        (acc2).x = ((val0*val10)+(acc2).x);
+        (acc3).x = ((val0*val11)+(acc3).x);
+        (acc4).x = ((val1*val8)+(acc4).x);
+        (acc5).x = ((val1*val9)+(acc5).x);
+        (acc6).x = ((val1*val10)+(acc6).x);
+        (acc7).x = ((val1*val11)+(acc7).x);
+        (acc8).x = ((val2*val8)+(acc8).x);
+        (acc9).x = ((val2*val9)+(acc9).x);
+        (acc10).x = ((val2*val10)+(acc10).x);
+        (acc11).x = ((val2*val11)+(acc11).x);
+        (acc12).x = ((val3*val8)+(acc12).x);
+        (acc13).x = ((val3*val9)+(acc13).x);
+        (acc14).x = ((val3*val10)+(acc14).x);
+        (acc15).x = ((val3*val11)+(acc15).x);
+        (acc16).x = ((val0*val12)+(acc16).x);
+        (acc17).x = ((val0*val13)+(acc17).x);
+        (acc18).x = ((val0*val14)+(acc18).x);
+        (acc19).x = ((val0*val15)+(acc19).x);
+        (acc20).x = ((val1*val12)+(acc20).x);
+        (acc21).x = ((val1*val13)+(acc21).x);
+        (acc22).x = ((val1*val14)+(acc22).x);
+        (acc23).x = ((val1*val15)+(acc23).x);
+        (acc24).x = ((val2*val12)+(acc24).x);
+        (acc25).x = ((val2*val13)+(acc25).x);
+        (acc26).x = ((val2*val14)+(acc26).x);
+        (acc27).x = ((val2*val15)+(acc27).x);
+        (acc28).x = ((val3*val12)+(acc28).x);
+        (acc29).x = ((val3*val13)+(acc29).x);
+        (acc30).x = ((val3*val14)+(acc30).x);
+        (acc31).x = ((val3*val15)+(acc31).x);
+        (acc0).y = ((val4*val8)+(acc0).y);
+        (acc1).y = ((val4*val9)+(acc1).y);
+        (acc2).y = ((val4*val10)+(acc2).y);
+        (acc3).y = ((val4*val11)+(acc3).y);
+        (acc4).y = ((val5*val8)+(acc4).y);
+        (acc5).y = ((val5*val9)+(acc5).y);
+        (acc6).y = ((val5*val10)+(acc6).y);
+        (acc7).y = ((val5*val11)+(acc7).y);
+        (acc8).y = ((val6*val8)+(acc8).y);
+        (acc9).y = ((val6*val9)+(acc9).y);
+        (acc10).y = ((val6*val10)+(acc10).y);
+        (acc11).y = ((val6*val11)+(acc11).y);
+        (acc12).y = ((val7*val8)+(acc12).y);
+        (acc13).y = ((val7*val9)+(acc13).y);
+        (acc14).y = ((val7*val10)+(acc14).y);
+        (acc15).y = ((val7*val11)+(acc15).y);
+        (acc16).y = ((val4*val12)+(acc16).y);
+        (acc17).y = ((val4*val13)+(acc17).y);
+        (acc18).y = ((val4*val14)+(acc18).y);
+        (acc19).y = ((val4*val15)+(acc19).y);
+        (acc20).y = ((val5*val12)+(acc20).y);
+        (acc21).y = ((val5*val13)+(acc21).y);
+        (acc22).y = ((val5*val14)+(acc22).y);
+        (acc23).y = ((val5*val15)+(acc23).y);
+        (acc24).y = ((val6*val12)+(acc24).y);
+        (acc25).y = ((val6*val13)+(acc25).y);
+        (acc26).y = ((val6*val14)+(acc26).y);
+        (acc27).y = ((val6*val15)+(acc27).y);
+        (acc28).y = ((val7*val12)+(acc28).y);
+        (acc29).y = ((val7*val13)+(acc29).y);
+        (acc30).y = ((val7*val14)+(acc30).y);
+        (acc31).y = ((val7*val15)+(acc31).y);
+      }
+    }
+  }
+  *((float2*)(data0+alu4)) = acc0;
+  *((float2*)(data0+alu4+12544)) = acc1;
+  *((float2*)(data0+alu4+25088)) = acc2;
+  *((float2*)(data0+alu4+37632)) = acc3;
+  *((float2*)(data0+alu4+112)) = acc4;
+  *((float2*)(data0+alu4+12656)) = acc5;
+  *((float2*)(data0+alu4+25200)) = acc6;
+  *((float2*)(data0+alu4+37744)) = acc7;
+  *((float2*)(data0+alu4+224)) = acc8;
+  *((float2*)(data0+alu4+12768)) = acc9;
+  *((float2*)(data0+alu4+25312)) = acc10;
+  *((float2*)(data0+alu4+37856)) = acc11;
+  *((float2*)(data0+alu4+336)) = acc12;
+  *((float2*)(data0+alu4+12880)) = acc13;
+  *((float2*)(data0+alu4+25424)) = acc14;
+  *((float2*)(data0+alu4+37968)) = acc15;
+  *((float2*)(data0+alu4+50176)) = acc16;
+  *((float2*)(data0+alu4+62720)) = acc17;
+  *((float2*)(data0+alu4+75264)) = acc18;
+  *((float2*)(data0+alu4+87808)) = acc19;
+  *((float2*)(data0+alu4+50288)) = acc20;
+  *((float2*)(data0+alu4+62832)) = acc21;
+  *((float2*)(data0+alu4+75376)) = acc22;
+  *((float2*)(data0+alu4+87920)) = acc23;
+  *((float2*)(data0+alu4+50400)) = acc24;
+  *((float2*)(data0+alu4+62944)) = acc25;
+  *((float2*)(data0+alu4+75488)) = acc26;
+  *((float2*)(data0+alu4+88032)) = acc27;
+  *((float2*)(data0+alu4+50512)) = acc28;
+  *((float2*)(data0+alu4+63056)) = acc29;
+  *((float2*)(data0+alu4+75600)) = acc30;
+  *((float2*)(data0+alu4+88144)) = acc31;
+}
+"""
+
+dev = "HIP"
+lib = Device[dev].compiler.compile(prefix+code)
+#lib = compile_hip(code)
+b0 = Buffer(dev, 1605632, dtypes.float)
+b1 = Buffer(dev, 301506, dtypes.float)
+b2 = Buffer(dev, 9408, dtypes.float)
+print(hex(b0._buf.value), hex(b0._buf.value+1605632*4))
+print(hex(b1._buf.value))
+print(hex(b2._buf.value))
+#prg = CompiledRunner("r_2_8_7_7_4_8_3_7_7_4_4_2_2", "", dev, [7, 1, 1], [8, 4, 1], precompiled=lib)
+prg = CompiledRunner("r_2_8_7_7_4_8_3_7_7_4_4_2_2", "", dev, [49, 8, 2], [8, 4, 1], precompiled=lib)
+print("compiled")
+prg([b0, b1, b2], {})
+print("ran")
+Device[dev].synchronize()
+print("sync")

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

@@ -63,7 +63,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
     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=())
   prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer())
-  prog = PythonProgram("run", PythonCompiler().compile(prg.src))
+  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 or "").tolist()[0]
 

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
@@ -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, 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

@@ -1,6 +1,5 @@
 import unittest
 import numpy as np
-from dataclasses import replace
 from tinygrad.device import Buffer, Device, is_dtype_supported
 from tinygrad.dtype import dtypes, ConstType
 from tinygrad.engine.realize import CompiledRunner, get_program
@@ -18,8 +17,8 @@ def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
   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]
-  prg = replace(prg, device=Device.DEFAULT)
-  if local_size is not None: prg = replace(prg, 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]
@@ -72,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)

test/test_uops.py

@@ -15,7 +15,6 @@ 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
-from dataclasses import replace
 
 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   sink = UOp.group(*u)
@@ -27,7 +26,9 @@ def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
 
 def _uops_to_prg(uops_list):
   prg = get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer)
-  return CompiledRunner(replace(prg, device=Device.DEFAULT))
+  # 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,7 +2,6 @@ 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.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
@@ -166,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
@@ -191,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:
@@ -206,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:
@@ -216,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:
@@ -230,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_elf.py

@@ -1,5 +1,5 @@
 import unittest, subprocess, platform
-from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler
+from tinygrad.runtime.ops_cpu import ClangJITCompiler
 from tinygrad.runtime.support.elf import elf_loader
 
 class TestElfLoader(unittest.TestCase):

tinygrad/codegen/__init__.py

@@ -1,12 +1,11 @@
 from typing import cast
 import itertools
-from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, getenv, TracingKey
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
+from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
-from tinygrad.renderer import Renderer, ProgramSpec
+from tinygrad.renderer import Renderer
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.helpers import panic
-from tinygrad.codegen.opt import Opt
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
@@ -29,8 +28,6 @@ pm_syntactic_sugar = PatternMatcher([
 def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
   if ren is None: ren = Renderer()
 
-  if getenv("VIZ"): graph_rewrite(sink, PatternMatcher([]), name="View Base AST")
-  if DEBUG >= 5: print(pyrender(sink))
   if SPEC: type_verify(sink, kernel_spec)
 
   # preprocess
@@ -135,40 +132,12 @@ 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:Renderer, prg:UOp, source:UOp) -> UOp|None:
-  if ctx.compiler is None: return None
-  lib = ctx.compiler.compile_cached(source.arg)
-  return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),))
-
 pm_to_program = PatternMatcher([
   (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),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
 ])
 
-@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, opts:list[Opt]|None=None) -> ProgramSpec:
-  """
-  Transform an AST into a ProgramSpec. May trigger BEAM search.
-
-  Args:
-    ast: The Ops.SINK rooted AST
-    renderer: The renderer used to generate the code
-
-  Returns:
-    The ProgramSpec of the program.
-  """
-
-  # fix up KernelInfo
-  if opts is not None:
-    assert ast.arg is None, "can't apply opts if sink has an arg"
-    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
-  if ast.arg is None: ast = ast.replace(arg=KernelInfo())
-
-  # rewrite to prg
-  full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
-  prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.device)))
-  prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
-
-  # create the ProgramSpec
-  return ProgramSpec.from_uop(prg)
+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, UOp(Ops.DEVICE, arg=ren.device)))
+  return graph_rewrite(sink, pm_to_program, ctx=ren, name="linearize/render")

tinygrad/codegen/opt/search.py

@@ -1,14 +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
-from tinygrad.codegen import 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)]
@@ -35,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 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)
@@ -59,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
@@ -67,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 = p.lib if p.lib is not None else 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:
@@ -155,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
@@ -168,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/device.py

@@ -278,7 +278,7 @@ class Compiler:
   def disassemble(self, lib:bytes): pass
 
 @dataclass(frozen=True)
-class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial|None; ctrl_var:ContextVar|None = None # noqa: E702
+class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial; ctrl_var:ContextVar|None = None # noqa: E702
 
 @dataclass(frozen=True)
 class CompilerSet: cset:list[CompilerPair]; ctrl_var:ContextVar|None = None # noqa: E702
@@ -290,25 +290,21 @@ class Compiled:
     self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id
 
     self.comps_ctrl_var = compilers.ctrl_var if compilers is not None else None
-    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial|None]]] = {}
-    self.cached_pair:dict[Any, tuple[Renderer, Compiler|None]] = {}
+    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial]]] = {}
+    self.cached_pair:dict[Any, tuple[Renderer, Compiler]] = {}
     for cpair in (compilers.cset if compilers is not None else [CompilerPair(Renderer, Compiler)]):
-      self.comp_sets[self._compiler_name(cpair.renderer, cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
+      self.comp_sets[self._compiler_name(cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
 
   @property
   def renderer(self) -> Renderer: return self._select_compiler_pair()[0]
 
   @property
-  def compiler(self) -> Compiler:
-    if (ret:=self.renderer.compiler or self._select_compiler_pair()[1]) is None: raise RuntimeError(f"no compiler for {self.device}")
-    return ret
+  def compiler(self) -> Compiler: return self._select_compiler_pair()[1]
 
-  def _compiler_name(self, r:type[Renderer]|functools.partial, c:type[Compiler]|functools.partial|None) -> str:
-    devname = self.device.split(':')[0].upper()
-    if c is None: return unwrap_class_type(r).__name__.upper().removesuffix("RENDERER").removeprefix(devname) or devname
-    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname) or devname
+  def _compiler_name(self, c:type[Compiler]|functools.partial) -> str:
+    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname
 
-  def _select_compiler_pair(self) -> tuple[Renderer, Compiler|None]:
+  def _select_compiler_pair(self) -> tuple[Renderer, Compiler]:
     # select forced compiler from global env var.
     forced_comps = set([self.comp_sets[val][1]] if self.comps_ctrl_var is not None and (val:=self.comps_ctrl_var.value) else [])
 
@@ -402,7 +398,7 @@ def enumerate_devices_str() -> Generator[str, None, None]:
             # d.renderer, d.compiler = r(), c()
             with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
             if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
-            set_text = f'({cc_ctrl_var.key}={d._compiler_name(r, c)} to make default)' if cc_ctrl_var is not None else ''
+            set_text = f'({cc_ctrl_var.key}={d._compiler_name(c)} to make default)' if cc_ctrl_var is not None else ''
             default_text = '(default)' if type(default_compiler) is type(d.compiler) else set_text
             compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
             any_works = True

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.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,13 +1,46 @@
 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, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
-from tinygrad.helpers import unwrap
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer
+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 ProgramSpec, Estimates
-from tinygrad.codegen import get_program
+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.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 PROGRAM UOp. May trigger BEAM search.
+
+  Args:
+    ast: The Ops.SINK rooted AST
+    renderer: The renderer used to generate the code
+
+  Returns:
+    The PROGRAM UOp with structure (SINK, DEVICE, LINEAR, SOURCE).
+  """
+
+  if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
+  if DEBUG >= 5: print(pyrender(ast))
+
+  # linearize
+  if opts is not None:
+    assert ast.arg is None, "can't apply opts if sink has an arg"
+    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)
+
+  # print
+  if DEBUG >= 6: print_uops(list(prg.src[2].src))  # LINEAR is src[2]
+
+  return prg
 
 # **************** Runners ****************
 
@@ -36,29 +69,44 @@ 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)
-    if p.lib is None:
-      with cpu_profile(TracingKey(f"compile {p.name}", (p.function_name,)), "TINY"):
-        p = replace(p, lib=Device[p.device].compiler.compile_cached(p.src))
-    self.p:ProgramSpec = p
-    assert self.p.lib is not None
-    if DEBUG >= 7: Device[p.device].compiler.disassemble(self.p.lib)
-    self._prg = Device[p.device].runtime(p.function_name, self.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 = {}
+    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 Device[self.p.device].renderer.has_local 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)
-    return self._prg(*[x._buf for x in rawbufs], global_size=tuple(global_size), local_size=tuple(local_size) if local_size else None,
-                     vals=tuple(var_vals[k.expr] for k in self.p.vars), wait=wait)
+      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)
+      assert len(global_size) == 3, "global size must have len 3"
+    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.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)
@@ -115,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)
-    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/helpers.py

@@ -115,12 +115,12 @@ def suppress_finalizing(func):
       if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
   return wrapper
 
-def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]|None]], err_msg:str, cache:dict|None=None):
+def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg:str, cache:dict|None=None) -> tuple[T,...]|T:
   excs = []
   for typ in candidates:
     if cache is not None and typ in cache: return cache[typ]
     try:
-      x = tuple([cast(Callable, t)() if t is not None else None for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
+      x = tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
       if cache is not None: cache[typ] = x
       return x
     except Exception as e: excs.append(e)

tinygrad/renderer/__init__.py

@@ -1,13 +1,10 @@
 from __future__ import annotations
-from typing import Callable, cast, TYPE_CHECKING
-import functools
-from dataclasses import dataclass, field
-from tinygrad.helpers import to_function_name, dedup, prod, DEBUG
-from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher, print_uops
+from typing import Callable, cast
+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
-if TYPE_CHECKING: from tinygrad.device import Compiler
 
 @dataclass(frozen=True)
 class Estimates:
@@ -58,74 +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
-
-  # filled in from uops (via from_uop)
-  global_size:list[int]=field(default_factory=lambda: [1,1,1])
-  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)
-
-  @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]
-    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
-
-  @staticmethod
-  def from_uop(prg:UOp) -> ProgramSpec:
-    """Construct ProgramSpec from a PROGRAM UOp."""
-    assert prg.op is Ops.PROGRAM, f"expected PROGRAM, got {prg.op}"
-    # SINK/DEVICE/LINEAR/SOURCE/BINARY?
-    sink, device, linear, source = prg.src[:4]
-    lib = prg.src[4].arg if len(prg.src) > 4 else None
-    uops = list(linear.src)
-    if DEBUG >= 6: print_uops(uops)  # LINEAR is src[2]
-
-    # single pass through the uops to extract metadata
-    _vars: list[Variable] = []
-    _globals: list[int] = []
-    outs: list[int] = []
-    ins: list[int] = []
-    global_size: list[int] = [1, 1, 1]
-    local_size: list[int]|None = [1, 1, 1]
-    for u in uops:
-      if u.op is Ops.DEFINE_VAR: _vars.append(u)
-      if u.op is Ops.DEFINE_GLOBAL: _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: (outs if u.op is Ops.STORE else ins).append(buf.arg)
-        # TODO: can else happen?
-      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
-        # 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())
-
-    return ProgramSpec(sink.arg.name, source.arg, device.arg, sink, uops, lib, global_size, local_size,
-                       sorted(_vars, key=lambda v: v.arg), sorted(dedup(_globals)), sorted(dedup(outs)), sorted(dedup(ins)))
-
 class Renderer:
   device: str = ""
   suffix: str = ""
@@ -142,7 +71,6 @@ class Renderer:
   pre_matcher: PatternMatcher|None = None
   extra_matcher: PatternMatcher|None = None
   code_for_op: dict[Ops, Callable] = {}
-  compiler: Compiler|None = None
 
   def __reduce__(self): return self.__class__, ()
   def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")

tinygrad/renderer/cstyle.py

@@ -8,7 +8,6 @@ from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, trunc
 from tinygrad.renderer import Renderer
 from tinygrad.codegen.late.devectorizer import no_vectorized_alu
 
-
 base_rewrite = PatternMatcher([
   (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
   (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"),
@@ -279,11 +278,6 @@ class ClangRenderer(CStyleLanguage):
     defines = '\n'.join(self._render_defines(uops))
     return defines + "\n" + self._render_body(function_name, kernel, bufs, uops, prefix) + "\n" + self._render_entry(function_name, bufs)
 
-class ClangJITRenderer(ClangRenderer):
-  def __init__(self):
-    from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler
-    self.compiler = ClangJITCompiler()
-
 class OpenCLRenderer(CStyleLanguage):
   device = "CL"
 
@@ -334,10 +328,7 @@ class IntelRenderer(OpenCLRenderer):
 class MetalRenderer(CStyleLanguage):
   device = "METAL"
   shared_max = 32768
-  def __init__(self):
-    self.tensor_cores = tc.metal if hasattr(os, 'uname') and os.uname().machine == "arm64" else []
-    from tinygrad.runtime.ops_metal import MetalCompiler
-    self.compiler = MetalCompiler()
+  def __init__(self): self.tensor_cores = tc.metal if hasattr(os, 'uname') and os.uname().machine == "arm64" else []
 
   # language options
   kernel_typedef = "kernel void"
@@ -449,18 +440,6 @@ class CUDARenderer(CStyleLanguage):
 
     return super().render_kernel(function_name, kernel, bufs, uops, prefix=prefix)
 
-class CUDACUDARenderer(CUDARenderer):
-  def __init__(self, arch:str):
-    super().__init__(arch)
-    from tinygrad.runtime.support.compiler_cuda import CUDACompiler
-    self.compiler = CUDACompiler(arch)
-
-class CUDANVCCRenderer(CUDARenderer):
-  def __init__(self, arch:str):
-    super().__init__(arch)
-    from tinygrad.runtime.support.compiler_cuda import NVCCCompiler
-    self.compiler = NVCCCompiler(arch)
-
 class AMDRenderer(CStyleLanguage):
   device = "AMD"
   shared_max = 65536
@@ -553,32 +532,6 @@ class AMDRenderer(CStyleLanguage):
   for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d;\n}""")
     return super().render_kernel(function_name, kernel, bufs, uops, prefix)
 
-class AMDHIPRenderer(AMDRenderer):
-  def __init__(self, arch:str):
-    super().__init__(arch)
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    self.compiler = HIPCompiler(arch)
-
-class AMDHIPCCRenderer(AMDRenderer):
-  def __init__(self, arch:str):
-    super().__init__(arch)
-    from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
-    self.compiler = HIPCCCompiler(arch)
-
 class NVRenderer(CUDARenderer): device = "NV"
-
-class NVNVRenderer(NVRenderer):
-  def __init__(self, arch:str):
-    super().__init__(arch)
-    from tinygrad.runtime.support.compiler_cuda import NVCompiler
-    self.compiler = NVCompiler(arch)
-
 class HIPRenderer(AMDRenderer): device = "HIP"
-
-class HIPHIPRenderer(HIPRenderer):
-  def __init__(self, arch:str):
-    super().__init__(arch)
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    self.compiler = HIPCompiler(arch)
-
 class QCOMRenderer(OpenCLRenderer): device = "QCOM"

tinygrad/renderer/llvmir.py

@@ -143,9 +143,6 @@ class LLVMRenderer(Renderer):
   if AMX: tensor_cores = tc.amx
 
   extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast
-  def __init__(self):
-    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
-    self.compiler = CPULLVMCompiler()
   def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops)))
   def _render_footer(self, uops: list[UOp]) -> str: return 'attributes #0 = { alwaysinline nounwind "no-builtins" "no-trapping-math"="true" }'
   def _render_fn(self, name:str, args:list[tuple[str,DType]], kernel:list[str], prefix:list[str]|None=None) -> str:
@@ -257,9 +254,7 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
                   f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"']
     return 'attributes #0 = { ' + ' '.join(attributes) + ' }'
   def __init__(self, arch:str):
-    from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler
     self.arch = arch
-    self.compiler = AMDLLVMCompiler(arch)
     self.tensor_cores = AMDRenderer.get_tensor_cores(arch)
     self.is_cdna = AMDRenderer.is_cdna(arch)
     self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])

tinygrad/renderer/nir.py

@@ -245,11 +245,6 @@ class LVPRenderer(NIRRenderer):
     srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz:
     setattr(self, "param_idx", self.param_idx+sz))(lambda self,b,x,sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))
 
-  def __init__(self):
-    from tinygrad.runtime.support.compiler_mesa import LVPCompiler
-    super().__init__()
-    self.compiler = LVPCompiler()
-
   def prerender(self, uops:list[UOp]):
     super().prerender(uops)
     self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)])

tinygrad/renderer/ptx.py

@@ -240,17 +240,3 @@ class PTXRenderer(Renderer):
 
       if u.op is Ops.SPECIAL: kernel = [f".reg .u32 %{u.arg};"] + kernel
     return self.render_kernel(kernel, name, bufs, c.items(), uops)
-
-class CUDAPTXRenderer(PTXRenderer):
-  def __init__(self, arch:str):
-    super().__init__(arch, "CUDA")
-    from tinygrad.runtime.support.compiler_cuda import PTXCompiler
-    self.compiler = PTXCompiler(arch)
-  def __reduce__(self): return self.__class__, (self.arch,)
-
-class NVPTXRenderer(PTXRenderer):
-  def __init__(self, arch:str):
-    super().__init__(arch, "NV")
-    from tinygrad.runtime.support.compiler_cuda import NVPTXCompiler
-    self.compiler = NVPTXCompiler(arch)
-  def __reduce__(self): return self.__class__, (self.arch,)

tinygrad/renderer/wgsl.py

@@ -1,7 +1,6 @@
 from tinygrad.dtype import DType, PtrDType, dtypes, AddrSpace
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
 from tinygrad.renderer.cstyle import CStyleLanguage, base_rewrite, extra_pm
-from tinygrad.device import Compiler
 from tinygrad.helpers import strip_parens
 
 def sign_extend(val:UOp, sext_am:int):
@@ -47,7 +46,6 @@ class WGSLRenderer(CStyleLanguage):
   global_max = (65535, 65535, 65535)
   local_max = (256, 256, 64)
   code_for_workitem = {"g": lambda x: f"i32(gindex.{'xyz'[int(x)]})", "l": lambda x: f"i32(lindex.{'xyz'[int(x)]})"}
-  def __init__(self): self.compiler = Compiler()
   extra_matcher = wgsl_matcher
   supports_float4 = False
   barrier = "workgroupBarrier();"

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

@@ -9,10 +9,11 @@ from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet, CompilerPair
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar
 from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, ceildiv
-from tinygrad.renderer.cstyle import AMDHIPRenderer, AMDHIPCCRenderer
+from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
 from tinygrad.runtime.autogen.am import am
+from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler, AMDLLVMCompiler
 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
@@ -930,9 +931,9 @@ class AMDDevice(HCQCompiled):
     max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
     self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))
 
-    compilers = CompilerSet([CompilerPair(functools.partial(AMDHIPRenderer, self.arch), None),
-                             CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), None, AMD_LLVM),
-                             CompilerPair(functools.partial(AMDHIPCCRenderer, self.arch), None)], ctrl_var=AMD_CC)
+    compilers = CompilerSet([CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch)),
+                             CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch), AMD_LLVM),
+                             CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCCCompiler, self.arch))], ctrl_var=AMD_CC)
 
     super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                      functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),

tinygrad/runtime/ops_cpu.py

@@ -5,9 +5,10 @@ from tinygrad.helpers import CPU_CC, CPU_LVP, CPU_LLVM
 from tinygrad.device import BufferSpec, DMACPURef, CompilerSet, CompilerPair
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.hcq import CLikeArgsState
-from tinygrad.renderer.cstyle import ClangJITRenderer
+from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import LLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
+from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler
 from tinygrad.runtime.support.compiler_mesa import LVPCompiler
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.uop.ops import sint
@@ -135,6 +136,6 @@ class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
     self.tasks:queue.Queue = queue.Queue()
     CPUWorker(self, self.tasks, thread_id=0).start()
-    compilers = CompilerSet([CompilerPair(ClangJITRenderer, None), CompilerPair(LLVMRenderer, None, ctrl_var=CPU_LLVM),
-                             CompilerPair(LVPRenderer, None, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC)
+    compilers = CompilerSet([CompilerPair(ClangRenderer, ClangJITCompiler), CompilerPair(LLVMRenderer, CPULLVMCompiler, ctrl_var=CPU_LLVM),
+                             CompilerPair(LVPRenderer, LVPCompiler, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC)
     super().__init__(device, CPUAllocator(self), compilers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue)

tinygrad/runtime/ops_cuda.py

@@ -2,10 +2,10 @@ from __future__ import annotations
 import ctypes, functools
 from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing, CUDA_CC, CUDA_PTX
 from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerPair, CompilerSet
-from tinygrad.renderer.cstyle import CUDACUDARenderer, CUDANVCCRenderer
-from tinygrad.renderer.ptx import CUDAPTXRenderer
+from tinygrad.renderer.cstyle import CUDARenderer
+from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.runtime.autogen import cuda
-from tinygrad.runtime.support.compiler_cuda import pretty_ptx
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler, NVCCCompiler
 if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl  # noqa: F401  # pylint: disable=unused-import
 if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported
 
@@ -117,9 +117,9 @@ class CUDADevice(Compiled):
     CUDADevice.devices.append(self)
 
     from tinygrad.runtime.graph.cuda import CUDAGraph
-    compilers = CompilerSet([CompilerPair(functools.partial(CUDACUDARenderer, self.arch), None),
-                             CompilerPair(functools.partial(CUDAPTXRenderer, self.arch), None, CUDA_PTX),
-                             CompilerPair(functools.partial(CUDANVCCRenderer, self.arch), None)], ctrl_var=CUDA_CC)
+    compilers = CompilerSet([CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(CUDACompiler, self.arch)),
+                             CompilerPair(functools.partial(PTXRenderer, self.arch), functools.partial(PTXCompiler, self.arch), CUDA_PTX),
+                             CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(NVCCCompiler, self.arch))], ctrl_var=CUDA_CC)
     super().__init__(device, CUDAAllocator(self), compilers, functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph)
 
   def synchronize(self):

tinygrad/runtime/ops_dsp.py

@@ -83,7 +83,7 @@ class DSPProgram:
   def __init__(self, dev:DSPDevice, name:str, lib:bytes):
     self.dev, self.lib = dev, lib
 
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
+  def __call__(self, *bufs, vals:tuple[int, ...]=(), wait=False):
     if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}")
 
     pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*4)), off_mv:=memoryview(bytearray(len(bufs)*4))],
@@ -289,7 +289,7 @@ class MockDSPRenderer(DSPRenderer):
 
 class MockDSPProgram:
   def __init__(self, name:str, lib:bytes): self.lib = lib
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
+  def __call__(self, *bufs, vals:tuple[int, ...]=(), wait=False):
     with tempfile.NamedTemporaryFile(suffix=".out") as dsp_lib:
       dsp_lib.write(self.lib)
       dsp_lib.flush()

tinygrad/runtime/ops_hip.py

@@ -2,7 +2,8 @@ import ctypes, functools
 from tinygrad.helpers import init_c_var, mv_address, init_c_struct_t, getenv
 from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet, CompilerPair
 from tinygrad.runtime.autogen import hip
-from tinygrad.renderer.cstyle import HIPHIPRenderer
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
+from tinygrad.renderer.cstyle import HIPRenderer
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
 def check(status):
@@ -14,7 +15,7 @@ class HIPDevice(Compiled):
     self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode()
     self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t(), lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
 
-    compilers = CompilerSet([CompilerPair(functools.partial(HIPHIPRenderer, self.arch), None)])
+    compilers = CompilerSet([CompilerPair(functools.partial(HIPRenderer, self.arch), functools.partial(HIPCompiler, self.arch))])
     super().__init__(device, HIPAllocator(self), compilers, functools.partial(HIPProgram, self))
   def synchronize(self):
     check(hip.hipSetDevice(self.device_id))

tinygrad/runtime/ops_metal.py

@@ -44,7 +44,7 @@ class MetalDevice(Compiled):
     from tinygrad.runtime.graph.metal import MetalGraph
     # NOTE: GitHub CI macOS runners use paravirtualized metal which is broken with graph.
     # This can be reproduced locally with any virtualization software (like utm) that can create macOS VMs with apple's own virtualization framework.
-    super().__init__(device, MetalAllocator(self), CompilerSet([CompilerPair(MetalRenderer, None)]),
+    super().__init__(device, MetalAllocator(self), CompilerSet([CompilerPair(MetalRenderer, MetalCompiler), CompilerPair(MetalRenderer, Compiler)]),
       functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None)
 
   def synchronize(self):

tinygrad/runtime/ops_nv.py

@@ -8,8 +8,9 @@ from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU
 from tinygrad.uop.ops import sint
 from tinygrad.device import BufferSpec, CompilerPair, CompilerSet
 from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, NV_CC, NV_PTX, NV_NAK
-from tinygrad.renderer.ptx import CUDAPTXRenderer, NVPTXRenderer
-from tinygrad.renderer.cstyle import NVNVRenderer, CUDACUDARenderer
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.cstyle import NVRenderer
+from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler
 from tinygrad.runtime.support.compiler_mesa import NAKCompiler
 from tinygrad.runtime.autogen import nv_570, nv_580, pci, mesa
 from tinygrad.runtime.support.elf import elf_loader
@@ -582,9 +583,9 @@ class NVDevice(HCQCompiled[HCQSignal]):
     self.arch: str = "sm_120" if self.sm_version==0xa04 else f"sm_{(self.sm_version>>8)&0xff}{(val>>4) if (val:=self.sm_version&0xff) > 0xf else val}"
     self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf)
 
-    nvr, ptxr = (CUDACUDARenderer, CUDAPTXRenderer) if MOCKGPU else (NVNVRenderer, NVPTXRenderer)
-    compilers = CompilerSet(ctrl_var=NV_CC, cset=[CompilerPair(functools.partial(nvr, self.arch), None),
-       CompilerPair(functools.partial(ptxr, self.arch), None, NV_PTX),
+    cucc, ptxcc = (CUDACompiler, PTXCompiler) if MOCKGPU else (NVCompiler, NVPTXCompiler)
+    compilers = CompilerSet(ctrl_var=NV_CC, cset=[CompilerPair(functools.partial(NVRenderer, self.arch),functools.partial(cucc, self.arch)),
+       CompilerPair(functools.partial(PTXRenderer, self.arch, device="NV"), functools.partial(ptxcc, self.arch), NV_PTX),
        CompilerPair(functools.partial(NAKRenderer, dev=self), functools.partial(NAKCompiler, self.arch, self.max_warps_per_sm), NV_NAK)])
     super().__init__(device, NVAllocator(self), compilers, functools.partial(NVProgram, self), HCQSignal, NVComputeQueue, NVCopyQueue)
 

tinygrad/runtime/ops_python.py

@@ -213,14 +213,10 @@ class PythonProgram:
         i += 1
     return time.perf_counter() - st
 
-class PythonCompiler(Compiler):
-  def compile(self, src:str) -> bytes: return base64.b64decode(src)
-
 class PythonRenderer(Renderer):
   device = "PYTHON"
   code_for_op = python_alu
   def __init__(self):
-    self.compiler = PythonCompiler()
     match cast(str, EMULATE.value):
       case "METAL": self.device, self.tensor_cores = "METAL", tc.metal
       case "AMD": self.device, self.tensor_cores = "AMD", tc.amd_rdna3
@@ -239,6 +235,9 @@ class PythonRenderer(Renderer):
     lops = [(u.op, u.dtype, [uops.index(v) for v in u.src if u.op is not Ops.SPECIAL], u.arg) for u in uops]
     return base64.b64encode(pickle.dumps(lops)).decode()
 
+class PythonCompiler(Compiler):
+  def compile(self, src:str) -> bytes: return base64.b64decode(src)
+
 class PythonAllocator(Allocator['PythonDevice']):
   def _alloc(self, size, options): return memoryview(bytearray(size))
   def _copyin(self, dest, src:memoryview): dest[:] = src
@@ -246,4 +245,4 @@ class PythonAllocator(Allocator['PythonDevice']):
 
 class PythonDevice(Compiled):
   def __init__(self, device:str):
-    super().__init__(device, PythonAllocator(self), CompilerSet([CompilerPair(PythonRenderer, None)]), PythonProgram)
+    super().__init__(device, PythonAllocator(self), CompilerSet([CompilerPair(PythonRenderer, PythonCompiler)]), PythonProgram)

tinygrad/runtime/ops_webgpu.py

@@ -1,5 +1,5 @@
 import functools, struct
-from tinygrad.device import Compiled, Allocator, BufferSpec, CompilerSet, CompilerPair
+from tinygrad.device import  Compiled, Allocator, Compiler, BufferSpec, CompilerSet, CompilerPair
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.helpers import round_up, suppress_finalizing
 from tinygrad.runtime.autogen import webgpu
@@ -215,7 +215,7 @@ class WebGpuDevice(Compiled):
     device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback,
       webgpu.WGPURequestDeviceStatus, 1, 2, adapter_res, dev_desc)
 
-    super().__init__(device, WebGpuAllocator(device_res), CompilerSet([CompilerPair(WGSLRenderer, None)]),
+    super().__init__(device, WebGpuAllocator(device_res), CompilerSet([CompilerPair(WGSLRenderer, Compiler)]),
       functools.partial(WebGPUProgram, (device_res, webgpu.WGPUFeatureName_TimestampQuery in supported)))
 
   def synchronize(self):

tinygrad/uop/__init__.py

@@ -28,8 +28,8 @@ class Ops(FastEnum):
   NOOP = auto(); REWRITE_ERROR = auto()
 
   # renderer
-  # LINEAR is a list of UOps, SOURCE has a str arg that's human readable, BINARY has bytes arg that's compiled
-  PROGRAM = auto(); LINEAR = auto(); SOURCE = auto(); BINARY = auto()
+  # 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 (SINK, DEVICE, 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),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE), UPat(Ops.BINARY))), lambda: True),
-  # codegen: standalone LINEAR/SOURCE/BINARY
+  # 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)):