usb amd: clean up copyin, surface async bulk errors in bulk_wait

- stream 240KB chunks to bridge SRAM with interleaved 512B sentinels
- GPU SDMA ring poll-gates on sentinels, so one doorbell per copyin
- mid-flight engine arm latching + device-side fence gate (0xF4)
- two 256KB SRAM bounce buffers, per-buffer drain counters
- 1MB SDMA ring for usb, host signals moved off the NVMe SQ region

.github/actions/setup-tinygrad/action.yml

@@ -42,7 +42,11 @@ inputs:
     required: false
     default: 'false'
   qemu:
-    description: "Install qemu"
+    description: "Install qemu?"
+    required: false
+    default: 'false'
+  ninja:
+    description: "Install ninja?"
     required: false
     default: 'false'
 runs:
@@ -130,7 +134,7 @@ runs:
 
     # ******************* apt *******************
     - name: Setup apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       shell: bash
       run: |
         sudo mkdir -p /var/cache/apt/archives
@@ -158,7 +162,7 @@ runs:
         echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list
 
     - name: Compute Package List + Hash
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       id: apt-pkgs
       shell: bash
       run: |
@@ -183,25 +187,29 @@ runs:
         if [[ "${{ inputs.qemu }}" == "true" ]]; then
           pkgs+=" qemu-user-static"
         fi
+        # **** ninja ****
+        if [[ "${{ inputs.ninja }}" == "true" ]]; then
+          pkgs+=" ninja-build"
+        fi
 
         echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
         echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
 
     - name: Cache apt (PR)
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name == 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request'
       uses: actions/cache/restore@v5
       with:
         path: /var/cache/apt/archives/
         key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
     - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name != 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request'
       uses: actions/cache@v5
       with:
         path: /var/cache/apt/archives/
         key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
 
     - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       shell: bash
       run: |
         sudo apt -qq update || true

.github/workflows/autogen.yml

@@ -35,7 +35,7 @@ jobs:
         key: 'autogen'
         amd: 'true'
         llvm: 'true'
-        pydeps: 'pyyaml mako'
+        deps: 'autogen'
     - name: Install autogen support packages
       run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
     - name: Regenerate autogen files

.github/workflows/benchmark.yml

@@ -541,7 +541,7 @@ jobs:
     - name: openpilot run_pickle big_driving_supercombo
       run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
     - name: Test copy speeds
-      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})

.github/workflows/docs.yml

@@ -8,7 +8,7 @@ permissions:
   contents: write
 jobs:
   deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v6
       - name: Configure Git Credentials

.github/workflows/python-publish.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
     - uses: actions/checkout@v6
     - name: Set up Python

.github/workflows/szdiff.yml

@@ -10,7 +10,7 @@ concurrency:
 jobs:
   checkbranch:
     name: Check PR Branch status
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     outputs:
       branchstat: ${{ steps.brstat.outputs.stat}}
     steps:
@@ -44,7 +44,7 @@ jobs:
     permissions:
       contents: read
       pull-requests: write
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'false'
     steps:
@@ -87,7 +87,7 @@ jobs:
     name: Core Library Line Difference
     permissions:
       pull-requests: write
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'true'
     steps:

.github/workflows/test.yml

@@ -31,8 +31,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
-        deps: docs
-        pydeps: "capstone torch"
+        deps: "docs testing_minimal"
     - name: Build wheel and show size
       run: |
         uv build --wheel
@@ -73,10 +72,7 @@ jobs:
         deps: testing_unit
         pydeps: "pillow torchvision expecttest"
         llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
+        ninja: 'true'
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
     - name: Test one op in torch tests
@@ -98,12 +94,8 @@ jobs:
       with:
         key: torch-backend-pillow-torchvision-et-pt
         deps: testing_unit
-        pydeps: "pillow torchvision expecttest"
         llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
+        ninja: 'true'
     - name: Test beautiful_mnist in torch with TINY_BACKEND
       run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
 

examples/mlperf/models/gpt_oss.py

@@ -183,10 +183,12 @@ class GPTOSS:
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
     xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)
 
+    fa_saves = []
     if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention
-      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
+      attn, _, l_vec = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
       attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
+      fa_saves = [xq, xk, xv, l_vec]
     elif sliding:
       attn = self._sliding_attention(xq, xk, xv, sinks)
     else:
@@ -200,7 +202,7 @@ class GPTOSS:
       attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
 
     out = matmul_mx(attn, wo, wo_scale) + wo_bias
-    return out, [x_normed, rrms, attn]
+    return out, [x_normed, rrms, attn] + fa_saves
 
   def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
                    w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
@@ -221,6 +223,7 @@ class GPTOSS:
       z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
           + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
       out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
+      return out, [x_normed, rrms, xg, h, y, z, r.weights, r.dest_row, r.off]
     else:
       thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
       weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)

pyproject.toml

@@ -111,6 +111,10 @@ docs = [
   "numpy",
 ]
 mesa = ["tinymesa==25.2.7.2"]
+autogen = [
+  "pyyaml",
+  "mako",
+]
 
 
 [tool.mutmut]

test/backend/test_const_folding.py

@@ -1,7 +1,7 @@
 import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops, UOp, GroupOp
 from tinygrad.codegen.decomp.op import threefry2x32
 import numpy as np
 from test.helpers import not_support_multi_device
@@ -17,7 +17,7 @@ def _check_ast_count(desired_count:int, t:Tensor):
 class TestMovedConstFolding(unittest.TestCase):
   def test_contiguous_deviceless_const(self):
     t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
-    self.assertIs(t.uop.op, Ops.CONST)
+    self.assertIs(t.uop, UOp.const(2.0, dtypes.float))
     self.assertIsNone(t.uop.device)
 
   def test_add_shrunk_zero(self):
@@ -169,8 +169,8 @@ class TestMultiConstFolding(unittest.TestCase):
 class TestThreefryConstFolding(unittest.TestCase):
   def test_threefry(self):
     # THREEFRY(const,const) folds to a const once decomposed
-    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
-    self.assertIs(x.simplify().op, Ops.CONST)
+    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64)).simplify()
+    self.assertEqual([u.op for u in x.toposort() if u.op in GroupOp.ALU], [])
 
 class TestTautologicalCompare(unittest.TestCase):
   # without const folding, these would have triggered -Wtautological-compare in clang

test/backend/test_custom_kernel.py

@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import assert_kernel_count
+from test.helpers import assert_kernel_count, KernelCountException
 
 # **** kernels ****
 
@@ -474,7 +474,7 @@ class TestCustomKernelInput(unittest.TestCase):
     y.realize()
     kernel_count = GlobalCounters.kernel_count
     self.assertEqual(y.tolist(), x.add(1).tolist())
-    self.assertLessEqual(kernel_count, max_kernels)
+    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
     # same test with @function, input is PARAM
     from tinygrad import function
     x0 = Tensor.arange(32).clone("CPU").realize()
@@ -487,7 +487,7 @@ class TestCustomKernelInput(unittest.TestCase):
     y = run(x0).realize()
     kernel_count = GlobalCounters.kernel_count
     self.assertEqual(y.tolist(), mop_fxn(x0).add(1).tolist())
-    self.assertLessEqual(kernel_count, max_kernels)
+    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
 
   def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
   def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)

test/backend/test_isel.py

@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext
 
 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype)
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int), dtype=y.dtype)
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -46,10 +46,10 @@ class TestIselX86(unittest.TestCase):
   # complex address is [base + index*scale + displacement]
   def test_complex_address(self):
     a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load()
+    load = UOp.param(0, dtypes.int32, (16,)).index(a + UOp.cconst(1, dtypes.int32)).load()
     n = self.isel_rewrite(load)
     # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
-    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4)
+    self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)
 
 if __name__ == "__main__":
   unittest.main()

test/backend/test_linearizer.py

@@ -16,8 +16,6 @@ from test.helpers import replace_opts, check_schedule
 from test.backend.test_softmax_fusion import single_kernel_softmax
 MOCKGPU = DEV.interface.startswith("MOCK")
 
-from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
-
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering")
 class TestLinearizer(unittest.TestCase):
   def test_arg_dedup(self):
@@ -248,7 +246,6 @@ class TestLinearizer(unittest.TestCase):
     uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
     begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
     end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
-    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
     for u in uops:
       if u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.REG:
         if uops.index(u) < begin_range:
@@ -261,7 +258,6 @@ class TestLinearizer(unittest.TestCase):
         assert end_range < uops.index(u)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
-  @unittest.skipIf(Device[Device.DEFAULT].renderer.casted_consts, "reads a literal, which is casted here. TODO: flip this")
   def test_default_global_reversed(self):
     # shrink so that the dims do not collapse
     t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
@@ -269,9 +265,9 @@ class TestLinearizer(unittest.TestCase):
     uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
     idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
     idxs = sorted(idxs, key=lambda uop: uop.arg)
-    assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0]
-    assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg
-    assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg
+    assert (idxs[0].arg, idxs[0].src[0].src[0].val) == ('gidx0', 6), idxs[0]
+    assert (idxs[1].arg, idxs[1].src[0].src[0].val) == ('gidx1', 5), idxs[1].arg
+    assert (idxs[2].arg, idxs[2].src[0].src[0].val) == ('gidx2', 4), idxs[2].arg
 
   def test_sum_collapse(self):
     t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()

test/backend/test_multitensor.py

@@ -6,7 +6,7 @@ from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import run_linear, compile_linear, pm_beam, pm_compile
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count
+from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -395,7 +395,7 @@ class TestMultiBufferView(unittest.TestCase):
     linear, var_vals = b_multi.linear_with_vars()
     if all(not d.startswith(("WEBGPU", "CL")) for d in b_multi.device):
       compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
-      self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
+      if len(compiled) != 0: raise KernelCountException(0, len(compiled))
     run_linear(linear, var_vals)
     np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())
 

test/backend/test_ops.py

@@ -2164,6 +2164,10 @@ class TestOps(unittest.TestCase):
   def test_roll(self):
     helper_test_op([(2, 4)], lambda x: x.roll(1))
     helper_test_op([(2, 4)], lambda x: x.roll((1,)))
+    helper_test_op([(0,)], lambda x: x.roll(1, 0))
+    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 0))
+    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 1))
+    helper_test_op([(2, 0, 3)], lambda x: x.roll(1))
     self.helper_test_exception([(2, 4)], lambda x: x.roll((1, 2)), expected=RuntimeError)
     helper_test_op([(2, 4)], lambda x: x.roll(1, 0))
     helper_test_op([(2, 4)], lambda x: x.roll(-1, 0))

test/backend/test_pickle.py

@@ -3,6 +3,7 @@ import numpy as np
 from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
 from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function
+from test.helpers import KernelCountException
 
 class TestPickle(unittest.TestCase):
   def test_pickle_code_object(self):
@@ -41,7 +42,7 @@ class TestPickle(unittest.TestCase):
     t2:Tensor = pickle.loads(st)
     np.testing.assert_equal(t_values, t2.numpy())
     # expect at most one COPY kernel
-    self.assertLessEqual(GlobalCounters.kernel_count, 1)
+    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)
 
   def test_pickle_realized_tensor_alt(self):
     print("** init")

test/external/external_test_usb_asm24.py

@@ -6,7 +6,7 @@ import numpy as np
 class TestDevCopySpeeds(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    cls.sz = getenv("SIZE", 2e6)
+    cls.sz = getenv("SIZE", 2000000)
     cls.dev = Device["AMD"]
     if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
 

test/null/test_const_folding.py

@@ -1,39 +1,10 @@
 import unittest, itertools, math
-from tinygrad import Tensor, dtypes, Context
+from tinygrad import dtypes, Context
 from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
 
-def _check_ast_count(desired_count:int, t:Tensor):
-  # NOTE: this has side effect because everything can be scheduled only once
-  linear = t.schedule_linear()
-  asts = [s for s in linear.src if s.src[0].op is Ops.SINK]
-  len(asts)
-  # NOT SUPPORTED ANYMORE
-  #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
-
-class TestUnaryOpsConstFolding(unittest.TestCase):
-  def test_all_consts_ops(self):
-    _check_ast_count(0, Tensor.ones(4).exp())
-    _check_ast_count(0, Tensor.ones(4).sqrt())
-    _check_ast_count(0, Tensor.ones(4) + Tensor.ones(4))
-    _check_ast_count(0, Tensor.ones(4) / Tensor.ones(4))
-
-  def test_cast(self):
-    _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
-    _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
-
-  def test_neg_folding(self):
-    _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
-    _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
-    _check_ast_count(0, Tensor([1, 2, 3]).neg().neg())
-
-  def test_neg_realized_no_fold(self):
-    x = Tensor.randn(32, 32)
-    x = x.clip(0, 1).realize()
-    _check_ast_count(1, x.neg())
-
 class TestWeakConstFolding(unittest.TestCase):
   def test_weakint_math(self):
     out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
@@ -51,84 +22,18 @@ class TestWeakConstFolding(unittest.TestCase):
   def test_invalid_poison(self):
     self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
 
-  def test_single_rounding_log10_backward(self):
-    # log10 backward folds log10(2)/log(2) = 1/log(10) in one rounding, not the double-rounded 1/float32(log(10))
-    x = Tensor([1.0, 2.0, 3.0])
-    ast = next(s.src[0] for s in x.log10().sum().gradient(x)[0].schedule_linear().src if s.src[0].op is Ops.SINK)
-    const = next(u.arg for u in full_rewrite(ast).toposort() if u.op is Ops.CONST and u.dtype is dtypes.float32)
-    # correctly rounded: within half a float32 ulp of the exact value (folding at float32 lands 0.66 ulp off)
-    self.assertLess(abs(const - 1/math.log(10)), 2**-26)
-
-class TestBinaryOpsConstFolding(unittest.TestCase):
-  def test_add_literal_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
-  def test_add_tensor_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(4))
-  def test_literal_zero_add(self):
-    _check_ast_count(0, 0 + Tensor([1.0, 2, 3, 4]))
-  def test_tensor_zero_add(self):
-    _check_ast_count(0, Tensor.zeros(4) + Tensor([1.0, 2, 3, 4]))
-
-  def test_sub_literal_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - 0)
-  def test_sub_tensor_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - Tensor.zeros(4))
-
-  def test_mul_literal_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 0)
-  def test_mul_tensor_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.zeros(4))
-  def test_literal_zero_mul(self):
-    _check_ast_count(0, 0 * Tensor([1.0, 2, 3, 4]) * 0)
-  def test_tensor_zero_mul(self):
-    _check_ast_count(0, Tensor.zeros(4) * Tensor([1.0, 2, 3, 4]))
-
-  def test_mul_literal_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 1)
-  def test_mul_tensor_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(4))
-  def test_literal_one_mul(self):
-    _check_ast_count(0, 1 * Tensor([1.0, 2, 3, 4]))
-  def test_tensor_one_mul(self):
-    _check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))
-
-  def test_bool_tensor_mul_bool(self):
-    _check_ast_count(0, Tensor([True, False]) * True)
-    _check_ast_count(0, Tensor([True, False]) * False)
-  def test_bool_mul_bool_tensor(self):
-    _check_ast_count(0, True * Tensor([True, False]))
-    _check_ast_count(0, False * Tensor([True, False]))
-
-  def test_div_literal_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / 1)
-  def test_div_tensor_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))
-
-  def test_floordiv_literal_one(self):
-    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
-  def test_floordiv_tensor_one(self):
-    _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))
-
-  def test_pow_literal_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 0)
-  def test_pow_tensor_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.zeros(4))
-
-  def test_pow_literal_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 1)
-  def test_pow_tensor_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.ones(4))
-  def test_literal_one_pow(self):
-    _check_ast_count(0, 1 ** Tensor([1.0, 2, 3, 4]))
-  def test_tensor_one_pow(self):
-    _check_ast_count(0, Tensor.ones(4) ** Tensor([1.0, 2, 3, 4]))
-
 class TestBitcastConstFolding(unittest.TestCase):
+  def test_out_of_range_source_value(self):
+    for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000),
+                                      (70000, dtypes.int16, dtypes.uint16, 4464),
+                                      (-5, dtypes.uint32, dtypes.int32, -5)):
+      self.assertEqual(UOp.const(val, src_dt).bitcast(dst_dt).simplify().val, bits)
+
   def test_scalar_bitcast(self):
     def t(cases: dict[DType, ConstType]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
-          r = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0]
+          r = UOp.const(from_v, from_dt).bitcast(to_dt).simplify()
           self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
           self.assertEqual(r.dtype, to_dt, msg)
           np.testing.assert_equal(r.val, to_v, msg)
@@ -152,24 +57,9 @@ class TestBitcastConstFolding(unittest.TestCase):
 
   def test_vec_bitcast(self):
     with Context(SPEC=0):
-      srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src
-    self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
-    self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75))
-
-# folds advance indexing into basic indexing
-class TestIndexingConstFolding(unittest.TestCase):
-  def test_scalar_index(self):
-    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
-    _check_ast_count(1, t[:,:,Tensor(1),:])
-    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
-    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
-
-  def test_const_tensor_index(self):
-    # TODO: these can be 0, implement const tensor folded indexing
-    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
-    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
-    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
-    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
+      result = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink())
+      expected = full_rewrite(UOp.const((2**32-1, 2**31, 75), dtypes.uint32).sink())
+    self.assertEqual(result.src, expected.src)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_elf.py

@@ -1,6 +1,7 @@
 import unittest, subprocess, platform
 from tinygrad.runtime.support.compiler_cpu import ClangCompiler
 from tinygrad.runtime.support.elf import elf_loader
+from tinygrad.runtime.support.c import DLL
 
 class TestElfLoader(unittest.TestCase):
   def test_load_clang_jit_strtab(self):
@@ -23,7 +24,7 @@ class TestElfLoader(unittest.TestCase):
       }
     '''
     with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
-      ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src)
+      elf_loader(ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src))
   def test_link(self):
     src = '''
       float powf(float, float); // from libm
@@ -32,7 +33,7 @@ class TestElfLoader(unittest.TestCase):
     args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
     obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode())
     with self.assertRaisesRegex(RuntimeError, 'powf'): elf_loader(obj)
-    elf_loader(obj, link_libs=['m'])
+    elf_loader(obj, link_libs=[DLL('m', 'm')])
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_graph_rewrite.py

@@ -1,8 +1,7 @@
 import unittest, math
 from tinygrad import dtypes
-from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import all_same, Context
-from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat
+from tinygrad.uop.ops import GroupOp, UOp, Ops, PatternMatcher, TrackedPatternMatcher, UPat
 from test.helpers import full_rewrite
 from hypothesis import given, strategies as strat
 
@@ -11,125 +10,14 @@ from hypothesis import given, strategies as strat
 def apply_rewrite(expr):
   return full_rewrite(expr.sink()).src[0]
 
-@Context(SPEC=0)
-def apply_rewrite_values(expr):
-  srcs = full_rewrite(expr.sink()).src
-  if len(srcs) == 1:
-    if srcs[0].op is Ops.CONST: return (srcs[0].val,)
-    if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src)
-  return tuple(s.val for s in srcs)
-
-def evaluate_uop(uop, variables):
-  if uop.op == Ops.CONST:
-    return uop.val
-  elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU:
-    return variables[uop.expr]
-  elif uop.op in GroupOp.ALU:
-    src_values = [evaluate_uop(src, variables) for src in uop.src]
-    return exec_alu(uop.op, uop.dtype, src_values)
-  else:
-    raise NotImplementedError(f"Unsupported UOp {uop.op}")
-
-class TestArithmeticSimplifications(unittest.TestCase):
-  def test_full_graph_rewrite_division_by_zero(self):
-    optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0))
-    self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val))
-
-  def test_full_graph_rewrite_redundant_operations(self):
-    optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0))
-    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 10.0)
-
-  def test_full_graph_rewrite_large_graph(self):
-    prev_uop = UOp.const(0)
-    for i in range(1, 101):
-      prev_uop += UOp.const(i)
-    optimized_uop = apply_rewrite(prev_uop)
-    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, sum(range(1, 101)))
-
-  def test_full_graph_rewrite_division_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0))
-    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 42.0)
-
-  def test_full_graph_rewrite_modulo_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
-    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 0)
-
-
-class TestFoldingAndReduction(unittest.TestCase):
-  @unittest.skip("reduce is removed now")
-  def test_full_graph_rewrite_constant_reduction_folding(self):
-    const1 = UOp.const(5)
-    const2 = UOp.const(10)
-    const3 = UOp.const(20)
-    optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
-    expected_sum = 5 + 10 + 20
-    self.assertEqual(optimized_sink.val, expected_sum)
-
-  @unittest.skip("reduce is removed now")
-  def test_full_graph_rewrite_reduction_with_unused_range(self):
-    const1 = UOp.const(15)
-    const2 = UOp.const(25)
-    rng = UOp.range(10, idx=0)
-    optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
-    expected_sum = 10 * (15 + 25)
-    self.assertEqual(optimized_sink.val, expected_sum)
-
-  @unittest.skip("currently failing")
-  def test_full_graph_rewrite_range_reduction(self):
-    simple_range = UOp.range(5, idx=0)
-    optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
-    expected_sum = sum(range(5))
-    self.assertEqual(optimized_sink.val, expected_sum)
-
-  @unittest.skip("currently failing")
-  def test_full_graph_rewrite_simple_reduction_folding(self):
-    simple_range = UOp.range(4, idx=0)
-    add_uop = simple_range + UOp.const(1)
-    optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
-    expected_sum = sum(i + 1 for i in range(4))
-    self.assertEqual(optimized_sink.val, expected_sum)
-
-  @unittest.skip("currently failing")
-  def test_full_graph_rewrite_nested_loop_collapse(self):
-    outer_range = UOp.range(8, 0)
-    inner_range = UOp.range(4, 1)
-    expr = (outer_range * 10) + inner_range
-    optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
-    self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
-    self.assertEqual(optimized_reduce_uop.val, sum((i * 10) + j for i in range(8) for j in range(4)))
-
+def const_value(uop:UOp):
+  if uop.op is Ops.CAST: uop = uop.src[0]
+  assert uop.op is Ops.CONST
+  return uop.val
 
 class TestModuloAndDivisionFolding(unittest.TestCase):
-  def test_full_graph_rewrite_modulo_folding_with_define_var(self):
-    # index dtype because div-mod rules only work on index
-    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
-    optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
-    self.assertEqual(optimized_mod_uop.op, Ops.CONST)
-    self.assertEqual(optimized_mod_uop.val, 2)
-
-  def test_full_graph_rewrite_division_folding_with_define_var(self):
-    # index dtype because div-mod rules only work on index
-    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
-    optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
-    self.assertEqual(optimized_div_uop.op, Ops.MUL)
-    self.assertEqual(optimized_div_uop.src[1].val, 2)
-
-  def test_full_graph_rewrite_complex_mod_div_folding(self):
-    # index dtype because div-mod rules only work on index
-    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
-    optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
-    self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertEqual(optimized_div_uop.val, 1)
-
   def test_graph_rewrite_div_folding_bug(self):
-    lhs = UOp(Ops.ADD, src=(
-      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4),
-      UOp.const((0, 256, 512, 768))))
+    lhs = UOp.stack(*(UOp.special(32, 'lidx0'),)*4) + UOp.const((0, 256, 512, 768))
     rhs = UOp.const((2,)*4)
     unopt = lhs<rhs
     opt = apply_rewrite(unopt)
@@ -137,74 +25,31 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
     print(opt)
     if opt.op is Ops.STACK: self.assertFalse(all_same(opt.src))
 
-  def test_full_graph_rewrite_modulo_large_divisor(self):
-    # index dtype because div-mod rules only work on index
-    x_var_uop = UOp.variable('x', 1, 5)
-    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))
-
-  def test_full_graph_rewrite_division_with_remainder(self):
-    x_var_uop = UOp.variable('x', 7, 9, param=True)
-    optimized_sink = apply_rewrite(x_var_uop // 2)
-    for x_value in range(7, 10):
-      self.assertEqual(x_value // 2, evaluate_uop(optimized_sink, {'x': x_value}))
-
-  def test_full_graph_rewrite_complex_mod_div_expression(self):
-    x_var_uop = UOp.variable('x', 1, 10, param=True)
-    optimized_sink = apply_rewrite(((x_var_uop * 5) % 3) // 2)
-    for x_value in range(1, 11):
-      original_result = ((x_value * 5) % 3) // 2
-      optimized_result = evaluate_uop(optimized_sink, {'x': x_value})
-      self.assertEqual(original_result, optimized_result)
-
-
 class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
   def test_full_graph_rewrite_transcendental_edge_cases(self):
     optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal()))
     optimized_log2_neg, optimized_recip_zero = optimized_sink.src
-    self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}")
-    self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0,
-                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}")
-
-  @unittest.skip("broken")
-  def test_full_graph_rewrite_modulo_negative_dividend(self):
-    x_var_uop = UOp.variable('x', -5, -1)
-    optimized_sink = full_rewrite((x_var_uop % 3).sink())
-    for x_value in range(-5, 0):
-      self.assertEqual(x_value % 3, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
-
-  @unittest.skip("broken")
-  def test_full_graph_rewrite_division_negative_divisor(self):
-    x_var_uop = UOp.variable('x', 1, 5)
-    optimized_sink = full_rewrite((x_var_uop // -2).sink())
-    for x_value in range(1, 6):
-      self.assertEqual(x_value // -2, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
+    log2_neg, recip_zero = const_value(optimized_log2_neg), const_value(optimized_recip_zero)
+    self.assertTrue(math.isnan(log2_neg), f"Expected NaN for log2(-1.0), got {log2_neg}")
+    self.assertTrue(math.isinf(recip_zero) and recip_zero > 0, f"Expected +inf for reciprocal(0.0), got {recip_zero}")
 
 class TestGEPAndVectorizeRewrite(unittest.TestCase):
   def test_gep_single_element_extraction(self):
     # GEP on a vector dtype to extract a single element
     base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0)
+    self.assertIs(apply_rewrite(base_vector.index(2)), apply_rewrite(base_vector.src[2]))
 
   def test_gep_tuple_extraction(self):
     # GEP on a vector dtype to extract multiple elements as a vector
     base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
-
-  def test_gep_on_const_stack(self):
-    # GEP on a const STACK to extract a single element
-    const_stack = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(const_stack.index(2)).val, 3.0)
-
-  def test_gep_tuple_on_const_stack(self):
-    # GEP on a const STACK using a tuple to extract multiple elements
-    const_stack = UOp.const((7.0, 8.0, 9.0, 10.0))
-    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
+    self.assertIs(apply_rewrite(UOp.stack(*[base_vector.index(i) for i in (2, 3)])),
+                  apply_rewrite(UOp.stack(base_vector.src[2], base_vector.src[3])))
 
   def test_vectorize_multiple_elements(self):
     # Vectorizing multiple elements using GEP
     base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
-    vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4)))
-    self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])
+    vectorized_uop = UOp.stack(*(base_vector.index(i) for i in range(4)))
+    self.assertIs(apply_rewrite(vectorized_uop), apply_rewrite(base_vector))
 
 
 import inspect
@@ -256,16 +101,6 @@ class TestSubstitute(unittest.TestCase):
     ret = substitute(ret, {a.sin():b})
     self.assertIs(ret, b.sin())
 
-  # broken due to infinite recursion
-  # NOTE: VIZ hangs and doesn't recover if you click this one
-  @unittest.skip("recursion error no longer raised")
-  def test_assert_inf_recurse(self):
-    a = UOp.variable('a', 0, 10)
-    n1 = a.sin()
-    ret = n1
-    with self.assertRaises(RecursionError):
-      ret = substitute(ret, {n1:n1.sqrt()})
-
   def test_sin_to_sqrt(self):
     a = UOp.variable('a', 0, 10, dtype=dtypes.float)
     n1 = a.sin()

test/null/test_mnist_dataset.py

@@ -1,6 +1,7 @@
 import unittest
 from tinygrad.helpers import GlobalCounters
 from tinygrad.nn.datasets import mnist
+from test.helpers import KernelCountException
 
 class TestDataset(unittest.TestCase):
   def test_dataset_is_realized(self):
@@ -8,7 +9,7 @@ class TestDataset(unittest.TestCase):
     X_train[0].contiguous().realize()
     GlobalCounters.reset()
     X_train[0].contiguous().realize()
-    self.assertLessEqual(GlobalCounters.kernel_count, 1)  # 0 if SLICE (zero-copy), 1 otherwise
+    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)  # 0 if SLICE (zero-copy), 1 otherwise
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_simplify_valid_idx.py

@@ -15,16 +15,12 @@ def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move
 def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")
 
 def get_gated_load_uop(valid:UOp, idx:UOp):
-  return UOp(Ops.LOAD, src=(
-    UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)),
-  ))
+  return UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load()
 
 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
-  return UOp(Ops.LOAD, src=(
-    UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)),
-  ))
+  return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load()
 
-def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr)
+def Special(expr, nmax): return UOp.special(nmax, expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax, param=True)
 def Range(n, nmax): return UOp.range(nmax, n)
 
@@ -512,7 +508,7 @@ class TestDropTrueGate(unittest.TestCase):
     buf = UOp.param(0, dtypes.int, (1,))
     idx = UOp.const(0)
     true_gate = UOp.const(True)
-    index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
+    index_with_gate = buf.index(idx.valid(true_gate))
     # apply the optimization
     result = graph_rewrite(index_with_gate, sym+indexing_simplify)
     # the True valid should be dropped (INDEX should only have 2 sources)
@@ -524,13 +520,17 @@ class TestRangeShrink(unittest.TestCase):
       result = full_rewrite(sink)
     return [u for u in result.toposort() if u.op is Ops.RANGE]
 
+  def assert_range_end(self, ranges:list[UOp], end:int):
+    self.assertEqual(len(ranges), 1)
+    with Context(NOOPT=1, SPEC=0): expected = full_rewrite(UOp.const(end, dtypes.int).sink()).src[0]
+    self.assertIs(ranges[0].src[0], expected)
+
   def test_range_shrink_single_guard(self):
     # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
     r = Range(0, 204)
     load = get_gated_load_uop(r < UOp.const(4), r)
     ranges = self.get_ranges(load.sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assert_range_end(ranges, 4)
 
   def test_range_shrink_picks_max_guard(self):
     # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
@@ -538,25 +538,22 @@ class TestRangeShrink(unittest.TestCase):
     load1 = get_gated_load_uop(r < UOp.const(4), r)
     load2 = get_gated_load_uop(r < UOp.const(8), r)
     ranges = self.get_ranges(UOp.sink(load1, load2))
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 8)
+    self.assert_range_end(ranges, 8)
 
   def test_range_no_shrink_guard_ge_max(self):
     # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
     r = Range(0, 204)
     load = get_gated_load_uop(r < UOp.const(300), r)
     ranges = self.get_ranges(load.sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assert_range_end(ranges, 204)
 
   def test_range_no_shrink_when_unguarded_elsewhere(self):
     # one load guards r < 4, but another load uses r without a gate -> no shrink
     r = Range(0, 204)
     load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
+    load2 = UOp.param(1, dtypes.float, (204,)).index(r).load()
     ranges = self.get_ranges(UOp.sink(load1, load2))
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assert_range_end(ranges, 204)
 
   def test_range_no_shrink_when_used_in_reduce(self):
     # range used in both a gated load AND directly in the reduce expression -> no shrink
@@ -564,8 +561,7 @@ class TestRangeShrink(unittest.TestCase):
     gated_load = get_gated_load_uop(r < UOp.const(4), r)
     red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
     ranges = self.get_ranges(red.sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assert_range_end(ranges, 204)
 
   def test_range_shrink_to_single_iteration(self):
     # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
@@ -580,8 +576,7 @@ class TestRangeShrink(unittest.TestCase):
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(1.0), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assert_range_end(ranges, 4)
 
   def test_range_shrink_store_where_invalid_flipped(self):
     # above, but flipped
@@ -589,8 +584,7 @@ class TestRangeShrink(unittest.TestCase):
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(1.0), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assert_range_end(ranges, 4)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_tensor.py

@@ -171,6 +171,11 @@ class TestTensorConstLike(unittest.TestCase):
     t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
     with self.assertRaises(RuntimeError): t.full_like(5, device="NULL")
 
+class TestTensorShape(unittest.TestCase):
+  def test_float_shape_raises(self):
+    for dim in (2.0, 2.5):
+      with self.subTest(dim=dim), self.assertRaisesRegex(RuntimeError, "shape must be int"): Tensor.ones(dim)
+
 class TestTensorDevice(unittest.TestCase):
   def test_create_from_single_device_tuple(self):
     (Tensor([1.0], device=(Device.DEFAULT,)) + Tensor([2.0])).realize()

test/null/test_uop_graph.py

@@ -1,10 +1,9 @@
 import unittest, pytest
 from tinygrad import dtypes, Variable, Device
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo
 from tinygrad.uop.symbolic import sym
-from test.helpers import to_uops_list
+from test.helpers import full_rewrite, to_uops_list
 from tinygrad.codegen import full_rewrite_to_sink
 
 simple_pm = PatternMatcher([
@@ -14,43 +13,27 @@ simple_pm = PatternMatcher([
   ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.val+c2.val)),
 ])
 
-def const_values(u:UOp):
-  if u.op is Ops.CONST: return (u.val,)
-  if u.op is Ops.STACK: return tuple(x.val for x in u.src)
-  raise AssertionError(f"expected const-like UOp, got {u.op}")
-
 class TestGraphRewriteConst(unittest.TestCase):
   def test_gep_const(self):
     v1 = UOp.const((0,1,2), dtypes.int)
     v2 = v1.index(1)
     ret = graph_rewrite(v2, sym)
-    self.assertEqual(ret.dtype, dtypes.int)
-    self.assertEqual(ret.val, 1)
+    self.assertIs(ret, UOp.const(1, dtypes.int))
 
   def test_add_const(self):
     v1 = UOp.const((0,1,2))
     v2 = UOp.const((5,6,7))
-    ret = graph_rewrite(v1+v2, sym)
-    self.assertEqual(ret.op, Ops.STACK)
-    self.assertEqual(const_values(ret), (5,7,9))
-
-  def test_add_const_lose_v(self):
-    v1 = UOp.const((0,1,2))
-    v2 = UOp.const((2,1,0))
-    ret = graph_rewrite(v1+v2, sym)
-    self.assertEqual(ret.op, Ops.STACK)
-    self.assertEqual(const_values(ret), (2,2,2))
+    self.assertIs(graph_rewrite(v1+v2, sym), UOp.const((5,7,9)))
 
 def xfail_broken_const_wraparound(fn):
   fn = pytest.mark.xfail(reason="const folding does not properly implement modular arithmetic")(fn)
   return unittest.expectedFailure(fn)
 class TestModularWraparound(unittest.TestCase):
   def _test(self, uop:UOp, expected:int):
-    results = to_uops_list([uop])
-    self.assertEqual(len(results), 2)  # +1 for SINK
-    self.assertEqual(results[0].op, Ops.CONST)
-    self.assertEqual(results[0].dtype, uop.dtype)
-    self.assertEqual(results[0].val, expected)
+    result = uop.simplify()
+    self.assertEqual(result.op, Ops.CONST)
+    self.assertEqual(result.dtype, uop.dtype)
+    self.assertEqual(result.val, expected)
 
   @xfail_broken_const_wraparound
   def test_cast(self):
@@ -191,63 +174,25 @@ class TestGraphRewrite(unittest.TestCase):
       self.assertEqual(len([x for x in sink.toposort() if x.op is Ops.CONST]), 1)
 
 class TestUOpGraph(unittest.TestCase):
-  def test_add_constant_fold(self):
-    c1 = UOp.const(1.0, dtypes.float)
-    c2 = UOp.const(2.0, dtypes.float)
-    out = c1+c2
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 3.0)
-
   def test_where_same_fold(self):
     v = UOp.variable('tmp', 0, 1)
     c0 = UOp.const(0)
     vc = v != c0
     c1 = UOp.const(1.0, dtypes.float)
     out = vc.where(c1, c1)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 1.0)
+    self.assertIs(out.simplify(), c1)
 
   def test_where_const_fold(self):
     bf = UOp.const(False)
     c1 = UOp.const(1.0, dtypes.float)
     c2 = UOp.const(2.0, dtypes.float)
     out = bf.where(c1, c2)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 2.0)
+    self.assertIs(out.simplify(), c2)
 
   def test_const_cast(self):
     bf = UOp.const(False)
     out = bf.cast(dtypes.int)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0)
-
-  def test_const_bitcast(self):
-    bf = UOp.const(1.0, dtypes.float)
-    out = bf.bitcast(dtypes.uint32)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0x3F800000)
-
-  @unittest.expectedFailure
-  def test_const_shape_change_bitcast(self):
-    bf = UOp.const(0x3F).cast(dtypes.uint8)
-    out = bf.bitcast(dtypes.half)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
+    self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
 
   def test_devectorize_derives_lane_dtype(self):
     from tinygrad.codegen import do_devectorize
@@ -257,66 +202,11 @@ class TestUOpGraph(unittest.TestCase):
     invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
     self.assertIs(invalid_lane_mul.dtype, dtypes.bool)
 
-  @unittest.skip("this test isn't valid uops")
-  def test_noop_vectorize_fold(self):
-    d0 = UOp.param(0, dtypes.float, (1,))
-    idx = UOp.const(0)
-    ld = d0.load(idx, dtype=dtypes.float)
-    vec = UOp(Ops.STACK, dtypes.float, (ld,))
-    x = vec.index(0)
-    alu = UOp(Ops.SQRT, src=(x, ))
-    out = UOp(Ops.STORE, src=(d0, idx, alu))
-    uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.STACK]), 0)
-
-  @unittest.skip("this test isn't valid uops")
-  def test_gep_vec_fold(self):
-    d0 = UOp.param(0, dtypes.float, (1,))
-    d1 = UOp.param(1, dtypes.float, (1,))
-    d2 = UOp.param(2, dtypes.float, (1,))
-    idx = UOp.const(0)
-    def _test_vec(geps, count=4):
-      vec = UOp(Ops.STACK, dtypes.float, geps)
-      out = d0.index(idx).store(vec)
-      uops = to_uops_list([out])
-      if DEBUG >= 4:
-        from tinygrad import Device
-        print(Device[Device.DEFAULT].renderer.render(uops))
-      return uops[-2].src[-1]  # -2 to skip SINK
-
-    # possible
-    val = d1.index(idx).load(dtype=dtypes.float)
-    xyzw = tuple(val.index(i) for i in range(4))
-    self.assertIs(_test_vec(xyzw).op, Ops.LOAD)
-
-    # unaligned
-    val = d1.index(idx).load(dtype=dtypes.float)
-    wzyx = tuple(val.index(i) for i in reversed(range(4)))
-    self.assertIs(_test_vec(wzyx).op, Ops.STACK)
-
-    # different_size
-    val = d1.index(idx).load(dtype=dtypes.float)
-    xy = tuple(val.index(i) for i in range(2))
-    self.assertIs(_test_vec(xy+xy).op, Ops.STACK)
-    val = d1.index(idx).load(dtype=dtypes.float)
-    xy = tuple(val.index(i) for i in range(2))
-    self.assertIs(_test_vec(xy, count=2).op, Ops.STACK)
-
-    # different vals
-    val1 = d1.index(idx).load(dtype=dtypes.float)
-    val2 = d2.index(idx).load(dtype=dtypes.float)
-    xy1 = tuple(val1.index(i) for i in range(2))
-    xy2 = tuple(val2.index(i) for i in range(2))
-    self.assertIs(_test_vec(xy1+xy2).op, Ops.STACK)
-
   def test_gep_vec_const_fold(self):
     for vec_size in [2, 4, 8]:
       consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
-      vec = UOp(Ops.STACK, src=tuple(consts))
-      with Context(SPEC=0):
-        uops = to_uops_list([vec.index(i) for i in range(vec_size)])
-        for uop, const in zip(uops, consts):
-          self.assertEqual(uop, const)
+      vec = UOp.stack(*consts)
+      for i, const in enumerate(consts): self.assertIs(vec.index(i), const)
 
   def test_cast_alu_fold(self):
     d0 = UOp.param(0, dtypes.bool, (1,))
@@ -326,7 +216,7 @@ class TestUOpGraph(unittest.TestCase):
     alu = (ld<1).cast(dtypes.bool)
     out = d0.index(idx).store(alu)
     uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 0)
 
   def test_double_cast_fold(self):
     d0 = UOp.param(0, dtypes.float, (1,))
@@ -336,7 +226,7 @@ class TestUOpGraph(unittest.TestCase):
     alu = ld.cast(dtypes.float).cast(dtypes.float)
     out = d0.index(idx).store(alu)
     uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 1)
 
   def test_depth_2_const_fold(self):
     v = UOp.variable("tmp", 0, 1, dtypes.int, param=True)
@@ -344,12 +234,7 @@ class TestUOpGraph(unittest.TestCase):
     c4 = UOp.const(4, dtypes.int)
     vc = v+c2
     out = vc+c4
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 5)  # +1 for SINK, +1 for the PARAM shape STACK
-    out = uops[-2]  # -2 to skip SINK
-    self.assertEqual(out.op, Ops.ADD)
-    self.assertEqual(out.src[1].op, Ops.CONST)
-    self.assertEqual(out.src[1].val, 6)
+    self.assertIs(out.simplify(), (v+UOp.const(6, dtypes.int)).simplify())
 
   def test_bitcast_to_same_dtype_fold(self):
     for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
@@ -360,9 +245,8 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_sub_with_cast_folds(self):
     a = Variable("a", 0, 5)
-    uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
-    assert uops[0] == UOp.const(0, dtypes.int)
-    assert uops[-1].op == Ops.SINK
+    out = a.cast(dtypes.int)+(-a).cast(dtypes.int)
+    self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
 
   def test_where_on_gated_load_fold(self):
     ridx0 = UOp.range(100, 0)
@@ -371,9 +255,10 @@ class TestUOpGraph(unittest.TestCase):
     w = (ridx0<50).where(ld, 5)
     out = UOp.param(1, dtypes.long, (100,))
     uops = to_uops_list([out.index(ridx0).store(w)])
+    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: self.assertIs(u.src[1], expected)
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
@@ -381,9 +266,10 @@ class TestUOpGraph(unittest.TestCase):
     ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
     w = (ridx0<50).where(5, ld)
     uops = to_uops_list([w])
+    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: assert u.src[1].val==5
+      if u.op is Ops.LOAD: self.assertIs(u.src[1], expected)
 
   def test_where_on_gated_load_with_cast(self):
     ridx0 = UOp.range(100, 0)
@@ -393,9 +279,10 @@ class TestUOpGraph(unittest.TestCase):
     w = (ridx0<50).where(ld, 5.0)
     out = UOp.param(1, dtypes.float, (100,))
     uops = to_uops_list([out.index(ridx0).store(w)])
+    expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: self.assertIs(u.src[1], expected)
 
   def test_where_on_casted_gated_load_extra_cond(self):
     ridx0 = UOp.range(100, 0)
@@ -425,9 +312,10 @@ class TestUOpGraph(unittest.TestCase):
     val = (ridx0<50).where(5, ld)
     st = idx.store(val).end(ridx0)
     uops = to_uops_list([st])
+    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.STORE: assert u.src[1].val==5
+      if u.op is Ops.STORE: self.assertIs(u.src[1], expected)
 
   def test_load_idx_becomes_int(self):
     # mnist indexing with split reduceop
@@ -501,13 +389,6 @@ class TestUOpGraph(unittest.TestCase):
     # only the second store happens
     self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1)
 
-  @unittest.skip("this is a uop type error")
-  def test_asserts_bad_gate(self):
-    glbl0 = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(0)
-    bad_gate = UOp.const(1)
-    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))])
-
   def test_after_end(self):
     r = UOp.range(10, 0)
 
@@ -575,7 +456,7 @@ class TestConstBufferize(unittest.TestCase):
     from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
     c = UOp.const(42.0)
     r1 = UOp.range(3, 0)
-    bufferize_with_range = UOp(Ops.STAGE, src=(c, r1), arg=BufferizeOpts(device="CPU"))
+    bufferize_with_range = c.bufferize(r1, arg=BufferizeOpts(device="CPU"))
     self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range
 
     result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
@@ -590,7 +471,7 @@ class TestConstBufferize(unittest.TestCase):
     c = UOp.const(3.14)
     r1 = UOp.range(3, 0)
     r2 = UOp.range(4, 1)
-    bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU"))
+    bufferize_with_ranges = c.bufferize(r1, r2, arg=BufferizeOpts(device="CPU"))
     self.assertEqual(len(bufferize_with_ranges.src), 3)  # const + 2 ranges
 
     result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')

test/null/test_uops.py

@@ -41,11 +41,6 @@ class TestDTypeFromUOp(unittest.TestCase):
     # an explicit (strong) const dtype is legal until the field is removed
     self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)
 
-  def test_weak_dtype_rejected_by_program_spec(self):
-    for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
-      with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program)
-      type_verify(UOp.const(value, concrete).sink(), spec_program)
-
   def test_invalid_stated_dtype(self):
     # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
     # and the spec is what rejects a non-bool Invalid
@@ -134,7 +129,7 @@ class TestConstFloatEq(unittest.TestCase):
     self.assertFalse(Invalid != HoldsInvalid())
 
   def test_matchers_agree_on_nan(self):
-    n = UOp.const(math.nan, dtypes.float32)
+    n = UOp.const(math.nan)
     for compiled in (False, True):
       pm = PatternMatcher([(UPat(Ops.CONST, arg=math.nan), lambda: True)], compiled=compiled)
       self.assertTrue(pm.rewrite(n), f"{compiled=}")
@@ -348,10 +343,9 @@ class TestFastIdiv(unittest.TestCase):
   def test_fast_idiv_remove_powers_of_two(self):
     ridx = UOp.range(2**20, 0)
     uops = to_uops_list([ridx//(7*64)], ren=Device[Device.DEFAULT].renderer)
-    ops = [x.op for x in uops]
     # this requires shifting out the powers of two before doing fast_idiv
     # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
-    self.assertNotIn(Ops.CAST, ops)
+    self.assertNotIn(dtypes.long, [x.dtype for x in uops])
 
   @unittest.expectedFailure
   def test_fast_idiv_overflow(self):

test/unit/test_allreduce.py

@@ -2,6 +2,7 @@ import unittest
 from tinygrad import Tensor, UOp, dtypes
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops
+from test.helpers import KernelCountException
 
 class TestRingAllReduce(unittest.TestCase):
   def test_schedule_ring(self):
@@ -13,7 +14,7 @@ class TestRingAllReduce(unittest.TestCase):
       copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
       pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]
       # N*(N-1) scatter reduce, and N*(N-1) allgather
-      self.assertEqual(len(pairs), N*(N-1)*2)
+      if len(pairs) != N*(N-1)*2: raise KernelCountException(N*(N-1)*2, len(pairs))
       # copy topology forms a ring
       self.assertEqual(len(set(pairs)), N)
 
@@ -25,8 +26,8 @@ class TestRingAllReduce(unittest.TestCase):
       linear = t.sum(0).mul(2.0).contiguous().linear_with_vars()[0]
       copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
       sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
-      self.assertEqual(len(copies), 24)
-      self.assertEqual(len(sinks), 26)
+      if len(copies) != 24: raise KernelCountException(24, len(copies))
+      if len(sinks) != 26: raise KernelCountException(26, len(sinks))
 
   @Context(RING=0, ALL2ALL=0)
   def test_schedule_naive(self):
@@ -39,8 +40,8 @@ class TestRingAllReduce(unittest.TestCase):
     sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
     pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]
 
-    self.assertEqual(len(pairs), N*(N-1))
-    self.assertEqual(len(sinks), 2)
+    if len(pairs) != N*(N-1): raise KernelCountException(N*(N-1), len(pairs))
+    if len(sinks) != 2: raise KernelCountException(2, len(sinks))
     self.assertTrue(all(dst != src for dst, src in pairs))
 
   def test_symbolic_shape(self):

test/unit/test_function.py

@@ -4,7 +4,7 @@ from tinygrad.function import function
 from tinygrad import Tensor, GlobalCounters, Device
 from tinygrad.dtype import Invalid
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
-from test.helpers import assert_kernel_count
+from test.helpers import assert_kernel_count, KernelCountException
 
 class TestFunction(unittest.TestCase):
   def test_simple(self):
@@ -516,7 +516,7 @@ class TestFunctionTuple(unittest.TestCase):
     Tensor.realize(a)
     c = f(a)
 
-    self.assertEqual(count_kernels(c), 1)
+    if count_kernels(c) != 1: raise KernelCountException(1, count_kernels(c))
 
     c.sum().backward()
     Tensor.realize(a.grad)

tinygrad/codegen/__init__.py

@@ -5,7 +5,7 @@ from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey,
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType
 from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak, pm_cast_weak
 from tinygrad.uop.render import pyrender
-from tinygrad.uop.spec import type_verify, spec_tensor, spec_program, spec_program_casted_consts
+from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
 from tinygrad.dtype import dtypes, AddrSpace
@@ -282,7 +282,7 @@ pm_implicit_barriers = PatternMatcher([
 ])
 
 pm_casted_consts = PatternMatcher([
-  (UPat(Ops.CONST, dtypes.all, name="c"), lambda c: UOp(Ops.CAST, c.dtype, src=(UOp.const(c.val),), arg=c.dtype)),
+  (UPat(Ops.CONST, dtypes.all, name="c"), lambda c: UOp.cconst(c.val, c.dtype)),
 ])
 
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
@@ -387,11 +387,12 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1])
   sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
 
-  # TODO: delete once migration are done
-  if ren.casted_consts: sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
+  # spell every literal as a casted const CAST(dt, CONST(value))
+  # TODO: remove once consts are always weak
+  sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
 
   if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
-  if SPEC: type_verify(sink, spec_program_casted_consts if ren.casted_consts else spec_program)
+  if SPEC: type_verify(sink, spec_program)
 
   # return the rewritten sink
   return sink

tinygrad/codegen/late/regalloc.py

@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
 from tinygrad.renderer.isa import ISARenderer, Register, greg
 from tinygrad.dtype import dtypes
 
-PSEUDO_OPS = {Ops.CONST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
+PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
 
 class LinearScanRegallocContext:
   # returns the uop that defines the virtual register
@@ -52,7 +52,7 @@ class LinearScanRegallocContext:
         # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes
         sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize)
         offset = self.stack_size + (sz - self.stack_size % sz) % sz
-        self.spills[v] = UOp.const(offset, dtypes.int32)
+        self.spills[v] = UOp.cconst(offset, dtypes.int32)
         self.stack_size = offset + sz
       r = alloc(cons if cons is not None else v.cons, i)
       self.insert_before.setdefault(i, []).append((v, r))
@@ -84,7 +84,7 @@ class LinearScanRegallocContext:
 
       # allocate stack array
       if u.op is Ops.BUFFER:
-        self.locals[u] = UOp.const(self.stack_size, dtypes.int32)
+        self.locals[u] = UOp.cconst(self.stack_size, dtypes.int32)
         self.stack_size += u.max_numel() * u.dtype.itemsize
 
       # loop prologue, avoid loading inside the loop
@@ -125,7 +125,7 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
   # alloc/dealloc stack
   if ctx.stack_size > 0:
     sp = ctx.ren.stack_pointer()
-    offset = UOp.const(ctx.stack_size, sp.dtype)
+    offset = UOp.cconst(ctx.stack_size, sp.dtype)
     if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before
     elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))]
 

tinygrad/engine/realize.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
-import time, random, itertools, math, contextlib, weakref, array
+import random, itertools, math, weakref, array, decimal
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.dtype import dtypes
@@ -17,6 +17,7 @@ def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call
 def get_call_var_uops(call:UOp, prg:UOp) -> list[UOp]:
   bound = {s.src[0].expr: s.src[1].src[1] for s in call.src[1:] if s.is_bound_var}
   return [bound.get(v.expr, v) for v in prg.arg.vars]
+
 def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
   ast = call.src[0]
   if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins)
@@ -24,6 +25,12 @@ def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
   return (), ()
 
+def get_call_kernels(call:UOp) -> list[tuple[str, UOp]]:
+  if (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return [(d, k) for devs, k, _ in call.arg.aux.kernels for d in devs]
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return [(to_tuple(ast.device)[0], call)]
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "validate": return []
+  return [(d, call) for d in to_tuple(call.src[1].device)]
+
 def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|None=None) -> str:
   def _uop_sz_to_str(uop:UOp) -> str: return size_to_str(sym_infer(prod(uop.shape) * uop.dtype.itemsize, var_vals or {}))
   def _dev_str(buf:Buffer|UOp) -> str: return ', '.join(d[:7] for d in to_tuple(buf.device))
@@ -39,49 +46,52 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
 # **************** Stat ****************
 
 def estimate_uop(call:UOp) -> Estimates:
-  ast = call.src[0]
-  if ast.op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
+  if (ast:=call.src[0]).op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
   if ast.op is Ops.COPY or (ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec"):
-    nbytes = prod(call.src[1].shape) * call.src[1].dtype.itemsize
-    return Estimates(lds=nbytes, mem=nbytes)
+    return Estimates(lds=(nbytes:=prod(call.src[1].shape) * call.src[1].dtype.itemsize), mem=nbytes)
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return get_graph_runtime(ast).estimates
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return call.arg.aux.estimates
   return Estimates()
 
 first_run_cache:set[bytes] = set()
-@contextlib.contextmanager
-def track_stats(ctx:ExecContext, call:UOp, device:str, bufs:list[Buffer], var_vals:dict[str, int]):
-  if PROFILE:
-    outputs, inputs = get_call_outs_ins(call)
-    cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"var_vals": var_vals,
-      "bufs": [b.trace_num for b in bufs], "name": get_call_name(call, bufs, var_vals), "outputs": outputs, "inputs": inputs}))
-  et: list[float|None] = [None]
-  if DEBUG >= 2: st = time.perf_counter()
-  yield et
-  if not ctx.update_stats: return
+def track_stats(ctx:ExecContext, call:UOp, st:decimal.Decimal, ets:list[float|None]):
+  if ctx.update_stats:
+    is_hcq = (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq"
+    estimates, n = estimate_uop(call), 1 if is_hcq else len(get_call_kernels(call))
+    GlobalCounters.kernel_count += len(call.arg.aux.kernels) if is_hcq else n
+    GlobalCounters.global_ops += n*sym_infer(estimates.ops, ctx.var_vals)
+    GlobalCounters.global_mem += n*sym_infer(estimates.mem, ctx.var_vals)
+    GlobalCounters.time_sum_s += sum(et for et in ets if et is not None)
+  if DEBUG < 2 and not PROFILE: return
 
-  if DEBUG >= 2 and et[0] is None:
-    Device[device].synchronize()
-    et[0] = time.perf_counter() - st
+  kernels = get_call_kernels(call) # everything below is the per kernel display: exec events for the profiler and DEBUG=2 lines
+  args = resolve_params(call, ctx.input_uops) if kernels and kernels[0][1] is call else []
+  lanes = list(unwrap_multi(call, [args[g] for g in call.src[0].arg.globals] if call.src[0].op is Ops.PROGRAM else args)) if args else []
+  for i, (device, kcall) in enumerate(kernels):
+    et, bufs = ets[i] if i < len(ets) else None, lanes[i][0] if i < len(lanes) else []
+    if PROFILE: # backdate the event to the start of the call, the viz matches a device range with the exec event before it
+      outputs, inputs = get_call_outs_ins(kcall)
+      cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"var_vals": ctx.var_vals,
+        "bufs": [b.trace_num for b in bufs], "name": get_call_name(kcall, bufs, ctx.var_vals), "outputs": outputs, "inputs": inputs}, ts=st))
+    if DEBUG < 2 or not ctx.update_stats: continue
+    if et is None:
+      Device[device].synchronize()
+      et, st = float(perf_counter_us() - st)*1e-6, perf_counter_us()
+      GlobalCounters.time_sum_s += et
 
-  estimates = estimate_uop(call)
-  GlobalCounters.kernel_count += 1
-  GlobalCounters.global_ops += (op_est:=sym_infer(estimates.ops, var_vals))
-  GlobalCounters.global_mem += (mem_est:=sym_infer(estimates.mem, var_vals))
-  if et[0] is not None: GlobalCounters.time_sum_s += et[0]
-  if DEBUG >= 2:
-    display_name = get_call_name(call, bufs, var_vals)
-    lds_est = sym_infer(estimates.lds, var_vals)
-    header_color = 'magenta' if ctx.jit else ('green' if call.src[0].key not in first_run_cache else None)
-    ptm = colored(time_to_str(et[0], w=9), "yellow" if et[0] > 0.01 else None) if et[0] is not None else ""
-    flops, membw, ldsbw = op_est/(et[0] or 1e-20), mem_est/(et[0] or 1e-20), lds_est/(et[0] or 1e-20)
+    estimates = estimate_uop(kcall)
+    display_name = get_call_name(kcall, bufs, ctx.var_vals)
+    op_est, mem_est, lds_est = (sym_infer(x, ctx.var_vals) for x in (estimates.ops, estimates.mem, estimates.lds))
+    header_color = 'magenta' if ctx.jit else ('green' if kcall.src[0].key not in first_run_cache else None)
+    ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
+    flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
     flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
     mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
       colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
     print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
       f" {display_name+' '*(46-ansilen(display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
-      ("" if et[0] is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})"))
-    first_run_cache.add(call.src[0].key)
+      ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})"))
+    first_run_cache.add(kcall.src[0].key)
 
 local_size_cache: dict[bytes, tuple[int, ...]] = {}
 def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:
@@ -154,33 +164,31 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
                    for x in call.src[0].toposort())
     for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_dnum else {}
 
-def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
     dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
-    with track_stats(ctx, call, dest.device, [dest, src], ctx.var_vals):
-      if hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]:
-        dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev)
-      elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
-           and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
-        dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
-      elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
-      else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
-  return None
+    if hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]:
+      dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev)
+    elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
+         and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
+      dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
+    elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
+    else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
+  return []
 
-def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  et = None
+def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
+  ets:list[float|None] = []
   resolved = resolve_params(call, ctx.input_uops)
   for device, (bufs, device_vars) in zip(to_tuple(call.src[1].device), unwrap_multi(call, [resolved[i] for i in ast.arg.globals])):
     var_vals = {**ctx.var_vals, **device_vars}
     prg_bufs = [b.ensure_allocated() for b in bufs]
     rt = get_runtime(device, ast, cache=ctx.cache)
     global_size, local_size = ast.arg.launch_dims(var_vals)
-    with track_stats(ctx, call, device, prg_bufs, var_vals) as tm:
-      et = tm[0] = rt(*[b.get_buf(device) for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals),
-                       wait=ctx.wait, timeout=ctx.timeout)
-  return et
+    ets.append(rt(*[b.get_buf(device) for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals),
+                  wait=ctx.wait, timeout=ctx.timeout))
+  return ets
 
-def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   import numpy as np
   for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
     bufs, dev_bufs = bufs[:len(bufs)//2], bufs[len(bufs)//2:]
@@ -189,43 +197,36 @@ def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
     global_size, local_size = prg.arg.launch_dims(var_vals)
     cpu_rt(*[bufs[i].ensure_allocated()._buf for i in prg.arg.globals], global_size=global_size, local_size=local_size, vals=prg.arg.vals(var_vals))
     for i in prg.arg.outs: np.testing.assert_allclose(dev_bufs[i].ensure_allocated().numpy(), bufs[i].numpy(), rtol=1e-3, atol=1e-3)
-  return None
+  return []
 
-def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)]
   shape, pos_var = tuple(s.val for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
-  with track_stats(ctx, call, bufs[0].device, bufs, ctx.var_vals):
-    bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
-  return None
+  bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
+  return []
 
-def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  rt = get_graph_runtime(ast, ctx.input_uops)
-  with track_stats(ctx, call, rt.device, [], ctx.var_vals) as t: t[0] = rt(ctx.input_uops, ctx.var_vals, wait=ctx.wait)
-  return t[0]
+def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
+  return [get_graph_runtime(ast, ctx.input_uops)(ctx.input_uops, ctx.var_vals, wait=ctx.wait)]
 
-def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   dev = cast(Any, Device[(info:= call.arg.aux).device[0]])
   addrs = [(b.bufs[j] if isinstance(b:=_resolve(ctx.input_uops[k], ctx.input_uops).buffer, MultiBuffer) else b).get_buf(dev_name).va_addr
            for devs, idxs in info.input_idxs for j, dev_name in enumerate(devs) for k in idxs]
   dev.rt_buffer._buf.cpu_view().view(offset=(base:=dev.rt_allocator.alloc(len(addrs) * 8)), fmt='Q')[:len(addrs)] = array.array('Q', addrs)
 
-  tables = [UOp.from_buffer(dev.rt_buffer.view(len(idxs), dtypes.uint64, base + j*len(idxs)*8), HCQ_RUNTIME_DEV.value)
-            for devs, idxs in info.input_idxs for j in range(len(devs))]
-  if info.inputs is not None: call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
-  exec_kernel(replace(ctx, update_stats=DEBUG>=3, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer._buf.va_addr + base}), call, ast)
+  if info.inputs is not None:
+    tables = [UOp.from_buffer(dev.rt_buffer.view(len(idxs), dtypes.uint64, base + j*len(idxs)*8), HCQ_RUNTIME_DEV.value)
+              for devs, idxs in info.input_idxs for j in range(len(devs))]
+    call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
+  exec_kernel(replace(ctx, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer._buf.va_addr + base}), call, ast)
 
-  tms = []
-  for devices, stat_call, prof in info.kernels:
-    for device in devices:
-      tm = None
-      if prof:
-        (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, stat_call.arg.name, *prof)
-        if ctx.wait:
-          d.synchronize(timeout=ctx.timeout)
-          st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
-          tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
-      with track_stats(ctx, stat_call, device, [], ctx.var_vals) as et: et[0] = tm
-  return max(tms) if tms else None
+  def _prof_tm(device:str, stat_call:UOp, prof:tuple[int, ...]) -> float|None:
+    (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, stat_call.arg.name, *prof)
+    if not ctx.wait: return None
+    d.synchronize(timeout=ctx.timeout)
+    st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
+    return float(en-st)/d.timestamp_divider/1e6
+  return [_prof_tm(device, k, prof) for devices, k, prof in info.kernels if prof for device in devices] if PROFILE or ctx.wait else []
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -280,7 +281,7 @@ def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequenc
   inputs = list(input_uops)
   if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs))
   ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2)
-  for call in linear.src: pm_exec.rewrite(call, ctx)
+  for call in linear.src: track_stats(ctx, call, perf_counter_us(), pm_exec.rewrite(call, ctx))
 
 def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> float:
   if clear_l2:
@@ -290,4 +291,4 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None
       with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
   ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
   linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache)
-  return max(pm_exec.rewrite(c, ctx) or 0.0 for c in linear.src)
+  return max(et for c in linear.src for et in pm_exec.rewrite(c, ctx) or [0.0])

tinygrad/mixin/movement.py

@@ -550,6 +550,7 @@ class MovementMixin:
     if dims is None: return self.flatten().roll(shifts, 0).reshape(self.shape)
     dims, shifts = tuple(self._resolve_dim(d) for d in make_tuple(dims, 1)), make_tuple(shifts, 1)
     if len(dims) != len(shifts): raise RuntimeError(f"{len(dims)=} != {len(shifts)=}")
+    if 0 in self.shape: return self
     shrink_arg: list[tuple[sint, sint]|None] = [None] * self.ndim
     for d, s in zip(dims, shifts): shrink_arg[d] = (delta:=self.shape[d]-s%self.shape[d], delta+self.shape[d])
     return self.repeat(*tuple(2 if i in dims else 1 for i in range(self.ndim))).shrink(tuple(shrink_arg))

tinygrad/renderer/__init__.py

@@ -72,8 +72,6 @@ class Renderer:
   tensor_cores: list[TensorCore] = []
   extra_matcher: PatternMatcher|None = None
   code_for_op: dict[Ops, Callable] = {}
-  # migration: this renderer consumes every literal as a casted const CAST(dt, CONST(value))
-  casted_consts: bool = False
 
   compiler: Compiler = Compiler()
 

tinygrad/renderer/cstyle.py

@@ -20,19 +20,19 @@ base_rewrite = PatternMatcher([
   (UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"),
 
   # const
-  (UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: None if math.isfinite(v:=x.val) else \
+  (UPat.cvar("c").cast(dtypes.floats, name="x"), lambda ctx,x,c: None if math.isfinite(v:=c.val) else \
     f"({ctx.render_cast(x, ctx.nan if math.isnan(v) else ctx.infinity if v > 0 else f'-{ctx.infinity}')})"),
-  (UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.val}f"),
-  (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.val}l"),
-  (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}ul"),
-  (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}u"),
-  (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.val else "0"),
+  (UPat.cvar("c").cast(dtypes.float), lambda ctx,c: f"{c.val}f"),
+  (UPat.cvar("c").cast(dtypes.int64), lambda ctx,c: f"{c.val}l"),
+  (UPat.cvar("c").cast(dtypes.uint64, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}ul"),
+  (UPat.cvar("c").cast(dtypes.uint32, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}u"),
+  (UPat.cvar("c").cast(dtypes.bool), lambda ctx,c: "1" if c.val else "0"),
   # consts are rendered to larger type and casted
-  (UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}f')})"),
-  (UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}u')})"),
-  (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, str(x.val))})"),
+  (UPat.cvar("c").cast((*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, f'{c.val}f')})"),
+  (UPat.cvar("c").cast((dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, f'{c.val}u')})"),
+  (UPat.cvar("c").cast((dtypes.int8, dtypes.int16), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, str(c.val))})"),
   # default const render
-  (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.val)),
+  (UPat.cvar("c").cast(), lambda ctx,c: str(c.val)),
 
   # casting
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" \
@@ -47,7 +47,7 @@ base_rewrite = PatternMatcher([
 
   # SHRINK/INDEX
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)),
-  (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat.cvar()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)),
+  (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat.cvar().cast()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)),
   (UPat(Ops.STACK, name="x"),
    lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_type(x))}" + \
                  f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"),
@@ -161,8 +161,8 @@ class CStyleLanguage(Renderer):
   def render_index(self, x:UOp, buf:UOp, idx:UOp):
     if buf.addrspace == AddrSpace.ALU:
       # this is lane access in C
-      if idx.op is not Ops.CONST: return f"({self[buf]})[{self[idx]}]"
-      return self[buf]+(f"[{idx.val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.val]}")
+      if not (idx.op is Ops.CAST and idx.src[0].op is Ops.CONST): return f"({self[buf]})[{self[idx]}]"
+      return self[buf]+(f"[{idx.src[0].val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.src[0].val]}")
     return f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})"
 
   def render_buffer(self, x:UOp):
@@ -208,7 +208,7 @@ class CStyleLanguage(Renderer):
     c: defaultdict[str, int] = defaultdict(int)
     name = "test"
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP}: continue
+      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST}: continue
       if u.op == Ops.STACK and len(u.src) == 0: continue
       if u.op is Ops.AFTER:
         r[u] = r[u.src[0]]
@@ -226,7 +226,7 @@ class CStyleLanguage(Renderer):
       if u.op is Ops.SPECIAL: r[u] = u.arg
       elif u.op is Ops.RANGE: r[u] = f"{axis_letters[u.arg[-1]]}idx"+range_str(u)
       else:
-        prefix = {Ops.WMMA: "wmma", Ops.CONST: "const", Ops.BUFFER: "buf", Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.STACK: "cast",
+        prefix = {Ops.WMMA: "wmma", Ops.BUFFER: "buf", Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.STACK: "cast",
                   Ops.INDEX: "bidx", Ops.LOAD: "val"}.get(u.op, "alu")
         r[u] = f"{prefix}{c[prefix]}"
 
@@ -234,7 +234,8 @@ class CStyleLanguage(Renderer):
       assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}"
 
       if u.op in {Ops.ENDIF, Ops.END}: depth -= 1
-      if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
+      if (u.op is not Ops.CAST or u.max_numel() == 1) and ((u.op is Ops.CAST and u.src[0].op is Ops.CONST) or \
+        u.op in {Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
         (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \
         (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
         (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
@@ -318,8 +319,7 @@ class OpenCLRenderer(CStyleLanguage):
   string_rewrite = PatternMatcher([
     (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
     # bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort
-    (UPat(Ops.CONST, dtypes.bfloat16, name="x"),
-      lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.val)))[0] >> 16)}u"),
+    (UPat.cvar("c").cast(dtypes.bfloat16), lambda ctx,c: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(c.val)))[0] >> 16)}u"),
     # load/store image (OpenCL)
     (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"),
     (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))),
@@ -493,8 +493,8 @@ class HIPRenderer(CStyleLanguage):
         (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]},"
           f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[0][2] == 128 else None),
         (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"),
-        (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x:
-          f"f32_to_fp8({ctx.nan if math.isnan(v:=x.val) else ctx.infinity if v == math.inf else f'-{ctx.infinity}' if v == -math.inf else f'{v}f'},"
+        (UPat.cvar("c").cast(dtypes.fp8s, name="x"), lambda ctx,x,c:
+          f"f32_to_fp8({ctx.nan if math.isnan(v:=c.val) else ctx.infinity if v == math.inf else f'-{ctx.infinity}' if v == -math.inf else f'{v}f'},"
           f" {fp8_index(x.dtype)})"),
         (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",),
           lambda ctx,x: f"f32_to_fp8({ctx[x.src[0]]}, {fp8_index(x.dtype)})"),
@@ -536,7 +536,7 @@ class HIPRenderer(CStyleLanguage):
     prefix, ockl = [], []
     type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" }
     used_dtypes = uops_to_dtypes(uops)
-    if any(u.op is Ops.CONST and not math.isfinite(u.val) for u in uops):
+    if any(u.op is Ops.CAST and u.src[0].op is Ops.CONST and not math.isfinite(u.src[0].val) for u in uops):
       prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"]
     if any(u.op is Ops.SPECIAL for u in uops):
       prefix.append("typedef long unsigned int size_t;")
@@ -550,7 +550,7 @@ class HIPRenderer(CStyleLanguage):
     if any(dt in dtypes.fp8s for dt, _ in used_dtypes):
       prefix += ["typedef unsigned char hip_bf8;", "typedef unsigned char hip_fp8;"]
     if any((u.op is Ops.CAST and u.dtype in dtypes.fp8s and u.src[0].dtype == dtypes.float) or
-           (u.op is Ops.CONST and u.dtype in dtypes.fp8s) for u in uops):
+           (u.op is Ops.CAST and u.src[0].op is Ops.CONST and u.dtype in dtypes.fp8s) for u in uops):
       prefix.append("""static inline __attribute__((device)) unsigned char f32_to_fp8(float v, int is_bf8) {
   v = (((*(unsigned*)&v)&0x7F800000)!=0x7F800000)?__builtin_amdgcn_fmed3f(v,is_bf8?57344.0f:448.0f,is_bf8?-57344.0f:-448.0f) : v;
   return (unsigned char)(is_bf8?__builtin_amdgcn_cvt_pk_bf8_f32(v,v,0,false):__builtin_amdgcn_cvt_pk_fp8_f32(v,v,0,false));\n}""")

tinygrad/renderer/isa/x86.py

@@ -166,7 +166,7 @@ def scratch_buffer(elem_dt:DType, count:int, slot:int) -> UOp:
 
 def gated_load(ctx, addr:UOp, alt:UOp, gate:UOp, x:UOp):
   local = scratch_buffer(addr.src[0].dtype, x.max_numel(), next(ctx))
-  local_idx = local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64)
+  local_idx = local.index(UOp.cconst(0, dtypes.int32), dtype=dtypes.uint64)
   # the selected address is a 64bit value, the AFTER orders the load after the scratch store and carries the element dtype for the encoder
   sel = gate.where(addr.replace(dtype=dtypes.uint64), local_idx)
   ptr = UOp(Ops.AFTER, addr.dtype, (sel, (local_idx if x.max_numel() == 1 else local).store(alt)))
@@ -174,7 +174,7 @@ def gated_load(ctx, addr:UOp, alt:UOp, gate:UOp, x:UOp):
 
 def gated_store(addr:UOp, gate:UOp, val:UOp):
   local = scratch_buffer(addr.src[0].dtype, val.max_numel(), -1)
-  sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64))
+  sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.cconst(0, dtypes.int32), dtype=dtypes.uint64))
   return UOp(Ops.AFTER, addr.dtype, (sel,)).store(val)
 
 # legalize the new style graph for isel. NOTE: this runs after the spec is verified, some of these rewrites violate it
@@ -195,7 +195,7 @@ pre_isel_matcher = PatternMatcher([
   # if gate in scalar int cmove is not a comparison need to add one to set the flag
   # NOTE: the 0 is int so the bool gate zero-extends and compares as int (a byte compare renders different kernels)
   (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")),
-   lambda m,a,b: m.ne(UOp.const(0, dtypes.int)).where(a,b) if m.op not in GroupOp.Comparison else None),
+   lambda m,a,b: m.ne(UOp.cconst(0, dtypes.int)).where(a,b) if m.op not in GroupOp.Comparison else None),
 ])
 
 # ***** X86 registers *****
@@ -221,15 +221,15 @@ reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"},
 
 # ***** X86 instruction selection *****
 def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX else s
-def lane(x:UOp, i:int) -> int: return s.src[1].val if (s:=x.src[i]).op is Ops.INDEX else 0
+def lane(x:UOp, i:int) -> int: return s.src[1].src[0].val if (s:=x.src[i]).op is Ops.INDEX else 0
 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt]
 def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, arg=X86Ops.DEFINE, tag=None if reg is None else (reg,))
-def imm(dt:DType, v:int) -> UOp: return UOp.const(truncate[dt](v), dt).rtag()
+def imm(dt:DType, v:int) -> UOp: return UOp.cconst(truncate[dt](v), dt).rtag()
 def to_imm(c:UOp) -> UOp|None:
-  if c.op is not Ops.CONST: return None
-  if c.dtype is dtypes.int64: return imm(dtypes.int32, c.val) if not c.overflows(dtypes.int32) else None
-  if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.val) if not c.overflows(dtypes.uint32) else None
-  if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.val)
+  if not (c.op is Ops.CAST and (v:=c.src[0]).op is Ops.CONST): return None
+  if c.dtype is dtypes.int64: return imm(dtypes.int32, v.val) if not v.overflows(dtypes.int32) else None
+  if c.dtype is dtypes.uint64: return imm(dtypes.uint32, v.val) if not v.overflows(dtypes.uint32) else None
+  if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, v.val)
   return None
 def cmp(x:UOp) -> UOp:
   if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void)
@@ -289,8 +289,9 @@ def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]:
   # buffers are indexed by element, everything else (the stack pointer) by byte
   scale = base.dtype.itemsize if base.op in {Ops.PARAM, Ops.BUFFER, Ops.AFTER} else 1
   sz = imm(dtypes.uint8, base.dtype.itemsize)
-  if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].val * scale), sz)
-  if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.val * scale), sz)
+  if idx.op is Ops.ADD and (c:=idx.src[1]).op is Ops.CAST and c.src[0].op is Ops.CONST:
+    return (base, _cast(idx.src[0]), _disp(c.src[0].val * scale), sz)
+  if idx.op is Ops.CAST and idx.src[0].op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.src[0].val * scale), sz)
   return (base, _cast(idx), _disp(0), sz)
 
 def abi(ctx:IselContext, x:UOp) -> UOp|None:
@@ -353,7 +354,7 @@ isel_matcher = PatternMatcher([
   # cast of void is a noop
   (UPat.var("y").cast(name="x"), lambda y,x: y if y.dtype == dtypes.void else None),
   # range is lowered to acc, cmp, jmp after regalloc
-  (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.val),) + x.src[1:])),
+  (UPat(Ops.RANGE, src=(UPat.cvar("c").cast(),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(x.dtype, c.val),) + x.src[1:])),
   (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None),
   # really all a backedge END is is an IF with a tag referencing the RANGE start label
   (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"),
@@ -367,10 +368,10 @@ isel_matcher = PatternMatcher([
   # function abi constraints
   (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi),
   # constants that can't be immediates, move them to registers
-  (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.val),)) if not x.tag else None),
-  (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.val),)) if not x.tag else None),
-  (UPat.cvar("x", dtypes.floats), lambda x:
-   UOp.const(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, x.val))[0], dt).bitcast(x.dtype) if not x.tag else None),
+  (UPat.cvar("c").cast(dtypes.int64s, name="x"), lambda c,x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, c.val),)) if not x.tag else None),
+  (UPat.cvar("c").cast(dtypes.ints+(dtypes.bool,), name="x"), lambda c,x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, c.val),)) if not x.tag else None),
+  (UPat.cvar("c").cast(dtypes.floats, name="x"), lambda c,x:
+   UOp.cconst(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, c.val))[0], dt).bitcast(x.dtype) if not x.tag else None),
   # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists
   (UPat.var("m").where(UPat.var("a", dtypes.int8s+dtypes.int16s+dtypes.int32s+(dtypes.int64,)), UPat.var("b")), lambda m,a,b:
    a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.max_numel() > 1 else None),
@@ -380,7 +381,7 @@ isel_matcher = PatternMatcher([
    a.ins(X86Ops.VBLENDVPD, src=(b, a, m.replace(dtype=m.src[0].dtype)))),
   # in this case we have a mask producing comparison whose user expects a bool, so we convert to bool
   (UPat(GroupOp.Comparison, dtypes.bool, (UPat.var("y", (dtypes.float32, dtypes.float64)), UPat()), name="x"), lambda y,x:
-   UOp(Ops.AND, src=(x.replace(dtype=y.dtype).bitcast(dt:=to_int(y.dtype)), UOp.const(1, dt))).f(Ops.NOOP, dtype=dtypes.bool)),
+   UOp(Ops.AND, src=(x.replace(dtype=y.dtype).bitcast(dt:=to_int(y.dtype)), UOp.cconst(1, dt))).f(Ops.NOOP, dtype=dtypes.bool)),
   # conditional moves that use flags
   (UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.sints), UPat()), name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b:
    a.ins(X86Ops.CMOVL, src=(b, a, cmp(m)))),
@@ -420,15 +421,15 @@ isel_matcher = PatternMatcher([
   (UPat(Ops.STACK, dtypes.float32, name="x"), vinsertps),
   (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), vpins),
   # INDEX on a vector register value extracts a single element
-  (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c"), name="x"),
+  (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c").cast(), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int16s).index(UPat.cvar("c"), name="x"),
+  (UPat.var("y", dtypes.int16s).index(UPat.cvar("c").cast(), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int32s).index(UPat.cvar("c"), name="x"),
+  (UPat.var("y", dtypes.int32s).index(UPat.cvar("c").cast(), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int64s).index(UPat.cvar("c"), name="x"),
+  (UPat.var("y", dtypes.int64s).index(UPat.cvar("c").cast(), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.floats).index(UPat.cvar("c"), name="x"),
+  (UPat.var("y", dtypes.floats).index(UPat.cvar("c").cast(), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.val * x.dtype.itemsize))) if _is_vec_xmm(y) else None),
   # packed bitwise
   ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.max_numel() > 1 else None),
@@ -453,15 +454,19 @@ isel_matcher = PatternMatcher([
   # scalar int binary
   ((UPat(dtype=dtypes.ints).alu(Ops.CDIV, UPat())).named("x"), idiv),
   # scalar int binary with immediate
-  (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.val)))),
-  (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))),
-  (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))),
-  (UPat.var("a", dtypes.ints) + UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat.var("a", dtypes.ints) * UPat.cvar("c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.cvar("c"), lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar("c"))), lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints) << UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.val)))),
+  (UPat.var("a", dtypes.uints) >> UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))),
+  (UPat.var("a", dtypes.sints) >> UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))),
+  (UPat.var("a", dtypes.ints) + UPat.cvar().cast(name="c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints) * UPat.cvar().cast(name="c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar().cast(name="c"),
+   lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.cvar().cast(name="c"),
+   lambda a,c: a.ins(X86Ops.ORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.cvar().cast(name="c"),
+   lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar().cast(name="c"))),
+   lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   # scalar int binary with register
   ((UPat(dtype=dtypes.ints) << UPat()).named("x"), lambda x: shift(x, X86Ops.SHL)),
   ((UPat(dtype=dtypes.uints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SHR)),
@@ -572,7 +577,7 @@ def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
   if x.dtype is dtypes.void: return (label, [label])
   else:
     acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:])
-    cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0]))
+    cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CAST else X86Ops.CMP, src=(acc, x.src[0]))
     jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}")
     ctx.loop_label[acc] = loop_label
     return (acc, [acc, label, cmp, jump_out])
@@ -591,7 +596,7 @@ def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
 # final rewrite to match the isa spec
 post_regalloc_matcher = PatternMatcher([
   # rewrite FRAME_INDEX to IMM now that the stack size is known
-  (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=x.const_like(ctx.stack_size + x.tag), [nx])),
+  (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=UOp.cconst(ctx.stack_size + x.tag, x.dtype), [nx])),
   # expand the cmp here so we can preserve rng src edge to get label from ctx
   (UPat(Ops.INS, arg=X86Ops.LOOP_CMP, name="x"), lower_loop),
   # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
@@ -614,7 +619,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     rm = cast(Register, greg(rm_uop)).index
     idx = cast(Register, greg(idx_uop)).index if idx_uop is not None and greg(idx_uop) is not None else 4
     # for a memory operand the rm size is the element size from the address, otherwise it's the size of the value in the register
-    rm_sz = sz_uop.val if sz_uop is not None else rm_uop.dtype.itemsize
+    rm_sz = sz_uop.src[0].val if sz_uop is not None else rm_uop.dtype.itemsize
     reg_sz = reg_uop.dtype.itemsize if reg_uop is not None else 0
     sz = reg_sz or rm_sz
 
@@ -647,10 +652,10 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # 0b10 -- signals memory access with 32bit displacement
     # 0b11 -- signals no memory access
     if disp_uop is not None:
-      assert disp_uop.op is Ops.CONST, "displacement must be a constant"
+      assert disp_uop.op is Ops.CAST, "displacement must be a literal"
       assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int"
       # rbp/r13 always require a displacement
-      if disp_uop.val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
+      if disp_uop.src[0].val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
       else: mod = 0b00
     else: mod = 0b11
     # x 0b0 and idx 0b100 means rsp which means no index exists
@@ -664,10 +669,10 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # DISP byte
     if mod == 0b01 or mod == 0b10:
       assert disp_uop is not None
-      inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.val)
+      inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.src[0].val)
     # IMM byte
     if imm_uop is not None:
-      if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.val)
+      if imm_uop.op is Ops.CAST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.src[0].val)
       elif isinstance(greg(imm_uop), Register): inst += bytes([(greg(imm_uop).index & 0b1111) << 4 | 0b0000])
     return inst
 
@@ -677,13 +682,13 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
   if x.arg in X86GroupOp.WriteMem:
     if len(x.src) > 4: address, rest = x.src[:4], x.src[4:]
     else: address, rest = (x, None, None, None), x.src
-    imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,)
+    imm_uop = rest[:1] if rest and rest[0].op is Ops.CAST else (None,)
     return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *imm_uop))
 
   if x.arg in X86GroupOp.Rm1st:
     if len(x.src) > 3: address, rest = x.src[:4], x.src[4:]
     else: address, rest = (x.src[0], None, None, None), x.src[1:]
-    imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,)
+    imm_uop = rest[:1] if rest and rest[0].op is Ops.CAST else (None,)
     return _encode(x, *address, *(None, *imm_uop)) if reg is None else _encode(None, *address, *(x if sel else None, *imm_uop))
 
   if x.arg in X86GroupOp.Rm2nd:
@@ -701,7 +706,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
 encodings = {
   # moves
   X86Ops.MOVABS: lambda x:
-   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].val),
+   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].src[0].val),
   X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0),
   X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D),
   X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1),
@@ -724,8 +729,8 @@ encodings = {
   X86Ops.VCVTPS2PD: lambda x: encode(x, 0x5A, pp=0, sel=1), X86Ops.VCVTPD2PS: lambda x: encode(x, 0x5A, pp=1, sel=1),
   X86Ops.VCVTTPS2DQ: lambda x: encode(x, 0x5B, pp=2, sel=1), X86Ops.VCVTTPD2DQ: lambda x: encode(x, 0xE6, pp=1, sel=1),
   # the int src is the 2nd src (the rm field), if it was folded into a memory operand its width is the element size of the address
-  X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
-  X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
+  X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].src[0].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
+  X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].src[0].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
   X86Ops.VCVTTSS2SI: lambda x: encode(x, 0x2C, pp=2, sel=1, we=x.dtype.itemsize == 8),
   X86Ops.VCVTTSD2SI: lambda x: encode(x, 0x2C, pp=3, sel=1, we=x.dtype.itemsize == 8),
   # int division
@@ -840,10 +845,10 @@ class X86Renderer(ISARenderer):
     def _format_op(x:UOp) -> str: return f"    {(o[7:-1] if (o:=str(x.arg))[-1] in ('i', 'm') else o[7:]).lower():7s}"
     def _format_operands(x:UOp) -> str:
       def _format(src:tuple[UOp, ...]) -> list[str]:
-        return [str(s.val) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize, o) if \
+        return [str(s.src[0].val) if s.op is Ops.CAST else reg_strs[o].get(s.dtype.itemsize, o) if \
                 (o:=str(greg(s))) in reg_strs else o for s in src if greg(s) is not None]
       def _mem_adress(base:UOp, idx:UOp, disp:UOp, sz:UOp) -> list[str]:
-        return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.val}" if greg(idx) else "") + (f" + {disp.val}" if disp.val else "") + "]"]
+        return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.src[0].val}" if greg(idx) else "") + (f" + {d}" if (d:=disp.src[0].val) else "") + "]"]
 
       if len(x.src) > 4 and x.arg in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:])
       elif len(x.src) > 3 and x.arg in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:])

tinygrad/renderer/llvmir.py

@@ -81,8 +81,8 @@ base_rewrite = PatternMatcher([
   (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat((Ops.BUFFER, Ops.PARAM, Ops.AFTER)),), allow_any_len=True, name="x"), lambda ctx,x:
    f"  {ctx[x]} = getelementptr inbounds {ldt(x.dtype)}, {ldt(x.dtype, ptr=True)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"),
   # register index
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("idx")), name="x"), lambda ctx,buf,idx,x:
-   f"  {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.val}" if buf.addrspace == AddrSpace.ALU else None),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("c").cast()), name="x"), lambda ctx,buf,c,x:
+   f"  {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {c.val}" if buf.addrspace == AddrSpace.ALU else None),
 
   # load/store
   (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"),
@@ -165,7 +165,7 @@ class LLVMRenderer(Renderer):
     local_args: list[str] = []
     name = "test"
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP}: continue
+      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST}: continue
       if u.op is Ops.AFTER:
         r[u] = r[u.src[0]]
         continue
@@ -185,7 +185,7 @@ class LLVMRenderer(Renderer):
           kernel.append(f"  {r[u]} = addrspacecast [{size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{size} x {ldt(u.dtype)}]*")
         else:
           kernel.append(f"  {r[u]} = alloca [{size} x {ldt(u.dtype)}], align 16")
-      elif u.op is Ops.CONST: r[u] = lconst(u.val, u.dtype)
+      elif u.op is Ops.CAST and u.src[0].op is Ops.CONST: r[u] = lconst(u.src[0].val, u.dtype)
       elif u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype):
         r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast
       else:

tinygrad/renderer/nir.py

@@ -145,7 +145,7 @@ class NIRRenderer(Renderer):
   ])
 
   def_rewrite = PatternMatcher([
-    (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.val, x.dtype)),
+    (UPat.cvar("c").cast(name="x"), lambda ctx,x,c: nimm(ctx.b, c.val, x.dtype)),
     (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, x.dtype.itemsize if x.addrspace is AddrSpace.ALU else 8)),
     (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
     (UPat(Ops.STORE, src=(UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val"))),
@@ -186,16 +186,17 @@ class NIRRenderer(Renderer):
 
   def render(self, uops:list[UOp]):
     self.prerender(uops)
-    for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].val
+    for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]:
+      self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].src[0].val
     self.r: dict[UOp, Any] = {}
     self.param_idx = 0
     ranges: list[mesa.nir_def|None] = []
 
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.STACK and len(u.src) == 0): pass
+      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST} or (u.op is Ops.STACK and len(u.src) == 0): pass
       elif u.op in {Ops.INDEX, Ops.SHRINK}:
         # INDEX on a register value picks the element, memory INDEX is handled in the LOAD/STORE patterns
-        if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].val)
+        if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].src[0].val)
       elif u.op is Ops.AFTER:
         self.r[u] = self.r[u.src[0]]
       elif u.op == Ops.SINK:

tinygrad/renderer/ptx.py

@@ -79,8 +79,8 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i
   (a.itemsize < b.itemsize or dtypes.is_int(b) or b == dtypes.bool) else ''
 
 string_rewrite = PatternMatcher([
-  (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.val, x.dtype)}, 0;"),
-  (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.val, x.dtype)};"),
+  (UPat.cvar("c").cast(dtypes.bool, name="x"), lambda ctx, x, c: f"setp.ne.s16 {ctx.r[x]}, {render_val(c.val, x.dtype)}, 0;"),
+  (UPat.cvar("c").cast(name="x"), lambda ctx, x, c: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(c.val, x.dtype)};"),
   (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"),
   (UPat(Ops.PARAM, name="x"), lambda ctx, x:
    f"ld.param.{ctx.types[dtypes.ulong] if x.addrspace is AddrSpace.GLOBAL else ctx.mem_types[x.dtype]} {ctx.r[x]}, [data{x.arg.slot}+0];"),
@@ -186,7 +186,7 @@ class PTXRenderer(Renderer):
 
     name = "test"
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP}: continue
+      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST}: continue
       if u.op is Ops.AFTER:
         self.r[u] = self.r[u.src[0]]
         continue
@@ -201,9 +201,9 @@ class PTXRenderer(Renderer):
         continue
       if u.op in {Ops.INDEX, Ops.SHRINK, Ops.LOAD} and u.src[0].addrspace in (AddrSpace.REG, AddrSpace.ALU):
         # on REG, INDEX/SHRINK pick the register (must be CONST) and LOAD is a noop
-        if u.op is not Ops.LOAD and u.src[1].op is not Ops.CONST:
+        if u.op is not Ops.LOAD and not (u.src[1].op is Ops.CAST and u.src[1].src[0].op is Ops.CONST):
           raise RuntimeError(f"PTX does not support dynamic register indexing: {u}")
-        r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].val]
+        r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].src[0].val]
         continue
       if u.op is Ops.SPECIAL: r[u] = "%" + u.arg
       elif u.op is Ops.LOAD:
@@ -216,7 +216,7 @@ class PTXRenderer(Renderer):
                        [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.itemsize)]]
         r[u] = [ssa("wmma", dtype=self.types[u.dtype]) for _ in range(u.max_numel())]
       prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.END: ("pred", "pred"), Ops.RANGE: ("ridx", None),
-        Ops.CONST: ("const", None), Ops.BUFFER: ("local", "u64"), Ops.INDEX: ("bidx", "u64"), Ops.SHRINK: ("bidx", "u64"),
+        Ops.BUFFER: ("local", "u64"), Ops.INDEX: ("bidx", "u64"), Ops.SHRINK: ("bidx", "u64"),
         Ops.PARAM: ("dat", "u64" if u.addrspace is AddrSpace.GLOBAL else None), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None))
       if u.op is Ops.RANGE and u.dtype == dtypes.void: prefix = None  # loop headers don't have a register
       if prefix: r[u] = ssa(prefix, u, dtype)

tinygrad/renderer/wgsl.py

@@ -69,10 +69,10 @@ class WGSLRenderer(CStyleLanguage):
 
   string_rewrite = PatternMatcher([
     (UPat(Ops.NEG, dtypes.uints, src=(UPat.var('x'))), lambda ctx,x: f"(0-{ctx[x]})"),
-    (UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.val else "false"),
-    (UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"),
-     lambda x: f"bitcast<u32>({x.val})" if x.val < 0 else f"{x.val&0xFFFFFFFF}u"),
-    (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}"),
+    (UPat.cvar("c").cast(dtypes.bool), lambda c: "true" if c.val else "false"),
+    (UPat.cvar("c").cast((dtypes.uchar, dtypes.ushort, dtypes.uint32)),
+     lambda c: f"bitcast<u32>({c.val})" if c.val < 0 else f"{c.val&0xFFFFFFFF}u"),
+    (UPat.cvar("c").cast(dtypes.int32, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}"),
     (UPat(Ops.BUFFER, name="x"), lambda ctx,x:
      f"var{'<workgroup>' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"),
     (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),

tinygrad/runtime/ops_amd.py

@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
 from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices, hcq_profile
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, BufferSpec, TinyELF
-from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey
+from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey, from_mv
 from tinygrad.helpers import VIZ, ceildiv, unwrap, pluralize
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
@@ -649,6 +649,151 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
 
   def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
 
+  def _copyin(self, dest:HCQBuffer, src:memoryview):
+    if not self.dev.is_usb(): return super()._copyin(dest, src)
+    from tinygrad.runtime.support.usb import alloc_cbuffer
+    # Pipelined USB copyin, self-synchronized between the 0xF2 engine and the GPU.
+    #
+    # The host streams 240KB payload chunks to the bridge SRAM over EP 0x02, alternating between the two 256KB
+    # bounce buffers, with a 512B sentinel sector after every 120KB part. The prebuilt SDMA ring holds, per part,
+    # a POLL packet that spins until the engine has written that part's sentinel, then copies it to VRAM: the GPU
+    # drains each part as it lands and the whole copyin needs a single doorbell. The arm for chunk c+1 is issued
+    # while chunk c is in flight (the engine latches it) and its data is pre-queued in the host controller, so
+    # the bulk stream runs back-to-back with no gaps (the engine misroutes data that arrives with no arm
+    # latched, so data is only ever submitted after its arm). Buffer reuse is gated on the GPU drain fence;
+    # because the engine writes each buffer sequentially, the fence of chunk c-1's first part is enough.
+    dev, usb = self.dev, self.dev.iface.pci_dev.usb
+    sdq, ts, sdma = dev.sdma_queue(0), dev.timeline_signal, dev.sdma
+    with hcq_profile(self.dev, queue_type=dev.hw_copy_queue_t, desc=TracingKey(f"TINY -> {dev.device}", ret=src.nbytes), enabled=PROFILE,
+                     dev_suff="SDMA:0"):
+      cp_size, src_mv = self.b[0].size - 0x4000, src.cast('B')  # 240KB payload per chunk, the last slot carries sentinels
+      nchunks = (src.nbytes + cp_size - 1) // cp_size
+      if nchunks == 0: return
+      PART = 0x1E000  # 120KB: part granularity, GPU copies overlap the stream in part units
+      slots = [(self.b[bi].cpu_view().addr - 0xf000) >> 14 for bi in range(len(self.b))]  # 0xF2 slot base per bounce buffer
+      # Fence flags: per-buffer GPU drain counters, written by the ring after each part copy. They live in the
+      # NVMe SQ window (PCIe 0x820800+, E4-readable at xdata 0xA800+); the 0x822000 window is the engine's
+      # completion queue and CQEs would clobber them there once the queue wraps.
+      # fences[0:2] = full-chunk fence (0xA800/0xA808), fences[2:4] = first-part fence (0xA810/0xA818)
+      fence_bufs = [dev.iface.sys_buf.offset(0x800 + bi * 8, 8) for bi in range(2 * len(self.b))]
+      if not hasattr(self, '_usb_seq'):
+        self._usb_seq, self._usb_prev_tail = 0, (0, 0)
+        self._usb_stage = [alloc_cbuffer(0x40000) for _ in range(2)]  # wire-image staging, one per bounce buffer
+        self._usb_buf_ctr = [0] * len(self.b)
+        for bi in range(2 * len(self.b)): usb.write(0xA800 + bi * 8, bytes(8))  # clear the flag slots
+        # zero both bounce regions once: sentinel polls use EQ semantics and each use carries a fresh sequence
+        # number, so every potential sentinel address must start at zero (stale values can never match).
+        for bi in range(len(self.b)): usb.scsi_write(bytes(0x40000), slot_start=slots[bi])
+
+      def next_sent():
+        self._usb_seq = (self._usb_seq + 1) & 0xFFFFFF
+        return 0x51000000 | self._usb_seq
+
+      # Build the whole ring: one wait on pre-copyin GPU work, then per chunk, per part: [POLL sentinel][copy].
+      q = dev.hw_copy_queue_t().wait(ts, dev.timeline_value - 1)
+      fences, blens, layouts = [], [], []
+      for c in range(nchunks):
+        bi, lsize = c % len(self.b), min(cp_size, src.nbytes - c * cp_size)
+        fence = dev.timeline_value  # value this chunk's signal will carry
+        self._usb_buf_ctr[bi] += 1
+        ctr = self._usb_buf_ctr[bi]  # fence flag value for this chunk
+        parts = []
+        poff, soff = 0, 0  # payload offset within the chunk, stream offset within the bounce buffer
+        while poff < lsize:
+          plen = min(PART, lsize - poff)
+          sval = next_sent()
+          # POLL (func 3 = EQ) until the engine writes this part's sentinel.
+          q.q(sdma.SDMA_OP_POLL_REGMEM | sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(3) | sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1),
+              *data64_le(self.b[bi].va_addr + soff + round_up(plen, 512)), sval, 0xFFFFFFFF,
+              sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
+          q.copy(dest.offset(c * cp_size + poff), self.b[bi].offset(soff), plen)
+          if poff == 0: q.write(fence_bufs[bi + 2], ctr, b64=True)  # first part of this buffer is drained
+          parts.append((poff, plen, soff, sval))
+          soff += round_up(plen, 512) + 512
+          poff += plen
+        q.signal(ts, dev.next_timeline())
+        q.write(fence_bufs[bi], ctr, b64=True)
+        self.b_timeline[bi] = fence
+        fences.append(ctr)
+        blens.append(lsize)
+        layouts.append(parts)
+      cmds = array.array('I', q._q).tobytes()
+
+      # Stage the ring in one EP2 write while the engine is disarmed. The GPU ring is circular, dword 0 is a NOP.
+      rb, pv_start = sdq.ring.nbytes, sdq.put_value
+      off = pv_start % rb
+      if off + len(cmds) > rb:  # zero-fill to the ring end and continue at offset 0
+        sdq.ring.view(off, rb - off, fmt='B')[:] = bytes(rb - off)
+        pv_start, off = pv_start + rb - off, 0
+      sdq.ring.view(off, len(cmds), fmt='B')[:] = cmds
+      sdq.put_value = pv_start + len(cmds)
+
+      def read_fences(bi):  # one E4 read covers both flags of buffer bi
+        d = usb.read(0xA800 + bi * 8, 24)
+        return int.from_bytes(d[0:8], 'little'), int.from_bytes(d[16:24], 'little')  # full, first-part
+
+      def wait_full(bi, ctr):  # wait until chunk `ctr` of buffer `bi` is fully copied to VRAM
+        for _ in range(100000):
+          if read_fences(bi)[0] >= ctr: return
+        raise RuntimeError("USB copyin GPU fence timeout")
+
+      def gate_part(bi, ctr):  # wait until the first part of chunk `ctr` of buffer `bi` is copied
+        for _ in range(100000):
+          full, part = read_fences(bi)
+          if full >= ctr or part >= ctr: return
+        raise RuntimeError("USB copyin GPU fence timeout")
+
+      wait_full(*self._usb_prev_tail)  # previous copyin's tail must be done before its ring entries are overwritten
+      sdq.write_ptr[0] = sdq.put_value
+      sdq.doorbell[0] = sdq.put_value
+
+      def build_chunk(c):
+        # Stage the wire image [part][sentinel][part][sentinel]... (arms cover exact sector counts, no padding).
+        sbuf, smv = self._usb_stage[c % 2]
+        saddr, sptr = ctypes.addressof(sbuf), ctypes.addressof(from_mv(src_mv[c * cp_size:c * cp_size + blens[c]]))
+        total = 0
+        for poff, plen, soff, sval in layouts[c]:
+          ctypes.memmove(saddr + soff, sptr + poff, plen)
+          struct.pack_into('<I', sbuf, soff + round_up(plen, 512), sval)
+          total = soff + round_up(plen, 512) + 512
+        return smv[:total]
+
+      def chunk_sectors(c):
+        return sum((round_up(plen, 512) + 512) // 512 for _, plen, _, _ in layouts[c])
+
+      bulk_async = usb.usb.bulk_write_async
+
+      # Prime: arm+submit chunk 0, then arm chunk 1 while chunk 0 is in flight and pre-queue it too.
+      usb.scsi_write_arm(chunk_sectors(0) * 512, slot_start=slots[0])
+      pending = [bulk_async(build_chunk(0))]
+      if nchunks > 1:
+        staged_next = build_chunk(1)
+        usb.scsi_write_arm(chunk_sectors(1) * 512, slot_start=slots[1])
+        pending.append(bulk_async(staged_next))
+
+      use_f4 = getenv("USB_F4GATE", 1)
+      for c in range(1, nchunks):
+        usb.bulk_wait(pending.pop(0))  # chunk c-1 has landed in SRAM
+        if c + 1 >= nchunks: continue
+        bi = (c + 1) % len(self.b)
+        staged_next = build_chunk(c + 1)  # memcpy overlaps the in-flight chunk (its staging buffer is free)
+        if use_f4:
+          # One blocking control transfer both waits for the fence of chunk c-1's first part and then arms
+          # chunk c+1: the device holds the arm until the fence lands, which is what stops the engine from
+          # streaming into the buffer early. On a wait bail the arm is skipped and the flags come back stale.
+          for _attempt in range(3):
+            resp = usb.scsi_write_arm_f4(chunk_sectors(c + 1) * 512, slot_start=slots[bi], buf_idx=bi, target=fences[c - 1])
+            if ((resp[bi * 8] - fences[c - 1]) & 0xFF) < 0x80 or ((resp[16 + bi * 8] - fences[c - 1]) & 0xFF) < 0x80: break
+          else: raise RuntimeError(f"0xF4 arm bailed waiting for GPU fence {fences[c-1]}, flags {resp.hex()}")
+        else:
+          usb.scsi_write_arm(chunk_sectors(c + 1) * 512, slot_start=slots[bi])
+          gate_part(bi, fences[c - 1])
+        pending.append(bulk_async(staged_next))
+
+      usb.bulk_wait(pending.pop(0))
+      self._usb_prev_tail = ((nchunks - 1) % len(self.b), fences[-1])
+      wait_full(*self._usb_prev_tail)  # the last chunk must be copied before returning
+
   def _copyout(self, dest:memoryview, src:HCQBuffer):
     if not self.dev.is_usb(): return super()._copyout(dest, src)
     self.dev.synchronize()
@@ -916,7 +1061,10 @@ class USBIface(PCIIface):
     self._compute_props()
 
     # special regions
-    self.copy_bufs = [self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)]
+    # Two 256KB SRAM bounce buffers (16 slots each), so USB bulk IN flight for chunk i overlaps the SDMA of chunk i-1.
+    # ctrl_addr tags encode the 16KB slot base for the 0xF2 engine: 0xf000 -> slot 0, 0x4f000 -> slot 16.
+    self.copy_bufs = [self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x40000),
+                      self._dma_region(ctrl_addr=0x4f000, sys_addr=0x240000, size=0x40000)]
     self.sys_buf, self.sys_next_off = self._dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0x200
     self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000)
 
@@ -926,9 +1074,8 @@ class USBIface(PCIIface):
 
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
     # usb allocates uncached and cpu_access in vram. vram writes are faster than sram writes
-    if host and self.sys_next_off + size < self.sys_buf.size:
-      self.sys_next_off += size
-      return self.sys_buf.offset(self.sys_next_off - size, size)
+    # NOTE: host allocs deliberately do NOT use sys_buf (the 0x820000 NVMe SQ region): the GPU's signal writes there
+    # collide with the 0xF2 engine mid-stream. Signals in VRAM are read back via 0xF0 streaming reads instead.
 
     # force devmem
     return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, force_devmem=True, **kwargs)
@@ -1048,7 +1195,8 @@ class AMDDevice(HCQCompiled):
     if getenv("AMD_DISABLE_SDMA"): return None
     if idx in self.sdma_queues: return self.sdma_queues[idx]
     with contextlib.suppress(OSError):
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
+      # USB: 1MB ring, so a full copyin can be staged in one write. GPU rings are circular and dword 0 is an SDMA NOP.
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, (1 << 20) if self.is_usb() else (16 << 20), idx=idx)
     return self.sdma_queues.get(idx, None)
 
   def _ensure_has_local_memory(self, private_segment_size):

tinygrad/runtime/ops_cpu.py

@@ -6,6 +6,7 @@ from tinygrad.helpers import to_mv, from_mv, OSX, WIN, Context, mv_address, supp
 from tinygrad.device import Buffer, BufferSpec, TinyELF, Program, Device
 from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface
 from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, make_cmdbuf, make_signal
+from tinygrad.runtime.support.c import DLL
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
@@ -111,9 +112,9 @@ def encode_queue(q:UOp) -> UOp:
 MAP_JIT = 0x0800
 
 class CPUProgram(Program['CPUDevice']):
-  rt_lib = None
-  try: rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1')
-  except OSError: pass
+  rt_lib, libm = DLL('rt', 'System' if OSX else 'kernel' if WIN else 'gcc_s'), DLL('m', 'm')
+
+  def _load(self, lib, base=0): return lib if lib[:4] != libc.ELFMAG.encode() else jit_loader(lib, base=base, link_libs=[self.libm, self.rt_lib])
 
   def __init__(self, dev:CPUDevice, obj:TinyELF):
     self.dev, self.name, self.signature = dev, obj.name, obj.signature
@@ -125,10 +126,10 @@ class CPUProgram(Program['CPUDevice']):
       ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p
       self.addr = ctypes.windll.kernel32.VirtualAlloc(ctypes.c_void_p(0), ctypes.c_size_t(len(obj.lib)), MEM_COMMIT | MEM_RESERVE,
                                                       PAGE_EXECUTE_READWRITE)
-      ctypes.memmove(self.addr, obj.lib, len(obj.lib))
+      ctypes.memmove(self.addr, (loaded:=self._load(obj.lib, self.addr)), len(loaded))
       ctypes.windll.kernel32.GetCurrentProcess.restype = ctypes.c_void_p
       proc = ctypes.windll.kernel32.GetCurrentProcess()
-      ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(obj.lib)))
+      ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(loaded)))
       self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)
     else:
       # On apple silicon with SPRR enabled (it always is in macos) RWX pages are unrepresentable: https://blog.svenpeter.dev/posts/m1_sprr_gxf/
@@ -137,18 +138,17 @@ class CPUProgram(Program['CPUDevice']):
       self.addr = mv_address(self.mem)
 
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(False)
-      lib = jit_loader(obj.lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m']) if self.lvp else obj.lib
-      self.mem.write(lib)
+      self.mem.write(loaded:=self._load(obj.lib, mv_address(self.mem)))
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True)
 
       # __clear_cache isn't a normal libc function, but a compiler support routine found in libgcc_s for gcc and compiler-rt for clang.
       # libgcc_s comes as shared library but compiler-rt is only a bunch of static library archives which we can't directly load, but fortunately
       # it somehow found its way into libSystem on macos (likely because it used __builtin_clear_cache) and libgcc_s is ~always present on linux
       # Using ["name"] instead of .name because otherwise name is getting mangled: https://docs.python.org/3.12/reference/expressions.html#index-5
-      if CPUProgram.rt_lib is not None: CPUProgram.rt_lib["__clear_cache"](ctypes.c_void_p(self.addr), ctypes.c_void_p(self.addr + len(lib)))
+      if 'rt' in DLL._loaded_: CPUProgram.rt_lib["__clear_cache"](ctypes.c_void_p(self.addr), ctypes.c_void_p(self.addr + len(loaded)))
       else:
         # msync should be a universal POSIX way to do this
-        libc.msync(ctypes.c_void_p(self.addr), len(lib), libc.MS_SYNC | libc.MS_INVALIDATE)
+        libc.msync(ctypes.c_void_p(self.addr), len(loaded), libc.MS_SYNC | libc.MS_INVALIDATE)
 
       self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)
 

tinygrad/runtime/ops_python.py

@@ -214,7 +214,6 @@ class PythonCompiler(Compiler):
 class PythonRenderer(Renderer):
   code_for_op = python_alu
   compiler = PythonCompiler()
-  casted_consts: bool = True
 
   def __init__(self, target:Target):
     assert (emu:=getenv("EMULATE", "")) == "", ("EMULATE is deprecated, use DEV=PYTHON::" +

tinygrad/runtime/ops_qcom.py

@@ -23,8 +23,8 @@ def dcache_flush():
   buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU)
   i = UOp.range(n, 0, dtype=dtypes.int)
   flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");')
-  sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush"))
-  prg = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())))), Device["CPU"].renderer)
+  sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush"), tag=1)
+  prg = to_program(sink, Device["CPU"].renderer)
   return Device["CPU"].runtime(prg.to_elf())
 
 #Parse C-style defines: <regname>_<field_x>__SHIFT and <regname>_<field_y>__MASK from the adreno module into the following format:

tinygrad/runtime/support/c.py

@@ -91,7 +91,7 @@ class DLL(ctypes.CDLL):
 
   @staticmethod
   def findlib(nm:str, paths:list[str], extra_paths=[]):
-    if nm == 'libc' and OSX: return '/usr/lib/libc.dylib'
+    if nm in ('libc', 'm') and OSX: return f'/usr/lib/lib{nm.removeprefix("lib")}.dylib'
     if pathlib.Path(path:=getenv(nm.replace('-', '_').upper()+"_PATH", '')).is_file(): return path
     for p in paths:
       libpaths = {"posix": [d for d in os.environ.get('LD_LIBRARY_PATH', '').split(os.pathsep) if d] + ["/usr/lib64", "/usr/lib", "/usr/local/lib"],

tinygrad/runtime/support/compiler_cpu.py

@@ -1,10 +1,9 @@
 import subprocess
 from tinygrad.device import Compiler
-from tinygrad.helpers import getenv, capstone_flatdump
-from tinygrad.runtime.support.elf import jit_loader
+from tinygrad.helpers import getenv, capstone_flatdump, cpu_objdump
 
 class ClangCompiler(Compiler):
-  def __init__(self, arch:list[str], cachekey="compile_clang_jit"):
+  def __init__(self, arch:list[str], cachekey="compile_clang_obj"):
     assert len(arch) >= 2, f"invalid arch string: {','.join(arch)!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
     self.arch, cpu, *feats = arch
     match self.arch:
@@ -16,15 +15,13 @@ class ClangCompiler(Compiler):
       case _: raise RuntimeError(f"unsupported arch: {self.arch!r}")
     super().__init__(f"{cachekey}_{'_'.join(arch)}")
 
-  def compile_to_obj(self, src:str) -> bytes:
+  def compile(self, src:str) -> bytes:
     """Compile C source to ELF object file (before linking)."""
     # -fno-math-errno is required for __builtin_sqrt to become an instruction instead of a function call
     return subprocess.check_output([getenv("CC", 'clang'), '-c', '-x', 'c', '-O2', '-fPIC', '-ffreestanding', '-fno-math-errno', '-nostdlib',
                                     '-fno-ident', f'--target={self.arch}-none-unknown-elf', *self.args, '-', '-o', '-'], input=src.encode('utf-8'))
 
-  def compile(self, src:str) -> bytes: return jit_loader(self.compile_to_obj(src))
-
-  def disassemble(self, lib:bytes): return capstone_flatdump(lib, self.arch)
+  def disassemble(self, lib: bytes): cpu_objdump(lib)
 
 
 class X86Compiler(Compiler):

tinygrad/runtime/support/compiler_llvm.py

@@ -1,7 +1,6 @@
 import ctypes
 from tinygrad.device import Compiler, CompileError
-from tinygrad.helpers import getenv, capstone_flatdump, amdgpu_disassemble, unwrap, DEBUG
-from tinygrad.runtime.support.elf import jit_loader
+from tinygrad.helpers import getenv, cpu_objdump, amdgpu_disassemble, unwrap, DEBUG
 from tinygrad.runtime.autogen import llvm
 
 def cerr(): return ctypes.pointer(ctypes.pointer(ctypes.c_char()))
@@ -11,7 +10,6 @@ def expect(x, err, ret=None):
   return ret
 
 class LLVMCompiler(Compiler):
-  jit = True
   def __init__(self, arch:str, processor:str, feats:str, cache_key=None):
     for component in ['Target', 'TargetInfo', 'TargetMC', 'AsmParser', 'AsmPrinter']:
       getattr(llvm, "LLVMInitialize" + {'arm64': 'AArch64', 'x86_64': 'X86', 'riscv64': 'riscv64'}.get(arch, "AMDGPU") + component)()
@@ -43,13 +41,13 @@ class LLVMCompiler(Compiler):
         self.diag_msgs.append(msg)
     self.handle_diag = handle_diag
     llvm.LLVMContextSetDiagnosticHandler(self.context, handle_diag, None)
-    super().__init__(cache_key or f"compile_llvm_{processor}_{feats}{'_jit' if self.jit else ''}{'_opt' if opt else ''}")
+    super().__init__(cache_key or f"compile_llvm_{processor}_{feats}{'_opt' if opt else ''}")
 
   def __del__(self):
     if hasattr(self, 'pbo'): llvm.LLVMDisposePassBuilderOptions(self.pbo)
     if hasattr(self, 'context'): llvm.LLVMContextDispose(self.context)
 
-  def compile_to_obj(self, src:str) -> bytes:
+  def compile(self, src:str) -> bytes:
     self.diag_msgs.clear()
     src_buf = llvm.LLVMCreateMemoryBufferWithMemoryRangeCopy(ctypes.create_string_buffer(src_bytes:=src.encode()), len(src_bytes), b'src')
     mod = expect(llvm.LLVMParseIRInContext(self.context, src_buf, ctypes.pointer(m:=llvm.LLVMModuleRef()), err:=cerr()), err, m)
@@ -64,9 +62,6 @@ class LLVMCompiler(Compiler):
     if self.diag_msgs: raise RuntimeError("llvm diagnostic: " + "\n".join(self.diag_msgs))
     return obj
 
-  def compile(self, src:str) -> bytes: return jit_loader(self.compile_to_obj(src)) if self.jit else self.compile_to_obj(src)
-
-
 class CPULLVMCompiler(LLVMCompiler):
   def __init__(self, arch:list[str], cache_key=None):
     assert len(arch) >= 2, f"invalid arch string: {','.join(arch)!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
@@ -78,10 +73,9 @@ class CPULLVMCompiler(LLVMCompiler):
     # +reserve-x18 here does the same thing as -ffixed-x18 in ClangCompiler, see comments there for why it's needed on arm osx
     super().__init__(self.arch, cpu, ('+reserve-x18,' if self.arch == "arm64" else '') + featstr, cache_key)
 
-  def disassemble(self, lib:bytes): capstone_flatdump(lib, self.arch)
+  def disassemble(self, lib: bytes): cpu_objdump(lib)
 
 class AMDLLVMCompiler(LLVMCompiler):
-  jit = False
   def __init__(self, arch: str):
     self.arch = arch
     super().__init__("AMDGPU", self.arch, "+cumode")

tinygrad/runtime/support/elf.py

@@ -1,4 +1,4 @@
-import struct, ctypes, ctypes.util
+import struct, ctypes
 from dataclasses import dataclass
 from tinygrad.helpers import getbits, i2u, unwrap
 from tinygrad.runtime.autogen import libc
@@ -6,13 +6,13 @@ from tinygrad.runtime.autogen import libc
 @dataclass(frozen=True)
 class ElfSection: name:str; header:libc.Elf64_Shdr|libc.Elf32_Shdr; content:bytes # noqa: E702
 
-def link_sym(sym:str, libs:list[str]) -> int:
+def link_sym(sym:str, libs:list[ctypes.CDLL]) -> int:
   for lib in libs:
-    try: return unwrap(ctypes.cast(getattr(ctypes.CDLL(ctypes.util.find_library(lib)), sym), ctypes.c_void_p).value)
+    try: return unwrap(ctypes.cast(getattr(lib, sym), ctypes.c_void_p).value)
     except (OSError, AttributeError): pass
   raise RuntimeError(f'Attempting to relocate against an undefined symbol {sym}')
 
-def elf_loader(blob:bytes, force_section_align:int=1, link_libs:list[str]|None=None) -> tuple[memoryview, list[ElfSection], list[tuple]]:
+def elf_loader(blob:bytes, force_section_align:int=1, link_libs:list[ctypes.CDLL]|None=None) -> tuple[memoryview, list[ElfSection], list[tuple]]:
   assert blob[:4] == libc.ELFMAG.encode(), "blob is not an ELF, missing magic bytes"
   ecls = {libc.ELFCLASS32: "Elf32", libc.ELFCLASS64: "Elf64"}[blob[libc.EI_CLASS]]
 
@@ -49,7 +49,7 @@ def elf_loader(blob:bytes, force_section_align:int=1, link_libs:list[str]|None=N
 
   return memoryview(image), sections, relocs
 
-def jit_loader(obj: bytes, base:int=0, link_libs:list[str]|None=None) -> bytes:
+def jit_loader(obj: bytes, base:int=0, link_libs:list[ctypes.CDLL]|None=None) -> bytes:
   image_, _, relocs = elf_loader(obj, link_libs=link_libs)
   image = bytearray(image_)
 

tinygrad/runtime/support/hcq2.py

@@ -252,7 +252,7 @@ def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
   devs, queue = get_submit(calls[0]).src[0].arg
   body = make_submit(*[cmd for c in calls for cmd in get_submit(c).src[0].src], devs=devs, queue=queue).sink()
   return make_call(f"submit {queue} ({len(calls)})", body,
-    replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates())))
+    replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()).simplify()))
 
 def merge_queues(linear:UOp) -> UOp:
   new_src:list[UOp] = []
@@ -426,7 +426,7 @@ def merge_batch(batch:list[UOp]) -> UOp:
   cmds = [c.src[0].src[0].call(*[_lane_arg(a.without_after, j, tables + off) for a in c.src[1:]], UOp.variable("_device_num", 0, 1 << 30).bind(j))
           for (c, j, _), off in zip(lanes, offs)]
 
-  info = HCQInfo((HCQ_RUNTIME_DEV.value,), sum((c.arg.aux.estimates for c in batch), start=Estimates()),
+  info = HCQInfo((HCQ_RUNTIME_DEV.value,), sum((c.arg.aux.estimates for c in batch), start=Estimates()).simplify(),
                  input_idxs=tuple(x for c in batch for x in c.arg.aux.input_idxs), kernels=tuple(k for c in batch for k in c.arg.aux.kernels))
   body = UOp.custom_function("hcq", make_submit(*cmds, devs=HCQ_RUNTIME_DEV.value, queue="SUBMIT:0").sink())
   return body.call(*[s for c in batch for s in c.src[1:] if s.without_after.tag != "inputs"], name=f"hcq_submitter ({len(batch)})", aux=info)
@@ -591,7 +591,7 @@ class HCQ2Compiled(Compiled):
 
   def new_buffer(self, b:UOp, cache:bool) -> Buffer:
     if cache or b.tag in HCQ_CACHE_TAGS:
-      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=b.tag != "program", cpu_access=True, nolru=True))
+      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=b.tag not in ("program","kernargs"), cpu_access=True,nolru=True))
     return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))
 
   @functools.cache

tinygrad/runtime/support/usb.py

@@ -1,6 +1,6 @@
 import ctypes, struct, time, functools, itertools
 from tinygrad.runtime.autogen import libusb
-from tinygrad.helpers import DEBUG, DEV, to_mv, round_up, ceildiv
+from tinygrad.helpers import DEBUG, DEV, to_mv, from_mv, round_up, ceildiv
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support import c
 
@@ -35,6 +35,12 @@ class USB3:
     self._tags, self._transferred = itertools.count(1), ctypes.c_int(0)
     self._bulk_buf, self._bulk_mv = alloc_cbuffer(4 << 20)
     self._ctrl_buf, self._ctrl_mv = alloc_cbuffer(0x1000)
+    # async bulk OUT state: tag -> (transfer ptr, keepalive memoryview)
+    self._async_seq = itertools.count(1)
+    self._async_pending: dict[int, tuple[c.POINTER[libusb.struct_libusb_transfer], memoryview]] = {}
+    self._async_free: list[c.POINTER[libusb.struct_libusb_transfer]] = []
+    self._async_err = 0
+    self._async_cb = libusb.libusb_transfer_cb_fn(self._on_bulk_done)
 
     self.handle = c.init_c_var(c.POINTER[libusb.struct_libusb_device_handle], lambda x: checked(libusb.libusb_open)(dev, x))
 
@@ -73,6 +79,29 @@ class USB3:
       (self.handle, 0x02, self._bulk_buf, len(payload), self._transferred, timeout)
     assert self._transferred.value == len(payload), f"bulk OUT short write: {self._transferred.value}/{len(payload)} bytes"
 
+  def _on_bulk_done(self, xfer):
+    # NOTE: runs inside libusb event handling; exceptions here are unraisable, so latch errors for bulk_wait.
+    if xfer.contents.status != 0 or xfer.contents.actual_length != xfer.contents.length: self._async_err = xfer.contents.status or -1
+    self._async_pending.pop(int(xfer.contents.user_data or 0), None)
+    self._async_free.append(xfer)  # transfers are reused (re-submit is cheaper than alloc/free)
+
+  def bulk_write_async(self, payload:memoryview, timeout:int=10000) -> int:
+    """Zero-copy async bulk OUT on EP 0x02. Returns a tag for bulk_wait. The payload must stay alive until bulk_wait."""
+    assert payload.contiguous, "bulk_write_async requires a contiguous buffer"
+    tr = self._async_free.pop() if self._async_free else libusb.libusb_alloc_transfer(0)
+    tag = next(self._async_seq)
+    tr.contents.dev_handle, tr.contents.endpoint, tr.contents.type = self.handle, 0x02, libusb.LIBUSB_TRANSFER_TYPE_BULK
+    tr.contents.timeout, tr.contents.length = timeout, len(payload)
+    tr.contents.buffer = ctypes.cast(from_mv(payload, ctypes.c_ubyte), ctypes.POINTER(ctypes.c_ubyte))
+    tr.contents.callback, tr.contents.user_data = self._async_cb, tag
+    self._async_pending[tag] = (tr, payload)
+    checked(libusb.libusb_submit_transfer, "async bulk OUT submit failed")(tr)
+    return tag
+
+  def bulk_wait(self, tag:int):
+    while tag in self._async_pending: checked(libusb.libusb_handle_events)(None)
+    if self._async_err: raise RuntimeError(f"async bulk OUT failed: status={self._async_err}")
+
   def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
     if length > len(self._bulk_mv): self._bulk_buf, self._bulk_mv = alloc_cbuffer(length)
     checked(libusb.libusb_bulk_transfer, "bulk IN 0x81 failed")(self.handle, 0x81, self._bulk_buf, length, self._transferred, timeout)
@@ -160,15 +189,34 @@ class CustomASM24Controller:
     """Write to chip XDATA via vendor control OUT (bRequest=0xE5). wValue=addr, wIndex=val."""
     for off, val in enumerate(data): self.usb.control_write(0xE5, value=base_addr + off, index=val)
 
-  def scsi_write(self, buf:bytes):
+  def scsi_write(self, buf:bytes, slot_start:int=0):
     """Write to SRAM via 0xF2 vendor command + bulk OUT."""
     buf_padded = buf + b'\x00' * (round_up(len(buf), 512) - len(buf))
-    sectors = len(buf_padded) // 512
-    num_slots = ceildiv(len(buf_padded), 0x4000)  # 16KB per slot
-    windex = (num_slots & 0xFF) << 8
-    self.usb.control_write(0xF2, value=sectors, index=windex)
+    self.scsi_write_arm(len(buf_padded), slot_start)
     self.usb.bulk_write(buf_padded)
 
+  def scsi_write_arm(self, nbytes:int, slot_start:int=0):
+    assert nbytes % 512 == 0, f"scsi write arm requires 512-byte aligned size, got {nbytes}"
+    sectors = nbytes // 512
+    num_slots = ceildiv(nbytes, 0x4000)  # 16KB per slot
+    self.usb.control_write(0xF2, value=sectors, index=(slot_start & 0xFF) | ((num_slots & 0xFF) << 8))
+
+  def bulk_wait(self, tag:int): self.usb.bulk_wait(tag)
+
+  def scsi_write_arm_f4(self, nbytes:int, slot_start:int, buf_idx:int, target:int) -> bytes:
+    """0xF4 (custom firmware): blocks device-side until the GPU fence flag for bounce buffer buf_idx reaches
+    target (a per-buffer counter, low byte compared mod 256), then arms the SRAM DMA engine for a bulk OUT of
+    nbytes. The read completing tells us the buffer is free; the late arm stops the engine from streaming into
+    it early. Returns the 32-byte fence flag window; if the device bailed on the wait the arm was skipped and
+    the returned flags will be stale (callers must check)."""
+    assert nbytes % 512 == 0, f"F4 arm requires 512-byte aligned size, got {nbytes}"
+    sectors = nbytes // 512
+    assert sectors <= 0x1FF, sectors
+    num_slots = ceildiv(nbytes, 0x4000)
+    wValue = sectors | ((target & 0x7F) << 9)  # target bit 7 rides in wIndex bit 14
+    wIndex = (slot_start & 0xFF) | ((num_slots & 0x3F) << 8) | (((target >> 7) & 1) << 14) | ((buf_idx & 1) << 15)
+    return bytes(self.usb.control_read(0xF4, 32, value=wValue, index=wIndex, timeout=30000))
+
   def scsi_read_arm(self, size:int):
     windex = (ceildiv(size, 0x4000) & 0xFF) << 8
     self.usb.control_write(0xF2, value=(ceildiv(size, 512) & 0x7FFF) | 0x8000, index=windex)
@@ -194,7 +242,9 @@ class USBMMIOInterface(MMIOInterface):
   def __setitem__(self, index, data):
     off, _ = self._off_from_index(index)
     data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data)
-    if not self.pcimem: self.usb.scsi_write(data) if self.addr == 0xf000 else self.usb.write(self.addr + off, data)
+    if not self.pcimem:
+      # addr tags 0xf000 + slot_start*0x4000 are SRAM bounce buffers, written via 0xF2 bulk DMA
+      self.usb.scsi_write(data, slot_start=(self.addr - 0xf000) >> 14) if self.addr in (0xf000, 0x4f000) else self.usb.write(self.addr + off, data)
     else: self.usb.pcie_mem_write(self.addr+off, data)
 
   def view(self, offset:int=0, size:int|None=None, fmt=None):

tinygrad/uop/ops.py

@@ -94,11 +94,10 @@ def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str:
   return ret
 
 def shape_to_shape_arg(arg:tuple[sint, ...]) -> UOp:
-  for x in arg:
-    if isinstance(x, UOp) and not dtypes.is_int(x.dtype): raise RuntimeError(f"shape must be int, got {x.dtype} in {arg}")
-  if len(arg) == 0: return UOp(Ops.STACK)
-  elif len(arg) == 1: return UOp.const(arg[0], dtypes.weakint)
-  else: return UOp(Ops.STACK, src=tuple(UOp.const(x) if isinstance(x, int) else x for x in arg))
+  src = tuple(x if isinstance(x, UOp) else UOp.const(x) for x in arg)
+  for x in src:
+    if not dtypes.is_int(x.dtype): raise RuntimeError(f"shape must be int, got {x.dtype} in {arg}")
+  return src[0] if len(src) == 1 else UOp(Ops.STACK, src=src)
 
 def consumer_map_from_toposort(lst:Iterable[UOp]):
   ret: dict[UOp, dict[UOp, None]] = {}
@@ -615,6 +614,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     # NOTE: it always has to be STACK now, even if they are all the same
     if isinstance(b, tuple): return UOp.stack(*[UOp.const(c, dtype) for c in b])
     return UOp(Ops.CONST, dtype, arg=dtype.const(b), src=())
+  # weak CONST with width on the CAST. TODO: this is the final const
+  @staticmethod
+  def cconst(b:ConstLike, dtype:DType): return UOp(Ops.CAST, dtype, src=(UOp.const(b),), arg=dtype)
   @staticmethod
   def range(end:sint, axis_id, axis_type=AxisType.WEAK, *arg, dtype=dtypes.weakint, src=(), **kwargs):
     return UOp(Ops.RANGE, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)

tinygrad/uop/render.py

@@ -39,6 +39,8 @@ renderer = PatternMatcher([
   (UPat(Ops.RANGE, dtypes.void, name="x"), lambda x: f"loop{x.arg[0]}"),
   (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
   (UPat(Ops.CONST, name="x"), lambda x: str(x.val)),
+  # CAST states the width, the weak CONST carries the value
+  (UPat.cvar("c", dtypes.weaks+(dtypes.bool,)).cast(), lambda c: str(c.val)),
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"),
   (UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"),
   (UPat(Ops.RECIPROCAL, name="x"), lambda ctx,x: f"(1/{ctx[x.src[0]]})"),

tinygrad/uop/spec.py

@@ -200,11 +200,13 @@ spec_tensor = PatternMatcher([
 
 # these ops can exist in programs but not the tensor spec. example: LOAD
 spec_program = PatternMatcher([
+  # a literal is CAST(dt, CONST(value)), so its inner CONST is the one weak node a program may contain
+  (UPat(Ops.CONST, dtype=dtypes.weaks, name="x"), lambda x: x.dtype is dtypes.from_py(x.val)),
   # index and weak dtypes are not allowed in programs
   (UPat(GroupOp.All, (dtypes.weakint, dtypes.weakfloat)), lambda: False),
 
   # allow special SHRINK
-  (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST))), lambda: True),
+  (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST).or_casted())), lambda: True),
 
   # movement ops are not allowed in programs
   (UPat(GroupOp.Movement), lambda: False),
@@ -223,12 +225,6 @@ spec_program = PatternMatcher([
   (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.int32),), name="s"), lambda s,x: matches_dtype(x, s.dtype) and isinstance(s.arg, str)),
 ])+spec_shared
 
-# migration: on a casted_consts renderer every literal is CAST(dt, CONST(value)) with a weak inner CONST
-spec_program_casted_consts = PatternMatcher([
-  (UPat(Ops.CONST, dtype=dtypes.weaks, name="x"), lambda x: x.dtype is dtypes.from_py(x.val)),
-  (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CAST, src=(UPat(Ops.CONST),)))), lambda: True),
-])+spec_program
-
 spec_hcq = PatternMatcher([
   (UPat(Ops.GETADDR, dtypes.uint64, src=(UPat((Ops.BUFFER, Ops.PARAM)).or_after(),), name="x"), lambda x: is_device(x.arg)),
   (UPat(Ops.PROGRAM, dtypes.void, src=(UPat((Ops.BUFFER, Ops.PARAM)).or_after(),)), lambda: True),

tinygrad/uop/symbolic.py

@@ -2,7 +2,7 @@
 import math
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid, bitcast
+from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid, bitcast, truncate
 from tinygrad.helpers import partition, all_same, prod, flatten, unwrap, IMAGE, dedup
 from tinygrad.uop.divandmod import div_and_mod_symbolic
 from tinygrad.uop.movement import mop_cleanup
@@ -22,7 +22,7 @@ def simplify_pow(x:UOp, c:UOp) -> UOp|None:
 
 def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
   if c.dtype.fmt is None or root.dtype.fmt is None or c.dtype.itemsize != root.dtype.itemsize: return None
-  return root.const_like(bitcast(c.val, c.dtype, root.dtype))
+  return root.const_like(bitcast(truncate[c.dtype](c.val), c.dtype, root.dtype))
 
 # const folding works for CONST, STACK, and casted CONST
 const_folding_pat = UPat.any(UPat((Ops.CONST, Ops.STACK)), UPat(Ops.CAST, src=(UPat(Ops.CONST),)))