UOp.dtype is dtype_from_uop

axes, not axis

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

@@ -194,21 +194,29 @@ runs:
         echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
         echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
 
+        installed=true
+        for pkg in $pkgs; do
+          info=$(dpkg-query -W -f='${db:Status-Abbrev} ${Version}' "$pkg" 2> /dev/null || true)
+          echo "${pkg}: ${info:-not in dpkg database}"
+          [[ "$info" == ii* ]] || installed=false
+        done
+        echo "installed=$installed" >> "$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' || inputs.ninja == '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' && steps.apt-pkgs.outputs.installed == 'false'
       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' || inputs.ninja == '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' && steps.apt-pkgs.outputs.installed == 'false'
       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' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && steps.apt-pkgs.outputs.installed == 'false'
       shell: bash
       run: |
         sudo apt -qq update || true

.github/workflows/benchmark.yml

@@ -94,7 +94,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -140,7 +140,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -188,7 +188,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -230,7 +230,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -275,7 +275,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code

.github/workflows/test.yml

@@ -390,10 +390,10 @@ jobs:
         run: |
           parallel --link --tagstring '[{1}]' '{2}' \
             ::: llama 'llama q4' qwen3.5 qwen \
-            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
+            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
                 # NOTE: qwen is dumb and only knows about female chickens
 
 # ****** Models Tests ******

conftest.py

@@ -2,7 +2,7 @@ import os, pytest, signal, threading
 
 @pytest.hookimpl(wrapper=True)
 def pytest_runtest_call(item):
-  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
+  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 90)), os.kill, args=(os.getpid(), signal.SIGABRT))
   t.start()
   try: yield
   finally:

examples/mlperf/model_train.py

@@ -1667,7 +1667,7 @@ def train_llama3():
 def train_gptoss():
   from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, fclip_grads
 
   BENCHMARK = getenv("BENCHMARK")
 
@@ -1785,12 +1785,10 @@ def train_gptoss():
 
     Tensor.realize(loss, *grads)
 
-    grad_norm = clip_grads(grads, 1, 1.0)
-    optim.fstep(grads, grad_norm)
+    clipped_grads, grad_norm = fclip_grads(grads, 1.0)
+    optim.fstep(clipped_grads, grad_norm)
     scheduler.step()
 
-    for g in grads: g.assign(0)
-
     loss_cpu = loss.flatten().float().to("CPU")
     lr_cpu = optim.lr.float().to("CPU")
     grad_norm_cpu = grad_norm.float().to("CPU")

examples/mlperf/models/gpt_oss.py

@@ -282,14 +282,14 @@ def apply_grad(grad_buf:Tensor, new_grad:UOp):
   pads = _get_pads(new_grad)
   if len(pads) <= 1:
     new_grad = new_grad.cast(grad_buf.dtype)
-    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
+    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(new_grad))
     return
   cur = grad_buf.uop
   for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
     if pad.op == Ops.PAD:
-      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
+      grad_shrink = tuple((p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg))
       buf_slice = cur.shrink(grad_shrink)
-      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
+      cur = cur.after(buf_slice.store(pad.src[0].cast(cur.dtype)))
     else:
       cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
   grad_buf.uop = cur

examples/mlperf/optim.py

@@ -27,6 +27,11 @@ def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
   for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
   return total_norm
 
+def fclip_grads(grads:list[Tensor], clip_norm) -> Tensor:
+  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+  scale = (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
+  return [(g * scale).cast(g.dtype) for g in grads], total_norm
+
 class GradAccClipAdamW(Optimizer):
   def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
     super().__init__(params, lr, device, fused)

extra/gemm/cdna_asm_gemm.py

@@ -215,7 +215,7 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   k = UOp.range(K, 0, AxisType.REDUCE)
   mul = (A.flatten().index((m*UOp.const(K)+k))*
          B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
-  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
+  red = mul.reduce(k, arg=Ops.ADD).cast(C.dtype)
   store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
   return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
 

extra/gemm/gemm_mxfp4.py

@@ -20,34 +20,39 @@ def v_mfma_fp4(dst, a, b, opsel, opsel_hi, scale_a, scale_b):
 def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
   k = Kernel()
   scale_k = K // 32
+  k.emit(s_and_b32(s[1], s[1], LIT, 65535))
   if (tile_m, tile_n) == (128, 512):
-    k.emit(s_and_b32(s[1], s[1], LIT, 65535))
     k.emit(s_mov_b32(s[47], s[2]))
     k.emit(s_mov_b32(s[48], s[3]))
-    k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[8], 0))
-    k.emit(s_mov_b32(s[9], 0))
-    k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[36], N))
-    k.emit(s_mov_b32(s[37], K))
-    k.emit(s_mov_b32(s[38], K))
-    k.emit(s_mov_b32(s[43], M))
-    k.emit(s_mov_b32(s[44], N))
-    k.emit(s_mov_b32(s[45], K))
-    k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[39], scale_k))
-    k.emit(s_mov_b32(s[40], scale_k))
-    k.emit(v_lshrrev_b32_e32(v[1], 10))
-    k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
-    k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
-    k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
-    k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
-    k.emit(v_lshrrev_b32_e32(v[3], 6))
-    k.emit(v_and_b32_e32(v[0], 63))
-    k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
-    k.emit(s_waitcnt(49279))
+  k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
+  k.emit(s_mov_b32(s[8], 0))
+  k.emit(s_mov_b32(s[9], 0))
+  k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
+  k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
+  k.emit(s_mov_b32(s[36], N))
+  k.emit(s_mov_b32(s[37], K))
+  k.emit(s_mov_b32(s[38], K))
+  k.emit(s_mov_b32(s[43], M))
+  k.emit(s_mov_b32(s[44], N))
+  k.emit(s_mov_b32(s[45], K))
+  k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
+  k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
+  k.emit(s_mov_b32(s[39], scale_k))
+  k.emit(s_mov_b32(s[40], scale_k))
+  k.emit(v_lshrrev_b32_e32(v[1], 10))
+  k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
+  k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
+  k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
+  k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
+  k.emit(v_lshrrev_b32_e32(v[3], 6))
+  k.emit(v_and_b32_e32(v[0], 63))
+  if (tile_m, tile_n) == (256, 256):
+    k.emit(s_mov_b32(s[49], s[2]))
+    k.emit(s_mov_b32(s[47], s[3]))
+  k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
+  k.emit(s_waitcnt(49279))
+
+  if (tile_m, tile_n) == (128, 512):
     for i in range(2):
       k.emit(s_mov_b32(s[6 + i * 8], -16))
       k.emit(s_mov_b32(s[10 + i * 12], -16))
@@ -1213,31 +1218,6 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
     k.emit(s_waitcnt())
     k.emit(s_endpgm())
   elif (tile_m, tile_n) == (192, 256):
-    k.emit(s_and_b32(s[1], s[1], LIT, 65535))
-    k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[8], 0))
-    k.emit(s_mov_b32(s[9], 0))
-    k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[36], N))
-    k.emit(s_mov_b32(s[37], K))
-    k.emit(s_mov_b32(s[38], K))
-    k.emit(s_mov_b32(s[43], M))
-    k.emit(s_mov_b32(s[44], N))
-    k.emit(s_mov_b32(s[45], K))
-    k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[39], scale_k))
-    k.emit(s_mov_b32(s[40], scale_k))
-    k.emit(v_lshrrev_b32_e32(v[1], 10))
-    k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
-    k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
-    k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
-    k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
-    k.emit(v_lshrrev_b32_e32(v[3], 6))
-    k.emit(v_and_b32_e32(v[0], 63))
-    k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
-    k.emit(s_waitcnt(49279))
     k.emit(s_mul_i32(s[63], LIT, 8, 192))
     k.emit(v_cvt_f32_u32_e32(v[4], s[63]))
     k.emit(s_sub_i32(s[62], 0, s[63]))
@@ -2234,49 +2214,22 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
     k.emit(s_waitcnt())
     k.emit(s_endpgm())
   elif (tile_m, tile_n) == (256, 256):
-    k.emit(s_and_b32(s[1], s[1], LIT, 65535))
-    k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[8], 0))
-    k.emit(s_mov_b32(s[9], 0))
-    k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[40], N))
-    k.emit(s_mov_b32(s[41], K))
-    k.emit(s_mov_b32(s[42], K))
-    k.emit(s_mov_b32(s[43], M))
-    k.emit(s_mov_b32(s[44], N))
-    k.emit(s_mov_b32(s[45], K))
-    k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
-    k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
-    k.emit(s_mov_b32(s[36], scale_k))
-    k.emit(s_mov_b32(s[37], scale_k))
-    k.emit(v_lshrrev_b32_e32(v[1], 10))
-    k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
-    k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
-    k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
-    k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
-    k.emit(v_lshrrev_b32_e32(v[3], 6))
-    k.emit(v_and_b32_e32(v[0], 63))
-    k.emit(s_mov_b32(s[46], s[2]))
-    k.emit(s_mov_b32(s[47], s[3]))
-    k.emit(v_readfirstlane_b32_e32(v[49], v[3]))
-    k.emit(s_waitcnt(49279))
     k.emit(s_add_u32(s[55], s[44], LIT, 255))
     k.emit(s_lshr_b32(s[54], s[55], 8))
     k.emit(s_mul_i32(s[48], s[54], s[47]))
-    k.emit(s_add_i32(s[48], s[48], s[46]))
+    k.emit(s_add_i32(s[48], s[48], s[49]))
     k.emit(s_add_u32(s[55], s[43], LIT, 255))
     k.emit(s_lshr_b32(s[52], s[55], 8))
     k.emit(s_lshl_b32(s[52], s[52], 5))
-    k.emit(s_mov_b32(s[46], 0))
+    k.emit(s_mov_b32(s[49], 0))
     k.label('L2_00E8')
     k.emit(s_cmp_lt_i32(s[48], s[52]))
     k.emit(s_cbranch_scc1(3), target='L2_00FC')
     k.emit(s_sub_i32(s[48], s[48], s[52]))
-    k.emit(s_add_i32(s[46], s[46], 32))
+    k.emit(s_add_i32(s[49], s[49], 32))
     k.emit(s_branch(65531), target='L2_00E8')
     k.label('L2_00FC')
-    k.emit(s_sub_i32(s[54], s[54], s[46]))
+    k.emit(s_sub_i32(s[54], s[54], s[49]))
     k.emit(s_cmp_lt_i32(s[54], 32))
     k.emit(s_cbranch_scc1(3), target='L2_0114')
     k.emit(s_lshr_b32(s[47], s[48], 5))
@@ -2311,7 +2264,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
     k.emit(s_mul_i32(s[52], s[54], s[47]))
     k.emit(s_sub_i32(s[52], s[48], s[52]))
     k.label('L2_0194')
-    k.emit(s_add_i32(s[46], s[52], s[46]))
+    k.emit(s_add_i32(s[49], s[52], s[49]))
     k.emit(s_mov_b32(s[6], -16))
     k.emit(s_mov_b32(s[10], -16))
     k.emit(s_mov_b32(s[18], -16))
@@ -2328,18 +2281,18 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
     k.emit(s_or_b32(s[9], s[9], LIT, 262144))
     k.emit(s_or_b32(s[17], s[17], LIT, 262144))
     k.emit(s_or_b32(s[13], s[13], LIT, 262144))
-    k.emit(s_lshr_b32(s[41], s[41], 1))
-    k.emit(s_mul_i32(s[52], s[41], s[43]))
+    k.emit(s_lshr_b32(s[37], s[37], 1))
+    k.emit(s_mul_i32(s[52], s[37], s[43]))
     k.emit(s_mov_b32(s[14], s[52]))
-    k.emit(s_lshr_b32(s[42], s[42], 1))
-    k.emit(s_mul_i32(s[52], s[42], s[44]))
+    k.emit(s_lshr_b32(s[38], s[38], 1))
+    k.emit(s_mul_i32(s[52], s[38], s[44]))
     k.emit(s_mov_b32(s[18], s[52]))
     k.emit(s_add_u32(s[52], s[43], 31))
     k.emit(s_lshr_b32(s[52], s[52], 5))
     k.emit(s_lshl_b32(s[52], s[52], 5))
-    k.emit(s_mul_i32(s[53], s[52], s[36]))
+    k.emit(s_mul_i32(s[53], s[52], s[39]))
     k.emit(s_mov_b32(s[22], s[53]))
-    k.emit(s_mul_i32(s[53], s[44], s[37]))
+    k.emit(s_mul_i32(s[53], s[44], s[40]))
     k.emit(s_mov_b32(s[26], s[53]))
     k.emit(s_mov_b32(s[23], LIT, 131072))
     k.emit(s_mov_b32(s[27], LIT, 131072))
@@ -2356,23 +2309,23 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
     k.emit(v_add_u32_e32(v[5], v[5], v[6]))
     k.emit(v_and_b32_e32(v[4], 1, v[4]))
     k.emit(v_add_u32_e32(v[5], v[5], v[4]))
-    k.emit(v_mul_lo_u32(v[212], s[41], v[5]))
+    k.emit(v_mul_lo_u32(v[212], s[37], v[5]))
     k.emit(v_and_b32_e32(v[4], 7))
     k.emit(v_lshlrev_b32_e32(v[4], 4, v[4]))
     k.emit(v_add_u32_e32(v[212], v[212], v[4]))
-    k.emit(s_lshr_b32(s[52], s[49], 1))
+    k.emit(s_lshr_b32(s[52], s[46], 1))
     k.emit(s_mul_i32(s[52], s[52], 8))
-    k.emit(s_and_b32(s[53], s[49], 1))
+    k.emit(s_and_b32(s[53], s[46], 1))
     k.emit(s_mul_i32(s[53], s[53], 2))
     k.emit(s_add_u32(s[52], s[52], s[53]))
     k.emit(s_mul_i32(s[53], s[47], LIT, 256))
     k.emit(s_add_u32(s[52], s[52], s[53]))
-    k.emit(s_mul_i32(s[52], s[41], s[52]))
+    k.emit(s_mul_i32(s[52], s[37], s[52]))
     k.emit(v_add_u32_e32(v[212], s[52], v[212]))
-    k.emit(s_mul_i32(s[52], s[41], 32))
+    k.emit(s_mul_i32(s[52], s[37], 32))
     for i in range(7):
       k.emit(v_add_u32_e32(v[213 + i * 1], s[52], v[212 + i * 1]))
-    k.emit(s_mul_i32(s[59], LIT, s[49], 1056))
+    k.emit(s_mul_i32(s[59], LIT, s[46], 1056))
     k.emit(s_add_u32(s[59], LIT, s[59], 4096))
     k.emit(v_and_b32_e32(v[4], 15))
     k.emit(v_lshrrev_b32_e32(v[5], 3, v[4]))
@@ -2396,35 +2349,35 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
     k.emit(v_add_u32_e32(v[221], LIT, v[220], 33792))
     k.emit(v_lshlrev_b32_e32(v[222], 2))
     k.emit(s_mul_i32(s[52], s[47], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[49], 32))
+    k.emit(s_mul_i32(s[53], s[46], 32))
     k.emit(s_add_i32(s[52], s[53], s[52]))
-    k.emit(s_mul_i32(s[53], s[52], s[36]))
+    k.emit(s_mul_i32(s[53], s[52], s[39]))
     k.emit(v_add_u32_e32(v[222], s[53], v[222]))
-    k.emit(s_mul_i32(s[53], LIT, s[36], 128))
+    k.emit(s_mul_i32(s[53], LIT, s[39], 128))
     k.emit(v_add_u32_e32(v[223], s[53], v[222]))
-    k.emit(s_mul_i32(s[60], s[49], LIT, 256))
+    k.emit(s_mul_i32(s[60], s[46], LIT, 256))
     k.emit(s_add_i32(s[60], s[60], 0))
     k.emit(v_lshlrev_b32_e32(v[224], 2))
     k.emit(v_add_u32_e32(v[224], 0, v[224]))
     k.emit(v_lshlrev_b32_e32(v[225], 4))
-    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[49], 64))
+    k.emit(s_mul_i32(s[52], s[49], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[46], 64))
     k.emit(s_add_u32(s[52], s[52], s[53]))
-    k.emit(s_mul_i32(s[52], s[52], s[42]))
+    k.emit(s_mul_i32(s[52], s[52], s[38]))
     k.emit(v_add_u32_e32(v[225], s[52], v[225]))
-    k.emit(s_mul_i32(s[52], 16, s[42]))
+    k.emit(s_mul_i32(s[52], 16, s[38]))
     k.emit(v_add_u32_e32(v[226], s[52], v[225]))
     k.emit(v_add_u32_e32(v[227], s[52], v[226]))
     k.emit(v_add_u32_e32(v[228], s[52], v[227]))
     for i in range(4):
       k.emit(v_add_u32_e32(v[229 + i * 1], LIT, v[225 + i * 1], 1024))
     k.emit(v_lshlrev_b32_e32(v[233], 2))
-    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[49], 64))
+    k.emit(s_mul_i32(s[52], s[49], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[46], 64))
     k.emit(s_add_i32(s[52], s[53], s[52]))
-    k.emit(s_mul_i32(s[53], s[52], s[37]))
+    k.emit(s_mul_i32(s[53], s[52], s[40]))
     k.emit(v_add_u32_e32(v[233], s[53], v[233]))
-    k.emit(s_mul_i32(s[52], 32, s[37]))
+    k.emit(s_mul_i32(s[52], 32, s[40]))
     k.emit(v_add_u32_e32(v[234], s[52], v[233]))
     k.emit(s_mov_b32(s[61], LIT, 128))
     k.emit(s_mov_b32(s[62], LIT, 2048))
@@ -2510,18 +2463,18 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
     k.emit(ds_read_b32(v[201], v[224], v[0], v[0], 0, 0, 1))
     k.emit(ds_read_b32(v[202], v[224], v[0], v[0], 0, 0, 2))
     k.emit(ds_read_b32(v[203], v[224], v[0], v[0], 0, 0, 3))
-    k.emit(s_lshl_b32(s[40], s[40], 1))
+    k.emit(s_lshl_b32(s[36], s[36], 1))
     k.emit(s_mul_i32(s[52], s[47], LIT, 256))
-    k.emit(s_mul_hi_u32(s[53], s[52], s[40]))
+    k.emit(s_mul_hi_u32(s[53], s[52], s[36]))
     k.emit(s_add_u32(s[5], s[5], s[53]))
-    k.emit(s_mul_i32(s[53], s[52], s[40]))
+    k.emit(s_mul_i32(s[53], s[52], s[36]))
     k.emit(s_add_u32(s[4], s[4], s[53]))
     k.emit(s_addc_u32(s[5], 0, s[5]))
     k.emit(s_sub_i32(s[52], s[43], s[52]))
-    k.emit(s_mul_i32(s[52], s[52], s[40]))
+    k.emit(s_mul_i32(s[52], s[52], s[36]))
     k.emit(s_mov_b32(s[6], s[52]))
     k.emit(v_and_b32_e64(v[235], v[0], 15))
-    k.emit(v_mul_lo_u32(v[235], v[235], s[40]))
+    k.emit(v_mul_lo_u32(v[235], v[235], s[36]))
     k.emit(v_lshrrev_b32_e32(v[4], 5))
     k.emit(v_mul_i32_i24_e32(v[4], 16, v[4]))
     k.emit(v_add_u32_e32(v[235], v[4], v[235]))
@@ -2529,12 +2482,12 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
     k.emit(v_and_b32_e32(v[4], 1, v[4]))
     k.emit(v_mul_i32_i24_e32(v[4], 32, v[4]))
     k.emit(v_add_u32_e32(v[235], v[4], v[235]))
-    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[49], 64))
+    k.emit(s_mul_i32(s[52], s[49], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[46], 64))
     k.emit(s_add_i32(s[52], s[52], s[53]))
     k.emit(s_lshl_b32(s[52], s[52], 1))
     k.emit(v_add_u32_e32(v[235], s[52], v[235]))
-    k.emit(s_mul_i32(s[53], s[40], 16))
+    k.emit(s_mul_i32(s[53], s[36], 16))
     for i in range(15):
       k.emit(v_add_u32_e64(v[236 + i * 1], v[235 + i * 1], s[53]))
     k.emit(s_mov_b32(s[50], 0))
@@ -2543,7 +2496,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
       k.emit(s_cmp_lt_u32(LIT, s[51], 512 + i * -256))
       k.emit(s_cselect_b32(s[61 + i * 1], s[61 + i * 1], 0))
       k.emit(s_cselect_b32(s[63 + i * 1], s[63 + i * 1], 0))
-    k.emit(s_cmp_lt_i32(s[49], 2))
+    k.emit(s_cmp_lt_i32(s[46], 2))
     k.emit(s_cbranch_scc0(1367), target='L2_25B8')
     k.label('L2_105C')
     k.emit(s_waitcnt(122))

extra/llama_kernels/__init__.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 import functools, pathlib
-from dataclasses import replace
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import shape_to_shape_arg
+from tinygrad import Tensor
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 
 FP8_MAX = 448.0
@@ -12,7 +10,7 @@ NUM_WG, THREADS_PER_WG = 1024, 256
 @functools.cache
 def _local_abs_max_fxn(x_p, device):
   x = Tensor(x_p, device=device)
-  inner = Tensor(x.uop.replace(src=(shape_to_shape_arg(x.uop.shard_shape),), arg=replace(x.uop.arg, axis=None))) if x.uop.axis is not None else x
+  inner = Tensor(x.uop.src[0]) if x.uop.axis is not None else x  # the per-shard view of the flat param
   return (inner.abs().max(),)
 
 def local_abs_max(x:Tensor) -> Tensor:

extra/onnx_helpers.py

@@ -1,10 +1,12 @@
 from tinygrad import Tensor
+from tinygrad.helpers import NUM_CPU_THREADS
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
 ort_options = ort.SessionOptions()
 ort_options.log_severity_level = 3
+ort_options.intra_op_num_threads = NUM_CPU_THREADS.value
 
 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
   """

spec/tinyspec.tex

@@ -50,10 +50,10 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Param} & $(\mathbf{s})$ & slot, dtype, device?, addrspace? &
-  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
-\op{Buffer} & $(\mathbf{s})$ & slot, dtype, device, addrspace &
-  Concrete buffer slot with shape $\mathbf{s}$. If device is a tuple, it creates the fully sized buffer across multiple devices. \\
+\op{Param} & () & slot, dtype, size?, device?, addrspace? &
+  Placeholder with flat storage of $\mathrm{size}$ elements. Substituted in \op{Call}. \\[4pt]
+\op{Buffer} & () & slot, dtype, size, device, addrspace &
+  Concrete buffer slot with flat storage of $\mathrm{size}$ elements. \\
 \op{Const} & () & value, dtype &
   A scalar constant with shape $(\ )$. \\
   & & & Form vector consts with \op{Stack} \\
@@ -256,9 +256,9 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
 \toprule
 \textbf{Op} & \textbf{dtype} & \textbf{shape} & \textbf{device} & \textbf{min\_max} \\
 \midrule
-\op{Buffer} & from arg & from $\mathrm{src}[0]$ & from arg & dtype range \\
+\op{Buffer} & from arg & from arg ($\mathrm{size}$) & from arg & dtype range \\
 \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
-\op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
+\op{Param} & from arg & from arg ($\mathrm{size}$) & from arg & from src or dtype range \\[3pt]
 Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]

test/amd/test_emu2_pcode.py

@@ -152,7 +152,7 @@ class TestDSPcodePatterns(unittest.TestCase):
 
   def test_global_atomic_add_f32_parsing(self):
     """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
-    vmem = UOp.param(2, dtypes.uint32, (1024,))
+    vmem = UOp.param(2, dtypes.uint32, 1024)
     srcs = {
       'ADDR': UOp.const(0, dtypes.uint64),
       'DATA': UOp.const(0x3f800000, dtypes.uint32),
@@ -183,7 +183,7 @@ class TestDSPcodePatterns(unittest.TestCase):
   def test_mem_read_parsing(self):
     """Test MEM[addr].type read expression parsing."""
     # Create a mock LDS buffer
-    lds = UOp.param(3, dtypes.uint32, (16384,))
+    lds = UOp.param(3, dtypes.uint32, 16384)
     addr = UOp.const(0, dtypes.uint32)
     vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)}
 
@@ -218,7 +218,7 @@ class TestDSPcodePatterns(unittest.TestCase):
     pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
     self.assertIsNotNone(pcode)
     assert pcode is not None
-    lds = UOp.param(3, dtypes.uint32, (16384,))
+    lds = UOp.param(3, dtypes.uint32, 16384)
     srcs = {
       'ADDR': UOp.const(0, dtypes.uint32),
       'OFFSET0': UOp.const(0, dtypes.uint32),
@@ -299,7 +299,7 @@ class TestConcatWidthParsing(unittest.TestCase):
       self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
 
   def test_permlane64_wave64_pcode_indices(self):
-    vgpr = UOp.param(0, dtypes.uint32, (256,))
+    vgpr = UOp.param(0, dtypes.uint32, 256)
     srcs = {
       'SRC0': UOp.const(0, dtypes.uint32),
       'VDST': UOp.const(1, dtypes.uint32),
@@ -330,7 +330,7 @@ class TestAllPcode(unittest.TestCase):
   def _make_srcs(self):
     """Create dummy source variables for pcode parsing."""
     u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64)
-    lds = UOp.param(3, dtypes.uint32, (16384,))
+    lds = UOp.param(3, dtypes.uint32, 16384)
     return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
             'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
             'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),

test/backend/test_dtype.py

@@ -258,6 +258,11 @@ class TestDoubleDType(TestDType):
       a = [2, 3, 4]
       np.testing.assert_allclose(func(Tensor(a, dtype=self.DTYPE)).numpy(), func(torch.tensor(a, dtype=torch.float64)), rtol=1e-12, atol=1e-12)
 
+  def test_float32_compare_selecting_float64(self):
+    a = Tensor([1.0, 2.0, 5.0, 9.0], dtype=dtypes.float32)
+    p, q = Tensor([10., 20., 30., 40.], dtype=self.DTYPE), Tensor([50., 60., 70., 80.], dtype=self.DTYPE)
+    _test_op(lambda: (a < 3.0).where(p, q), self.DTYPE, [10., 20., 70., 80.])
+
   def test_float64_to_float32_cast_inf(self):
     _test_op(lambda: Tensor([3.4e40, 3.4e38, 1, 0], dtype=dtypes.float64).cast(dtypes.float32),
              dtypes.float32, [float('inf'), 3.4e38, 1, 0])
@@ -323,7 +328,11 @@ class TestUint16DType(TestDType):
 class TestInt32DType(TestDType): DTYPE = dtypes.int32
 class TestUint32DType(TestDType): DTYPE = dtypes.uint32
 
-class TestInt64DType(TestDType): DTYPE = dtypes.int64
+class TestInt64DType(TestDType):
+  DTYPE = dtypes.int64
+  # the high bits a narrowing cast drops can't come back through a widening cast
+  def test_int64_to_uint32_to_int64(self):
+    _test_op(lambda: Tensor([0x12345678ABCDEF01], dtype=dtypes.int64).cast(dtypes.uint32).cast(dtypes.int64), dtypes.int64, [2882400001])
 
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedInt64DType(TestInt64DType):
@@ -423,6 +432,11 @@ class TestDtypeUsage(unittest.TestCase):
         t = Tensor([[1, 2], [3, 4]], dtype=d)
         (t*t).max().item()
 
+  def test_where_float16_compare_to_const(self):
+    # t > 0 is CMPLT(0, t): the float16 operand is on the right
+    t = Tensor([-1.0, 1.0], dtype=dtypes.float16)
+    np.testing.assert_equal((t > 0).where(Tensor.ones(2, dtype=dtypes.float16), Tensor.zeros(2, dtype=dtypes.float16)).numpy(), [0.0, 1.0])
+
 @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
   def test_cast(self):

test/backend/test_encodings.py

@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops
 from tinygrad.dtype import dtypes
 from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
 
-def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)
+def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=(op, dt), src=src, tag=tag)
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86")
 class TestEncodingsX86(unittest.TestCase):
@@ -100,13 +100,6 @@ class TestEncodingsX86(unittest.TestCase):
     # vaddss xmm0, xmm0, xmm8
     self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0"))
 
-  # test ymm encoding
-  def test_ymm_encoding(self):
-    xmm0, xmm1 = def_reg(dtypes._uint256, XMM[0]), def_reg(dtypes._uint256, XMM[1])
-    add = ins(X86Ops.VADDPS, dtypes._uint256, (xmm0, xmm1), XMM[0])
-    # vaddps ymm0, ymm0, ymm1
-    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FC 58 C1"))
-
   # test encoding where register is in the immediate field
   def test_reg_in_imm_field(self):
     xmm0, xmm1, xmm2 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]), def_reg(dtypes.float32, XMM[2])
@@ -143,9 +136,9 @@ class TestEncodingsX86(unittest.TestCase):
 
   # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
   def test_cmove_ignore_cmp(self):
-    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX)
+    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=(X86Ops.CMP, dtypes.void))), RDX)
     # cmove edx, eax
     self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
 
 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()

test/backend/test_isel.py

@@ -20,7 +20,7 @@ class TestIselX86(unittest.TestCase):
       with self.subTest(dtype=dt):
         v = [UOp.variable(str(i), 0, 0, dt) for i in range(nargs)]
         n = self.isel_rewrite(expr(*v))
-        self.assertIs(n.arg, op)
+        self.assertIs(n.arg[0], op)
 
   def test_cmove(self):
     a = UOp.variable("a", 0, 0, dtypes.int32)
@@ -29,9 +29,9 @@ class TestIselX86(unittest.TestCase):
     d = (a != b).where(a, b)
     f = c + d
     n = self.isel_rewrite(f)
-    self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg is X86Ops.CMOVNE)
+    self.assertTrue(n.src[0].arg[0] is X86Ops.CMOVL and n.src[1].arg[0] is X86Ops.CMOVNE)
     # both comparisons become the same instruction
-    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP)
+    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg[0] is X86Ops.CMP)
 
   def test_vinsertps(self):
     a = UOp.variable("a", 0, 0, dtypes.float32)
@@ -41,12 +41,12 @@ class TestIselX86(unittest.TestCase):
 
     valid = [UOp.stack(lane(a, 0), lane(b, 1), lane(a, 2), lane(b, 3)),
              UOp.stack(lane(a, 3), lane(b, 2), lane(c, 1), d)]
-    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg[0], X86Ops.VINSERTPS)
 
   # 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 + UOp.cconst(1, dtypes.int32)).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].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)

test/backend/test_jit.py

@@ -6,7 +6,7 @@ from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_
 from test.unit.test_jit import _simple_test
 from tinygrad import Tensor, Variable, TinyJit, Device, dtypes
 from tinygrad.engine.jit import graph_class
-from tinygrad.helpers import JIT, DEV, GlobalCounters
+from tinygrad.helpers import JIT, DEV, GlobalCounters, HCQ2
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer.isa.x86 import X86Renderer
 
@@ -235,6 +235,7 @@ class TestJitPrune(unittest.TestCase):
     assert_jit_cache_len(w2_prune, 1)
 
 class TestJitFree(unittest.TestCase):
+  @unittest.skipIf(HCQ2, "hcq2 keeps refs to intermediate buffers")
   def test_free_intermediates(self):
     ext_tensor = Tensor([1,24,23,45,1])
     @TinyJit

test/backend/test_linearizer_dumb.py

@@ -11,16 +11,16 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailure(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
   def test_failure_beam_mnist(self):
-    c0 = UOp.param(0, dtypes.uchar, (4014080,))
+    c0 = UOp.param(0, dtypes.uchar, 4014080)
     c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL)
     c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL)
     c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL)
-    c4 = UOp.param(1, dtypes.int, (512,))
+    c4 = UOp.param(1, dtypes.int, 512)
     c5 = c4.index(c1.valid(UOp.const(True)))
     c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE)
     c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE)
     c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE)
-    c9 = UOp.param(2, dtypes.uchar, (47040000,))
+    c9 = UOp.param(2, dtypes.uchar, 47040000)
     c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True)))
     c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD)
     c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3)

test/backend/test_multitensor.py

@@ -187,6 +187,13 @@ class TestMultiTensor(unittest.TestCase):
         a,b = jit_allreduce(Tensor.rand(256, 256))
         np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
+  def test_allreduce_all2all_jit(self):
+    with Context(ALL2ALL=2):
+      jit_allreduce = TinyJit(_test_allreduce)
+      for _ in range(5):
+        a,b = jit_allreduce(Tensor.rand(256, 256))
+        np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
+
   def test_multitensor_jit_input(self):
     @TinyJit
     def f(x): return (x+1).contiguous().sum()

test/backend/test_nn.py

@@ -135,7 +135,7 @@ class TestNN(unittest.TestCase):
   def test_conv2d_same_padding_large_kernel(self):
     self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[28, 33], C2=32, K=9, S=1, P='same')
   def test_conv2d_same_padding_with_dilation(self):
-    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=3)
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 31], C2=32, K=(3,5), S=1, P='same', D=(2,3))
 
   def test_conv2d_same_padding_invalid_stride(self):
     self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same')

test/backend/test_ops.py

@@ -359,6 +359,13 @@ class TestOps(unittest.TestCase):
       lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32).permute((1, 0)),
       lambda x: (x > 0.5).where(4, 2).clone().permute((1, 0)), forward_only=True)
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "software vulkan evaluates a NaN != x as false")
+  def test_where_nan_cond(self):
+    # a NaN compares false against everything except !=.
+    for fxn in (lambda x: x<1, lambda x: x>1, lambda x: x!=1, lambda x: x==1):
+      helper_test_op(None, lambda x,a,b: torch.where(fxn(x), a, b), lambda x,a,b: fxn(x).where(a, b), forward_only=True,
+                     vals=[[math.nan, 1.0, 2.0, -1.0], [10, 20, 30, 40], [-1, -2, -3, -4]])
+
   def _test_cmp(self, fxn, reverse=True):
     # test different dtypes
     helper_test_op(None, fxn, fxn, forward_only=True, vals=[[0.,1,2], [2.,1,0]])
@@ -2810,7 +2817,7 @@ class TestOps(unittest.TestCase):
         lambda x: Tensor.interpolate(x, size=out_sz, mode="linear"))
 
   def test_interpolate_linear_corners_aligned(self):
-    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,))]:
+    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,)), ((29,),(1,))]:
       helper_test_op([(2,3)+in_sz],
         lambda x: torch.nn.functional.interpolate(x, size=out_sz, mode="linear", align_corners=True),
         lambda x: Tensor.interpolate(x, size=out_sz, mode="linear", align_corners=True))
@@ -2963,6 +2970,10 @@ class TestOps(unittest.TestCase):
     data = [math.inf, -math.inf, math.nan]
     helper_test_op((), lambda: torch.tensor(data)[torch.tensor([0, 1, 2])], lambda: Tensor(data)[Tensor([0, 1, 2])])
 
+  def test_fancy_indexing_index_dtypes(self):
+    helper_test_op((), lambda: torch.tensor([10., 20., 30., 40.])[torch.tensor([1, 2, 3, 0])],
+                       lambda: Tensor([10., 20., 30., 40.])[Tensor([1, 2, 3, 0], dtype=dtypes.uint8)])
+
   @slow_test
   def test_slice_fancy_indexing_no_dim_collapse(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()

test/backend/test_renderer_failures.py

@@ -23,8 +23,8 @@ def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):
 
 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   dtype = alu_src_uops[0].dtype
-  a = UOp.param(0, dtype, (1,))
-  b = UOp.param(1, dtype, (1,))
+  a = UOp.param(0, dtype, 1)
+  b = UOp.param(1, dtype, 1)
   idx = UOp.const(0)
   ld = b.index(idx).load()
   alu = ld.alu(alu_op, *alu_src_uops)
@@ -34,7 +34,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
 class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
   def test_gated_store_with_alu(self):
-    a = UOp.param(0, dtypes.int, (4,))
+    a = UOp.param(0, dtypes.int, 4)
     gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
     gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int)))
     sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
@@ -43,7 +43,7 @@ class TestRendererFailures(unittest.TestCase):
 
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
   def test_gated_store_with_alu_2d(self):
-    a = UOp.param(0, dtypes.int, (8,))
+    a = UOp.param(0, dtypes.int, 8)
     gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
     gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0)
     gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int)))
@@ -78,7 +78,7 @@ class TestCStyleFailures(unittest.TestCase):
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
 class TestWGSLFailures(unittest.TestCase):
   def test_folded_packed_store(self):
-    b = UOp.param(0, dtypes.char, (4,))
+    b = UOp.param(0, dtypes.char, 4)
     idx = b.index(UOp.const(0).cast(dtypes.int))
     store = UOp.store(idx, idx.cast(dtypes.uint32).load() & UOp.const(0xffffff00).cast(dtypes.uint32))
     src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
@@ -93,9 +93,9 @@ class TestWGSLFailures(unittest.TestCase):
 
   # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
   def test_gated_load(self):
-    a = UOp.param(0, dtypes.int, (4,))
-    b = UOp.param(1, dtypes.int, (4,))
-    c = UOp.param(2, dtypes.int, (4,))
+    a = UOp.param(0, dtypes.int, 4)
+    b = UOp.param(1, dtypes.int, 4)
+    c = UOp.param(2, dtypes.int, 4)
     lidx0 = UOp.special(4, "lidx0")
     gate = lidx0.ne(0)
     alt = c.index(lidx0).load()
@@ -110,7 +110,7 @@ class TestWGSLFailures(unittest.TestCase):
 class TestPTXFailures(unittest.TestCase):
   @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
   def test_gated_store_with_if(self):
-    a = UOp.param(0, dtypes.int, (4,))
+    a = UOp.param(0, dtypes.int, 4)
     gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
     val = UOp.const(1).cast(dtypes.int)
     if_uop = UOp(Ops.IF, src=(gate_alu,))

test/backend/test_schedule.py

@@ -379,27 +379,30 @@ class TestCopyFolding(unittest.TestCase):
     check_schedule(a.clone(), 1, filter_sink=False)
 
   def test_shrink_copy(self):
-    a = Tensor.arange(4)
-    view = a.shrink(((0, 2),))
-    b = view.clone()
-    run_linear(*check_schedule(b, 1, filter_sink=False))
-    self.assertEqual(b.uop.base.buffer.size, 2)
-    self.assertEqual(b.uop.numel(), 2)
-    self.assertListEqual(b.tolist(), [0, 1])
+    a = Tensor.arange(4).clone("CPU:1").realize()
+    b = a.to("CPU:2").shrink(((1, 3),)).to("CPU:3")
+    GlobalCounters.reset()
+    run_linear(*check_schedule(b, 3, filter_sink=False))
+    # extra E kernel, copy exactly 4 bytes
+    self.assertEqual(GlobalCounters.global_mem, 4*4 + 2*4*2 + 2*4)
+    self.assertListEqual(b.tolist(), [1, 2])
 
   def test_expanded_copy(self):
-    a = Tensor.arange(2)
-    view = a.reshape(2, 1).expand(2, 2)
-    b = view.clone()
-    run_linear(*check_schedule(b, 1, filter_sink=False))
-    self.assertEqual(b.uop.base.buffer.size, 4)
-    self.assertEqual(b.uop.numel(), 4)
-    self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])
+    a = Tensor.arange(4).clone("CPU:1").realize()
+    b = a.to("CPU:2").reshape(4, 1).expand(4, 2).to("CPU:3")
+    GlobalCounters.reset()
+    run_linear(*check_schedule(b, 3, filter_sink=False))
+    # TODO: expands before copy
+    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 8*4) + 8*4)
+    self.assertListEqual(b.tolist(), [[0, 0], [1, 1], [2, 2], [3, 3]])
 
   def test_permuted_copy(self):
-    a = Tensor.arange(4)
-    b = a.reshape(2, 2).permute(1, 0)
-    b.realize()
+    a = Tensor.arange(4).clone("CPU:1").realize()
+    b = a.to("CPU:2").reshape(2, 2).permute(1, 0).to("CPU:3")
+    GlobalCounters.reset()
+    run_linear(*check_schedule(b, 3, filter_sink=False))
+    # permutes before copy
+    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 4*4) + 4*4)
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
   def test_permute_on_disk(self):

test/backend/test_uops.py

@@ -20,8 +20,8 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
 
 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
   if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype))
-  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
-  else: uops.append(UOp(op, dtype, tuple(src), arg))
+  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, 1))
+  else: uops.append(UOp(op, tuple(src), arg))
   return uops[-1]
 
 def _test_single_value(vals, op, dts):
@@ -62,7 +62,7 @@ def _test_uops_result(output_dtype, uops, res):
 class TestBitcastBufferView(unittest.TestCase):
   @Context(SPEC=2)
   def test_render(self):
-    buf = UOp.param(0, dtypes.uint32, (4,))
+    buf = UOp.param(0, dtypes.uint32, 4)
     uops = to_uops_list([buf.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0).store(1)], ren=Device[Device.DEFAULT].renderer)
     idx = next(u for u in uops if u.op is Ops.INDEX and u.src[0].op is Ops.BITCAST)
     self.assertEqual(idx.src[0].src[0].op, Ops.SHRINK)
@@ -71,7 +71,7 @@ class TestBitcastBufferView(unittest.TestCase):
   @Context(SPEC=2)
   def test_load(self):
     val = 0x1122334455667788
-    src, out = UOp.param(0, dtypes.uint32, (4,)), UOp.param(1, dtypes.uint64, (1,))
+    src, out = UOp.param(0, dtypes.uint32, 4), UOp.param(1, dtypes.uint64, 1)
     ibuf = Buffer(Device.DEFAULT, 4, dtypes.uint32, initial_value=np.array([0, 0x55667788, 0x11223344, 0], dtype=np.uint32).tobytes())
     obuf = Buffer(Device.DEFAULT, 1, dtypes.uint64).allocate()
     run_uops([out.index(0).store(src.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0))], [ibuf, obuf])
@@ -80,7 +80,7 @@ class TestBitcastBufferView(unittest.TestCase):
   @Context(SPEC=2)
   def test_store(self):
     val = 0x1122334455667788
-    dst = UOp.param(0, dtypes.uint32, (6,))
+    dst = UOp.param(0, dtypes.uint32, 6)
     buf = Buffer(Device.DEFAULT, 6, dtypes.uint32, initial_value=bytes(24))
     view = dst.shrink(((1, 5),)).bitcast(dtypes.uint64)  # two stores through one view: it must inline, not get a declared vector-pointer
     run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
@@ -249,8 +249,8 @@ class TestLocalAccess(unittest.TestCase):
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
 class TestAssembly(unittest.TestCase):
   def test_bitshift_left(self):
-    g1 = UOp.param(0, dtypes.int32, shape=(3,))
-    out = UOp.param(1, dtypes.int32, shape=(2,))
+    g1 = UOp.param(0, dtypes.int32, 3)
+    out = UOp.param(1, dtypes.int32, 2)
     c1 = UOp.const(2)
     c2 = UOp.const(3)
     l1 = g1.index(c1)
@@ -278,7 +278,7 @@ class TestAssembly(unittest.TestCase):
     self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)
 
   def test_mulacc_shl(self):
-    g1 = UOp.param(0, dtypes.int32, shape=(2,))
+    g1 = UOp.param(0, dtypes.int32, 2)
     c1 = UOp.const(0)
     c2 = UOp.const(1)
     expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
@@ -287,7 +287,7 @@ class TestAssembly(unittest.TestCase):
     self.assertIn(Ops.MULACC, [x.op for x in uops])
 
   def test_use_cmpeq(self):
-    g = UOp.param(0, dtypes.uint32, shape=(8,))
+    g = UOp.param(0, dtypes.uint32, 8)
     c = UOp.const(7)
     comp = g.index(c).ne(c).ne(True)
     uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)

test/backend/test_wait_loop.py

@@ -43,7 +43,7 @@ def nested_loop_kernel(C:UOp) -> UOp:
   return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))
 
 def wait_ext_kernel() -> UOp:
-  sig = UOp.param(0, dtypes.int, (1,), volatile=True)
+  sig = UOp.param(0, dtypes.int, 1, volatile=True)
   l = UOp.loop(0)
   v = sig.after(l)[0].load()
   e = v.end(l, v < 1)

test/device/test_hcq2.py

@@ -3,10 +3,10 @@ from unittest.mock import patch
 from tinygrad import Device, Tensor
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import getenv
+from tinygrad.helpers import HCQ2
 from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in
 
-@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
+@unittest.skipUnless(HCQ2 and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
 class TestHCQ2(unittest.TestCase):
   def test_copy_without_copy_queue(self):
     with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):

test/external/external_benchmark_op_conv.py

@@ -28,7 +28,7 @@ def vision_conv_143():
   c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
   c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, (128,))
+  c63 = UOp.param(3, dtypes.float, 128)
   c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
   c67 = c0.index((c2*128+c5+c8*4096)).store(c65).end(c8, c2, c5)
 
@@ -54,7 +54,7 @@ def vision_conv_153():
   c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
   c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, (256,))
+  c63 = UOp.param(3, dtypes.float, 256)
   c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
   c67 = c0.index((c2*256+c5+c8*4096)).store(c65).end(c8, c2, c5)
 
@@ -73,7 +73,7 @@ def dm_conv_172():
   c18 = UOp.range(8, 2, AxisType.REDUCE)
   c23 = UOp.param(2, dtypes.half, shape=(240, 128, 4))
   c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp.param(3, dtypes.float, (960,))
+  c37 = UOp.param(3, dtypes.float, 960)
   c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
   c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
   c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666

test/external/external_benchmark_pyrender.py

@@ -7,8 +7,8 @@ BENCHMARK_OPS = {Ops.INDEX, Ops.STAGE}
 
 @functools.cache
 def create_uop(a:int) -> UOp:
-  op, dtype, src, arg, *rest = trace.uop_fields[a]
-  return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)
+  op, src, arg, *rest = trace.uop_fields[a]
+  return UOp(op, tuple(create_uop(s) for s in src), arg, *rest)
 
 if __name__ == "__main__":
   # load rewrite trace

test/external/external_test_gpu_crash.py

@@ -36,7 +36,7 @@ class TestGPUCrash(unittest.TestCase):
 
   def _run_insts(self, insts: list[Inst]):
     buf = UOp.new_buffer("AMD", 64, dtypes.uint8)
-    sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
+    sink = UOp.sink(UOp.param(0, dtypes.uint8, 64, device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
     prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts))))
     run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)
 

test/external/external_test_onnx_ops.py

@@ -54,6 +54,12 @@ class TestMainOnnxOps(TestOnnxOps):
     outputs = ["squeezed"]
     self.helper_test_single_op("Squeeze", inputs, attributes, outputs)
 
+  def test_mean_variance_normalization_axes(self):
+    inputs = {"x": np.random.randn(2, 3, 4, 5).astype(np.float32)}
+    attributes = {"axes": [2, 3]}
+    outputs = ["out"]
+    self.helper_test_single_op("MeanVarianceNormalization", inputs, attributes, outputs)
+
   def test_conv(self):
     # test VALID auto_pad
     inputs = {
@@ -235,6 +241,15 @@ class TestMainOnnxOps(TestOnnxOps):
     outputs = ["y"]
     self.helper_test_single_op("MaxUnpool", inputs, attributes, outputs)
 
+  def test_maxunpool_pads(self):
+    # per-axis pads shrink the output: spatial dim is (i-1)*stride + kernel - pad_begin - pad_end -> (2, 4), and indices index into that output
+    # NOTE: indices must be in bounds of that output; ORT aborts the process on out-of-bounds indices
+    xT = np.array([[[[5, 6], [7, 8]]]], dtype=np.float32)
+    xI = np.array([[[[0, 3], [4, 7]]]], dtype=np.int64)
+    inputs = {"x": xT, "indices": xI}
+    attributes = {"kernel_shape": [2, 2], "strides": [2, 2], "pads": [1, 0, 1, 0]}
+    self.helper_test_single_op("MaxUnpool", inputs, attributes, ["y"])
+
   def test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True(self):
     # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-13
     inputs = {"x": np.random.randn(1, 1, 32, 32, 32).astype(np.float32)}

test/helpers.py

@@ -122,7 +122,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
   for buf_dt, data in inputs or []:
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
-  g = UOp.param(0, uop.dtype, (1,))
+  g = UOp.param(0, uop.dtype, 1)
   prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
   prog = dev.runtime(prg.to_elf())
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)

test/mockgpu/amd/emu.py

@@ -321,10 +321,10 @@ def _int_clamp(op_name: str, srcs: dict) -> UOp | None:
 class _Ctx:
   """Context for instruction compilation - holds buffers and helpers."""
   __slots__ = ('inst_size', 'dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr')
-  sgpr = UOp.param(0, dtypes.uint32, (SGPR_COUNT,))
-  vmem = UOp.param(2, dtypes.uint32, (1 << 46,))
-  lds = UOp.param(3, dtypes.uint32, (16384,))
-  scratch = UOp.param(4, dtypes.uint8, (1 << 30,))
+  sgpr = UOp.param(0, dtypes.uint32, SGPR_COUNT)
+  vmem = UOp.param(2, dtypes.uint32, 1 << 46)
+  lds = UOp.param(3, dtypes.uint32, 16384)
+  scratch = UOp.param(4, dtypes.uint8, 1 << 30)
   # Cache PARAM UOps by wave_size so all _Ctx instances with same wave_size share identical UOp references
   _vgpr_cache: dict[int, UOp] = {}
   _accvgpr_cache: dict[int, UOp] = {}
@@ -332,10 +332,10 @@ class _Ctx:
   def __init__(self, inst_size: int, wave_size: int = 32):
     self.inst_size, self._axis_id, self.wave_size = inst_size, 0, wave_size
     self.dyn_fields: list[tuple[int, int]] = []  # (lo, hi) of fields read dynamically
-    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, (256 * wave_size,))
+    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, 256 * wave_size)
     self.vgpr = _Ctx._vgpr_cache[wave_size]
     if wave_size == 64:
-      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, (256 * wave_size,))
+      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, 256 * wave_size)
       self.accvgpr = _Ctx._accvgpr_cache[wave_size]
     else:
       self.accvgpr = self.vgpr

test/null/test_gpudims.py

@@ -107,7 +107,7 @@ class TestGroupedDims(unittest.TestCase):
 
   def test_global_prod_max(self):
     g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
-    sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
+    sink = UOp.param(0, dtypes.float, 512).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
     class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
     specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
     self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)

test/null/test_linearizer_failures.py

@@ -7,14 +7,14 @@ from tinygrad.codegen import to_program
 
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
-    c0 = UOp.param(0, dtypes.float, (64,))
+    c0 = UOp.param(0, dtypes.float, 64)
     c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK)
     c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK)
     c3 = ((c1*UOp.const(32))+c2)
-    c4 = UOp.param(1, dtypes.float, (163840,))
+    c4 = UOp.param(1, dtypes.float, 163840)
     c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE)
     c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920))))
-    c7 = UOp.param(2, dtypes.float, (64,))
+    c7 = UOp.param(2, dtypes.float, 64)
     c8 = c7.index(c3)
     c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)

test/null/test_pattern_matcher.py

@@ -40,7 +40,7 @@ class TestPatternMatcher(unittest.TestCase):
       assert len(x.src) == 0
       return UOp(Ops.CONST, src=(UOp(Ops.CONST),))
     matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
-    c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
+    c1 = UOp(Ops.CONST, arg=1.0)
     # second rewrite shouldn't match anything
     c1 = matcher.rewrite(c1)
     c1 = matcher.rewrite(c1)

test/null/test_simplify_valid_idx.py

@@ -15,7 +15,7 @@ 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.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load()
+  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.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load()
@@ -505,7 +505,7 @@ class TestDropTrueGate(unittest.TestCase):
     from tinygrad.codegen.late.coalesce import indexing_simplify
     from tinygrad.uop.ops import graph_rewrite
     from tinygrad.uop.symbolic import sym
-    buf = UOp.param(0, dtypes.int, (1,))
+    buf = UOp.param(0, dtypes.int, 1)
     idx = UOp.const(0)
     true_gate = UOp.const(True)
     index_with_gate = buf.index(idx.valid(true_gate))
@@ -551,7 +551,7 @@ class TestRangeShrink(unittest.TestCase):
     # 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.param(1, dtypes.float, (204,)).index(r).load()
+    load2 = UOp.param(1, dtypes.float, 204).index(r).load()
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assert_range_end(ranges, 204)
 
@@ -575,7 +575,7 @@ class TestRangeShrink(unittest.TestCase):
     from tinygrad.dtype import Invalid
     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())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r < 4).where(x, Invalid)).sink())
     self.assert_range_end(ranges, 4)
 
   def test_range_shrink_store_where_invalid_flipped(self):
@@ -583,7 +583,7 @@ class TestRangeShrink(unittest.TestCase):
     from tinygrad.dtype import Invalid
     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())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r >= 4).where(Invalid, x)).sink())
     self.assert_range_end(ranges, 4)
 
 if __name__ == '__main__':

test/null/test_transcendental_helpers.py

@@ -9,7 +9,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
   def test_payne_hanek_reduction(self):
     # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
     # Load input value from a buffer to prevent constant folding
-    input_buf = UOp.param(1, dtypes.double, (1,))
+    input_buf = UOp.param(1, dtypes.double, 1)
     loaded_value = input_buf.index(UOp.const(0)).load()
     def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
       return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))

test/null/test_uop_graph.py

@@ -214,8 +214,8 @@ class TestUOpGraph(unittest.TestCase):
       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,))
-    d1 = UOp.param(1, dtypes.int, (1,))
+    d0 = UOp.param(0, dtypes.bool, 1)
+    d1 = UOp.param(1, dtypes.int, 1)
     idx = UOp.const(0)
     ld = d1.index(idx)
     alu = (ld<1).cast(dtypes.bool)
@@ -224,8 +224,8 @@ class TestUOpGraph(unittest.TestCase):
     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,))
-    d1 = UOp.param(1, dtypes.int, (1,))
+    d0 = UOp.param(0, dtypes.float, 1)
+    d1 = UOp.param(1, dtypes.int, 1)
     idx = UOp.const(0, dtypes.int)
     ld = d1.index(idx)
     alu = ld.cast(dtypes.float).cast(dtypes.float)
@@ -243,7 +243,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_bitcast_to_same_dtype_fold(self):
     for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
-      d0 = UOp.param(0, dt, (1,))
+      d0 = UOp.param(0, dt, 1)
       v = d0.index(UOp.const(0))
       uops = to_uops_list([v.bitcast(dt)])
       self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
@@ -255,10 +255,10 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_gated_load_fold(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, (100,))
+    d0 = UOp.param(0, dtypes.long, 100)
     ld = d0.index(ridx0.valid(ridx0<50))
     w = (ridx0<50).where(ld, 5)
-    out = UOp.param(1, dtypes.long, (100,))
+    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:
@@ -267,7 +267,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, (100,))
+    d0 = UOp.param(0, dtypes.long, 100)
     ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
     w = (ridx0<50).where(5, ld)
     uops = to_uops_list([w])
@@ -278,11 +278,11 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_gated_load_with_cast(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.int, (100,))
+    d0 = UOp.param(0, dtypes.int, 100)
     gate_idx = ridx0.valid((ridx0<50))
     ld = d0.index(gate_idx).cast(dtypes.float)
     w = (ridx0<50).where(ld, 5.0)
-    out = UOp.param(1, dtypes.float, (100,))
+    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:
@@ -291,27 +291,27 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_on_casted_gated_load_extra_cond(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, (100,))
+    d0 = UOp.param(0, dtypes.float, 100)
     ld = d0.index(ridx0.valid(ridx0<50))
     w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, (100,))
+    out = UOp.param(1, dtypes.half, 100)
     uops = to_uops_list([out.index(ridx0).store(w)])
     for u in uops:
       assert u.op is not Ops.WHERE
 
   def test_where_on_casted_gated_load_extra_cond_swapped(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, (100,))
+    d0 = UOp.param(0, dtypes.float, 100)
     ld = d0.index(ridx0.valid(ridx0<50))
     w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, (100,))
+    out = UOp.param(1, dtypes.half, 100)
     uops = to_uops_list([out.index(ridx0).store(w)])
     for u in uops:
       assert u.op is not Ops.WHERE
 
   def test_where_in_store_becomes_gate(self):
     ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, (100,))
+    d0 = UOp.param(0, dtypes.long, 100)
     idx = d0.index(ridx0)
     ld = idx.load()
     val = (ridx0<50).where(5, ld)
@@ -325,14 +325,14 @@ class TestUOpGraph(unittest.TestCase):
   def test_load_idx_becomes_int(self):
     # mnist indexing with split reduceop
     # Make sure we are not doign math on the loaded index, which would promote it to long
-    c0 = UOp.param(0, dtypes.uchar, (128000,))
+    c0 = UOp.param(0, dtypes.uchar, 128000)
     c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
     c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
-    c3 = UOp.param(1, dtypes.int, (512,))
+    c3 = UOp.param(1, dtypes.int, 512)
     c4 = c3.index(c1)
     c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
     c6 = ((c2*UOp.const(240))+c5)
-    c7 = UOp.param(2, dtypes.uchar, (60000,))
+    c7 = UOp.param(2, dtypes.uchar, 60000)
     c8 = c7.index(c6)
     c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
     c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
@@ -342,14 +342,14 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_load_idx_no_math_on_loaded(self):
     # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
-    c0 = UOp.param(0, dtypes.uchar, (128000,))
+    c0 = UOp.param(0, dtypes.uchar, 128000)
     c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
     c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
-    c3 = UOp.param(1, dtypes.int, (512,))
+    c3 = UOp.param(1, dtypes.int, 512)
     c4 = c3.index(c1)  # c4 is a load
     c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
     c6 = ((c2*UOp.const(240))+c5)
-    c7 = UOp.param(2, dtypes.uchar, (60000,))
+    c7 = UOp.param(2, dtypes.uchar, 60000)
     c8 = c7.index(c6)
     # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
     loaded_idx = c4.cast(dtypes.weakint)
@@ -361,9 +361,9 @@ class TestUOpGraph(unittest.TestCase):
       self.assertNotEqual(u.dtype, dtypes.long)
 
   def test_fold_gated_load(self):
-    glbl0 = UOp.param(0, dtypes.int, (1,))
-    glbl1 = UOp.param(1, dtypes.int, (1,))
-    glbl2 = UOp.param(2, dtypes.int, (1,))
+    glbl0 = UOp.param(0, dtypes.int, 1)
+    glbl1 = UOp.param(1, dtypes.int, 1)
+    glbl2 = UOp.param(2, dtypes.int, 1)
     idx = UOp.const(0)
     ld0 = glbl1.index(UOp.invalid())
     ld1 = glbl2.index(idx.valid(UOp.const(True)))
@@ -372,7 +372,7 @@ class TestUOpGraph(unittest.TestCase):
     self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)
 
   def test_fold_gated_load_local(self):
-    glbl0 = UOp.param(0, dtypes.int, (16,))
+    glbl0 = UOp.param(0, dtypes.int, 16)
     smem = UOp.placeholder((18,), dtypes.int, slot=0, addrspace=AddrSpace.LOCAL)
     lidx = UOp.special(16, "lidx0")
     st = smem.index(lidx).store(glbl0.index(lidx).load())
@@ -385,7 +385,7 @@ class TestUOpGraph(unittest.TestCase):
     self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 2)
 
   def test_fold_gated_store(self):
-    glbl = UOp.param(0, dtypes.int, (1,))
+    glbl = UOp.param(0, dtypes.int, 1)
     idx0 = UOp.const(0)
     val = UOp.const(42)
     st0 = glbl.index(UOp.invalid()).store(val)
@@ -425,7 +425,7 @@ class TestReduceCollapse(unittest.TestCase):
 
   def test_reduce_shapeless_const_unroll(self):
     """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander"""
-    out = UOp.param(0, dtypes.float, (1,))
+    out = UOp.param(0, dtypes.float, 1)
     red = UOp.const(3.0).cast(dtypes.float).reduce(UOp.range(4, 0, AxisType.UNROLL), arg=(Ops.ADD, 0))
     ast = UOp.sink(out.index(UOp.const(0)).store(red)).replace(arg=KernelInfo())
     uops = full_rewrite_to_sink(ast, Device["CPU"].renderer, optimize=False).toposort()
@@ -441,7 +441,8 @@ class TestMovementOps(unittest.TestCase):
     self.assertEqual(result.op, Ops.INDEX)
     self.assertIs(result.src[0], src)
     self.assertEqual(result.shape, (4,))
-    self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])
+    # the only RESHAPE is src itself: the view of the flat param, the extra reshape was folded into the INDEX
+    self.assertEqual([u for u in result.toposort() if u.op is Ops.RESHAPE], [src])
 
   def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
     from tinygrad.schedule.prepare import pm_mops

test/null/test_uop_repr.py

@@ -4,21 +4,21 @@ from tinygrad import UOp
 class TestUOpRepr(unittest.TestCase):
   def test_simple_const(self):
     a = UOp.const(42)
-    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.weakint, arg=42, src=())")
+    self.assertEqual(repr(a), "UOp(Ops.CONST, arg=42, src=())")
   def test_different_consts(self):
     a, b = UOp.const(42), UOp.const(3)
     expected = (
-      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "  UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
-      "  UOp(Ops.CONST, dtypes.weakint, arg=3, src=()),))"
+      "UOp(Ops.ADD, arg=None, src=(\n" +
+      "  UOp(Ops.CONST, arg=42, src=()),\n" +
+      "  UOp(Ops.CONST, arg=3, src=()),))"
     )
     self.assertEqual(repr(a+b), expected)
   def test_walrus_operator_indentation(self):
     # The reference should have the same indentation as the definition
     a = UOp.const(42)
     expected = (
-      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
+      "UOp(Ops.ADD, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.CONST, arg=42, src=()),\n" +
       "  x0,))"
     )
     self.assertEqual(repr(a+a), expected)
@@ -26,9 +26,9 @@ class TestUOpRepr(unittest.TestCase):
     # Ensure indentation is consistent at multiple levels
     b = (a:=UOp.const(1)) + a
     expected = (
-      "UOp(Ops.MUL, dtypes.weakint, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "    x1:=UOp(Ops.CONST, dtypes.weakint, arg=1, src=()),\n" +
+      "UOp(Ops.MUL, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.ADD, arg=None, src=(\n" +
+      "    x1:=UOp(Ops.CONST, arg=1, src=()),\n" +
       "    x1,)),\n" +
       "  x0,))"
     )

test/null/test_uop_symbolic.py

@@ -1392,7 +1392,7 @@ class TestInvalidIndex(unittest.TestCase):
 
   def test_gated_load_keeps_index_valid(self):
     # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
-    buf = UOp.param(0, dtypes.bool, (17,))
+    buf = UOp.param(0, dtypes.bool, 17)
     ridx = Variable("ridx", 0, 31)
     cond = ridx < 17
     load = buf.index(ridx.valid(cond))
@@ -1404,7 +1404,7 @@ class TestStoreLoadFolding(unittest.TestCase):
   """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
   def test_store_load_folding(self):
     # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
-    buf = UOp.param(0, dtypes.int, (1,))
+    buf = UOp.param(0, dtypes.int, 1)
     index = buf.index(UOp.const(0))
     # Direct: store(idx, load(idx)) -> NOOP
     self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
@@ -1417,7 +1417,7 @@ class TestStoreLoadFolding(unittest.TestCase):
 
 class TestMoveWhereOnLoad(unittest.TestCase):
   def test_bool_index_preserves_dtype(self):
-    buf = UOp.param(0, dtypes.bool, (8,))
+    buf = UOp.param(0, dtypes.bool, 8)
     a = Variable("a", 0, 7)
     r = UOp.range(8, 0)
     # cond has a range that the rewrite can move into the valid: gate (a<4) goes into load valid
@@ -1475,7 +1475,7 @@ class TestRangeSplitting(unittest.TestCase):
     from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
     r0 = UOp.range(uconst(8), 0)
     # create a simple expression using the range with mod: store range%2 to a buffer
-    buf = UOp.param(0, dtypes.int, (1,))
+    buf = UOp.param(0, dtypes.int, 1)
     val = (r0 % uconst(2)).cast(dtypes.int)
     store = UOp(Ops.STORE, src=(buf.index(uconst(0)), val))
     sink = UOp(Ops.SINK, src=(UOp(Ops.END, src=(store, r0)),))

test/null/test_uop_vmin_vmax.py

@@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase):
   def test_vmin_vmax_multiplication_0_inf(self):
     # vmin and vmax for multiplication with a variable
     x = UOp.const(0.0)
-    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0))
+    y = UOp.load(UOp.param(0, dtypes.float, 1), UOp.const(0))
     uop = x * y
     # TODO: these should be 0, but definitely should not be nan
     self.assertEqual(uop.vmin, -math.inf)
@@ -332,7 +332,7 @@ class TestVminVmaxVConst(unittest.TestCase):
 
   def test_vmin_vmax_vector_with_gep(self):
     # vmin and vmax for a vector constant of bool values
-    d1 = UOp.param(1, dtypes.int, (1,))
+    d1 = UOp.param(1, dtypes.int, 1)
     idx = UOp.const(0)
     val = UOp(Ops.LOAD, src=(d1.index(idx),))
     uop = (val // 32)

test/null/test_uops.py

@@ -38,15 +38,9 @@ class TestDTypeFromUOp(unittest.TestCase):
     self.assertEqual(UOp(Ops.CONST, arg=ConstFloat(3.0)).dtype, dtypes.weakfloat)
     self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
     self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
-    # an explicit (strong) const dtype is legal until the field is removed
+    # UOp.const at a strong dtype builds the CAST that carries it
     self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)
 
-  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
-    self.assertIs(UOp.const(Invalid, dtypes.float32), UOp.invalid())
-    with self.assertRaises(RuntimeError): type_verify(UOp(Ops.CONST, dtypes.float32, arg=Invalid), spec_shared)
-
   def test_invalid_dtype_and_consumers(self):
     invalid = UOp.invalid()
     self.assertIs(invalid.dtype, dtypes.bool)
@@ -58,7 +52,7 @@ class TestDTypeFromUOp(unittest.TestCase):
     for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid)
     for u in (UOp(Ops.STACK, src=(value, invalid)), UOp(Ops.ADD, src=(value, invalid)),
               UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
-              UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
+              UOp.param(0, dtypes.float32, 4).index(invalid)): type_verify(u, spec_shared)
     gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
     self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float))
     type_verify(out.sink(), spec_program)
@@ -73,7 +67,7 @@ class TestLowerIndexDtype(unittest.TestCase):
   def test_gated_shrink_lowers_to_selected_width(self):
     # coalesce builds gated SHRINKs for masked vectorized loads; lowering must resolve them at the
     # width the offset bounds select (this one needs long)
-    buf = UOp.param(0, dtypes.float, (2**31+64,))
+    buf = UOp.param(0, dtypes.float, 2**31+64)
     i = UOp.variable("i", 0, 2**28)
     shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4)))
     lowered = graph_rewrite(shrink.sink(), pm_lower_weak)
@@ -84,11 +78,10 @@ class TestLowerIndexDtype(unittest.TestCase):
 
   def test_reg_buffer_size_lowers(self):
     reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG)
-    self.assertEqual(reg.src[0].dtype, dtypes.weakint)
+    self.assertEqual(reg.arg.size, 4)
     lowered = graph_rewrite(reg.sink(), pm_lower_weak)
     self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks),
                     "lowering must resolve every weak width, except a typed literal's value half")
-    self.assertEqual(next(u for u in lowered.backward_slice_with_self if u.op is Ops.BUFFER).src[0].dtype, dtypes.int)
 
 class TestSafeCast(unittest.TestCase):
   def test_cast_folds(self):
@@ -221,7 +214,7 @@ class TestExecALU(unittest.TestCase):
 
 class TestGatedStoreRewrite(unittest.TestCase):
   def test_tiny_gate_store(self):
-    gmem = UOp.param(0, dtypes.float, (8,))
+    gmem = UOp.param(0, dtypes.float, 8)
     gidx0 = UOp.special(4, 'gidx0')
     gate = gidx0<UOp.const(1)
     idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(2)).valid(gate)))
@@ -237,8 +230,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
     self.assertEqual(len(gated_uops[-1].src), 2)
 
   def test_gate_some_stores(self):
-    gmem0 = UOp.param(0, dtypes.float, (8,))
-    gmem1 = UOp.param(1, dtypes.float, (8,))
+    gmem0 = UOp.param(0, dtypes.float, 8)
+    gmem1 = UOp.param(1, dtypes.float, 8)
     gidx0 = UOp.special(4, 'gidx0')
     idx = gidx0 * UOp.const(2)
     idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(1))))
@@ -257,8 +250,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
   # scaled down version of TestLinearizerDumb.test_unmerged_ifs
   @unittest.skip("we don't merge ifs anymore")
   def test_merge_ifs_alt(self):
-    gmem0 = UOp.param(0, dtypes.float, (8,))
-    gmem1 = UOp.param(1, dtypes.float, (8,))
+    gmem0 = UOp.param(0, dtypes.float, 8)
+    gmem1 = UOp.param(1, dtypes.float, 8)
     gidx0 = UOp.special(4, 'gidx0')
     idx = gidx0*UOp.const(2)
     gate = gidx0<UOp.const(1)
@@ -281,7 +274,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
 class TestFastIdiv(unittest.TestCase):
   def test_division_power_of_two(self):
     for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp.param(0, dt, (3,))
+      g = UOp.param(0, dt, 3)
       c = UOp.const(2)
       l = g.index(c)
       a = UOp(Ops.CDIV, src=(l, c))
@@ -294,7 +287,7 @@ class TestFastIdiv(unittest.TestCase):
   def test_floormod_power_of_two(self):
     # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
     for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp.param(0, dt, (9,))
+      g = UOp.param(0, dt, 9)
       c = UOp.const(8)
       a = UOp(Ops.FLOORMOD, src=(g.index(c), c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -306,7 +299,7 @@ class TestFastIdiv(unittest.TestCase):
   def test_floordiv_power_of_two(self):
     # FLOORDIV by a power of two lowers to a shift, with no round toward zero correction (a shift is exactly floor division)
     for dt in (dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
-      g = UOp.param(0, dt, (3,))
+      g = UOp.param(0, dt, 3)
       c = UOp.const(2)
       a = UOp(Ops.FLOORDIV, src=(g.index(c), c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -319,7 +312,7 @@ class TestFastIdiv(unittest.TestCase):
   @Context(DISABLE_FAST_IDIV=0)
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
   def test_fast_idiv_and_mod(self):
-    g = UOp.param(0, dtypes.uint32, (4,))
+    g = UOp.param(0, dtypes.uint32, 4)
     c = UOp.const(3)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
@@ -353,7 +346,7 @@ class TestFastIdiv(unittest.TestCase):
   @unittest.expectedFailure
   def test_fast_idiv_overflow(self):
     # This will be possible with a slightly different method for fast_idiv
-    g = UOp.param(0, dtypes.uint32, (8,))
+    g = UOp.param(0, dtypes.uint32, 8)
     c = UOp.const(7).cast(dtypes.uint)
     l = UOp(Ops.LOAD, src=(g.index(c),))
     a = UOp(Ops.CDIV, src=(l, c))
@@ -364,7 +357,7 @@ class TestFastIdiv(unittest.TestCase):
     self.assertNotIn(Ops.CDIV, ops)
 
   def test_disable_fast_idiv(self):
-    g = UOp.param(0, dtypes.uint32, (4,))
+    g = UOp.param(0, dtypes.uint32, 4)
     c = UOp.const(3)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
@@ -407,8 +400,8 @@ class TestUOpMethod(unittest.TestCase):
       self.assertIs((x != x).simplify(), x.const_like(False, dtypes.bool))
 
   def test_replace(self):
-    x = UOp.param(0, dtypes.int, (1,))
-    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, (1,)).arg).arg.slot, 1)
+    x = UOp.param(0, dtypes.int, 1)
+    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, 1).arg).arg.slot, 1)
     with self.assertRaises(AssertionError): x.replace(field="a")
 
   def test_const_zero_neg_zero_different(self):
@@ -463,7 +456,7 @@ class TestUOpRender(unittest.TestCase):
   def test_render_ssimplified_marg_outside_toposort(self):
     r = UOp.range(UOp.const(16, dtypes.int), 2, AxisType.WEAK, dtype=dtypes.int)
     offset = (r * 2) + (r * 2)
-    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, (32,)), offset, UOp.const(2, dtypes.int)))
+    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, 32), offset, UOp.const(2, dtypes.int)))
     self.assertIsNot(shrink.src[1], shrink.marg[0][0])
     self.assertEqual(shrink.render(simplify=False), "p0.shrink((((r2*4), 2),))")
     self.assertEqual(UOp.range(1, 0, src=(shrink,), dtype=dtypes.int).render(simplify=False), "r0")

test/null/test_uops_stats.py

@@ -146,7 +146,7 @@ class TestUOpsStats(unittest.TestCase):
 
   #MULACC should have the same stats as MUL + ADD
   def test_mulacc(self):
-    globl = UOp.param(0, dtypes.int, (3,))
+    globl = UOp.param(0, dtypes.int, 3)
     o1 = UOp.const(1, dtypes.int)
     o2 = UOp.const(2, dtypes.int)
     u1 = globl.index(o1)
@@ -156,7 +156,7 @@ class TestUOpsStats(unittest.TestCase):
     u5 = UOp(Ops.ADD, src=(u4,u3))
     uops = tuple(u5.toposort())
 
-    globl = UOp.param(0, dtypes.int, (3,))
+    globl = UOp.param(0, dtypes.int, 3)
     o1 = UOp.const(1, dtypes.int)
     o2 = UOp.const(2, dtypes.int)
     u1 = globl.index(o1)

test/null/test_validate_oob.py

@@ -13,7 +13,7 @@ class TestValidateOOB(unittest.TestCase):
   # basic index patterns
   def test_const_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       to_uops_list([buf.index(UOp.const(0)).load()])  # valid
       to_uops_list([buf.index(UOp.const(15)).load()])  # valid (last element)
       with self.assertRaises(RuntimeError):
@@ -23,7 +23,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_variable_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       to_uops_list([buf.index(Variable("i", 0, 15)).load()])  # valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(Variable("i", 0, 20)).load()])  # oob
@@ -32,7 +32,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_range_with_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       r = UOp.range(42, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r.valid(r < 16)).load()])  # valid
       with self.assertRaises(RuntimeError):
@@ -40,7 +40,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_variable_with_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       v = Variable("v", -5, 80)
       to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load()])  # valid
       with self.assertRaises(RuntimeError):
@@ -48,7 +48,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_gated_store(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       v = Variable("v", 0, 20)
       to_uops_list([buf.index(v.valid(v < 16)).store(0)])  # valid
       with self.assertRaises(RuntimeError):
@@ -57,14 +57,14 @@ class TestValidateOOB(unittest.TestCase):
   # ALU ops in index
   def test_floordiv(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load()])  # 0..15 valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load()])  # 0..16 oob
 
   def test_mod(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       r = UOp.range(100, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r % 16).load()])  # 0..15 valid
       with self.assertRaises(RuntimeError):
@@ -72,14 +72,14 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_shr(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load()])  # 0..15 valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load()])  # 0..31 oob
 
   def test_shl(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (64,))
+      buf = UOp.param(0, dtypes.int, 64)
       r = UOp.range(8, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r << 2).load()])  # 0..28 valid
       with self.assertRaises(RuntimeError):
@@ -87,7 +87,7 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_and(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       r = UOp.range(100, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r & 15).load()])  # 0..15 valid
       with self.assertRaises(RuntimeError):
@@ -102,14 +102,14 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_max(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load()])  # 0..15 valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load()])  # 0..20 oob
 
   def test_xor_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       r = UOp.range(32, 0, AxisType.GLOBAL)
       to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load()])  # 0..15 valid
       with self.assertRaises(RuntimeError):
@@ -118,22 +118,22 @@ class TestValidateOOB(unittest.TestCase):
   # cast patterns
   def test_float_cast_in_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (16,))
+      buf = UOp.param(0, dtypes.int, 16)
       r = UOp.range(20, 0)
       i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
       to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()])
 
   def test_bool_cast_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, (1,))
+      buf = UOp.param(0, dtypes.int, 1)
       r = UOp.range(20, 0)
       to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load()])  # only r=0 valid
 
   # load result as index/mask
   def test_load_as_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp.param(0, dtypes.int, (16,))
-      buf1 = UOp.param(1, dtypes.int, (64,))
+      buf0 = UOp.param(0, dtypes.int, 16)
+      buf1 = UOp.param(1, dtypes.int, 64)
       r = UOp.range(42, 0, AxisType.GLOBAL)
       ld0 = buf0.index(r.valid(r < 8)).load().cast(dtypes.weakint)
       to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load()])  # valid
@@ -142,16 +142,16 @@ class TestValidateOOB(unittest.TestCase):
 
   def test_load_from_shrink_as_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp.param(0, dtypes.int, (16,))
-      buf1 = UOp.param(1, dtypes.int, (64,))
+      buf0 = UOp.param(0, dtypes.int, 16)
+      buf1 = UOp.param(1, dtypes.int, 64)
       shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4)))
       ld0 = shrink.load().index(0)
       to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load()])
 
   def test_load_bool_as_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
-      buf_bool = UOp.param(0, dtypes.bool, (16,))
-      buf_int = UOp.param(1, dtypes.int, (8,))
+      buf_bool = UOp.param(0, dtypes.bool, 16)
+      buf_int = UOp.param(1, dtypes.int, 8)
       gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0")
       ld_bool = buf_bool.index(gidx).load()
       with self.assertRaises(RuntimeError):
@@ -162,7 +162,7 @@ class TestValidateOOB(unittest.TestCase):
   def test_in_bounds_access_gated_local(self):
     with Context(CHECK_OOB=1):
       # Define buffers
-      gbuf = UOp.param(0, dtypes.uint, (400,))
+      gbuf = UOp.param(0, dtypes.uint, 400)
       sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)
 
       # Define indices, valids and barrier
@@ -186,8 +186,8 @@ class TestValidateOOB(unittest.TestCase):
   @unittest.skip("Bool load is not supported yet")
   def test_load_mask(self):
     with Context(CHECK_OOB=1):
-      glbl0 = UOp.param(0, dtypes.int, (16,))
-      mask = UOp.param(0, dtypes.bool, (16,))
+      glbl0 = UOp.param(0, dtypes.int, 16)
+      mask = UOp.param(0, dtypes.bool, 16)
       ridx = UOp.range(20, 0)
       ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
       to_uops_list([ld0])

test/null/test_viz.py

@@ -244,7 +244,7 @@ class TestViz(unittest.TestCase):
     self.assertEqual(len(lst), 1)
     graphs = [x["graph"] for x in viz.get_details(0, 0)]
     # const is always in the graph, client side hides exclude=True nodes by default
-    self.assertEqual(list(graphs[0]), [id(a.src[0]), id(a), id(z), id(alu), id(y), id(sink)])
+    self.assertEqual(list(graphs[0]), [id(a), id(z), id(alu), id(y), id(sink)])
     self.assertTrue(graphs[0][id(z)]["exclude"])
     self.assertTrue(graphs[0][id(y)]["exclude"])
     self.assertFalse(graphs[0][id(alu)]["exclude"])

test/unit/test_attention.py

@@ -73,10 +73,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
 
   def _make_config(self, **kwargs):
-    return TransformerConfig(**({"num_blocks":1, "dim":32, "hidden_dim":64, "n_heads":1, "n_kv_heads":1,
-                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":32, "rope_theta":10000.0,
-                                 "rope_dim":32, "v_head_dim":32, "max_context":4, "ssm_layers":(True,),
-                                 "ssm":SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32)} | kwargs))
+    return TransformerConfig(**({"num_blocks":1, "dim":8, "hidden_dim":16, "n_heads":1, "n_kv_heads":1,
+                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":8, "rope_theta":10000.0,
+                                 "rope_dim":8, "v_head_dim":8, "max_context":4, "ssm_layers":(True,),
+                                 "ssm":SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4)} | kwargs))
 
   def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
     block = GatedDeltaNetBlock(config, config.ssm)
@@ -229,7 +229,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha[..., None], rtol=1e-5, atol=1e-5)
 
   def test_kda_prefill_matches_decode(self):
-    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=True))
+    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=True))
     block = GatedDeltaNetBlock(config, config.ssm)
     for p in nn.state.get_parameters(block):
       p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
@@ -245,7 +245,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
 
   def test_varied_chunk_sizes_match_decode(self):
     for kda in (False, True):
-      ssm = SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=kda)
+      ssm = SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=kda)
       config = self._make_config(ssm=ssm)
       if kda:
         block = GatedDeltaNetBlock(config, config.ssm)
@@ -267,7 +267,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
         np.testing.assert_allclose(chunked_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
 
   def test_start_zero_resets_realized_state(self):
-    config, x = self._make_config(max_context=3), self._tensor_linspace(-1, 1, (1, 3, 32))
+    config = self._make_config(max_context=3)
+    x = self._tensor_linspace(-1, 1, (1, 3, config.dim))
     block = self._make_block(config)
     self._run_attention(block, x, 0)
     restarted = self._run_attention(block, x[:, :2], 0)

test/unit/test_call.py

@@ -52,6 +52,11 @@ class TestCall(unittest.TestCase):
     np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
     np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
 
+  def test_call_scalar_param_shape_mismatch(self):
+    scalar_fxn = UOp.param(0, dtypes.float, ()) * 2
+    with self.assertRaisesRegex(TypeError, "shape mismatch: expected scalar"):
+      Tensor.call(Tensor.ones(2), fxn=scalar_fxn).realize()
+
   def test_call_gemm(self):
     M, K, N = 4, 8, 4
     a = Tensor.randn(M, K)

test/unit/test_dtype_weak.py

@@ -116,7 +116,7 @@ class TestWeakPromotion(unittest.TestCase):
 
   def test_store_weak_value_uses_destination_dtype(self):
     with Context(DEFAULT_FLOAT=dtypes.float16):
-      dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32))
+      dst = UOp.param(0, dtypes.bfloat16, 1).index(UOp.const(0).cast(dtypes.int32))
       gate = UOp.const(True)
       out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_commit_weak)
     # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it
@@ -132,7 +132,7 @@ class TestWeakPromotion(unittest.TestCase):
 
   def test_derivable_const_rounds_at_the_derived_width(self):
     # re-rounds a derivable const in place (still bare) so value-keyed folds (x*1 -> x, x*-1 -> NEG) still fire
-    x = UOp.param(0, dtypes.float32, (1,)).index(UOp.const(0).cast(dtypes.int32)).load()
+    x = UOp.param(0, dtypes.float32, 1).index(UOp.const(0).cast(dtypes.int32)).load()
     mul = graph_rewrite(x * UOp.const(-0.9999999893980771), symbolic_simple+pm_commit_weak)
     self.assertIs(mul.src[1], UOp.const(-1.0))
     self.assertIs(graph_rewrite(x * UOp.const(1.0000000106), symbolic_simple+pm_commit_weak), x)
@@ -191,7 +191,7 @@ class TestWeakPromotion(unittest.TestCase):
     # float bitwise builds, the spec rejects it
     with Context(SPEC=1):
       f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0)
-      for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, dtypes.float32, (f32, f32)), UOp(Ops.AND, dtypes.int32, (wf, wf))):
+      for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, (f32, f32)), UOp(Ops.AND, (wf, wf))):
         with self.assertRaises(RuntimeError): type_verify([bad], spec_shared)
 
   def test_integer_values(self):

test/unit/test_gguf.py

@@ -5,17 +5,33 @@ from tinygrad.llm.gguf import _ggml_iq_grid, ggml_data_to_tensor, gguf_load
 from tinygrad.runtime.autogen import ggml_common as _ggml
 import numpy as np
 from gguf import GGUFReader, GGUFValueType, GGMLQuantizationType, GGML_QUANT_SIZES, dequantize, quantize
-from gguf.quants import IQ2_S, IQ3_S, IQ3_XXS
+from gguf.quants import IQ1_S, IQ2_S, IQ2_XS, IQ2_XXS, IQ3_S, IQ3_XXS
 
 ggml_test_block_count = 4
 supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
 
 class TestGGUFTables(unittest.TestCase):
+  def test_iq2_xxs_grid_matches_gguf_py(self):
+    IQ2_XXS.init_grid()
+    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq2xxs_grid, (256, 8)).numpy()
+    np.testing.assert_equal(grid, IQ2_XXS.grid.reshape(256, 8))
+
+  def test_iq2_xs_grid_matches_gguf_py(self):
+    IQ2_XS.init_grid()
+    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq2xs_grid, (512, 8)).numpy()
+    np.testing.assert_equal(grid, IQ2_XS.grid.reshape(512, 8))
+
   def test_iq2_s_grid_matches_gguf_py(self):
     IQ2_S.init_grid()
     grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq2s_grid, (1024, 8)).numpy()
     np.testing.assert_equal(grid, IQ2_S.grid.reshape(1024, 8))
 
+  def test_iq1_s_grid_matches_gguf_py(self):
+    IQ1_S.init_grid()
+    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq1s_grid, (2048, 8)).numpy()
+    grid = np.where(grid > 127, grid - 256, grid)
+    np.testing.assert_equal(grid, IQ1_S.grid.reshape(2048, 8))
+
   def test_iq3_xxs_grid_matches_gguf_py(self):
     IQ3_XXS.init_grid()
     grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq3xxs_grid, (256, 4)).numpy()
@@ -39,6 +55,56 @@ class TestGGUF(unittest.TestCase):
     expected = np.arange(1, 33, dtype=np.float32) * 2.0
     np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.Q8_0.value).numpy().flatten(), expected)
 
+  def test_dequantization_q2_k_hardcoded(self):
+    # Q2_K: scales[16] + qs[64] + d(fp16) + dmin(fp16). 16 sub-blocks of 16, x = d*(scale&0xF)*q - dmin*(scale>>4)
+    scales, qs = bytes([0x11]*16), bytes([0x55]*64)  # scale=1, min=1; qs=0x55 -> 2-bit quants of 1
+    d, dmin = np.float16(1.0).tobytes(), np.float16(0.0).tobytes()
+    block = np.frombuffer(scales + qs + d + dmin, dtype=np.uint8).copy()
+    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 10).numpy().flatten(), np.ones(256, dtype=np.float32))
+
+  def test_dequantization_q3_k_hardcoded(self):
+    # Q3_K: hmask[32] + qs[64] + scales[12] + d(fp16). 16 sub-blocks of 16, x = d * (scale-32) * (q - (hbit?0:4))
+    # 6-bit scales 32..47 so (scale-32) = 0..15; qs=0x55 -> 2-bit quants of 1; d=1.0
+    scales = bytes([0x80, 0x91, 0xA2, 0xB3, 0xC4, 0xD5, 0xE6, 0xF7, 0xAA, 0xAA, 0xAA, 0xAA])
+    d = np.float16(1.0).tobytes()
+    qs, ones = bytes([0x55]*64), np.ones(16, dtype=np.float32)
+    # hmask all-ones: high bit set, q=1; hmask zeros: subtract 4, q=-3
+    for hmask, q in ((bytes([0xFF]*32), 1.0), (bytes([0x00]*32), -3.0)):
+      block = np.frombuffer(hmask + qs + scales + d, dtype=np.uint8).copy()
+      expected = np.concatenate([q * s * ones for s in range(16)])
+      np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 11).numpy().flatten(), expected)
+
+  def test_dequantization_iq2_xxs_hardcoded(self):
+    # IQ2_XXS: d + 8 groups of (4 grid bytes + uint32 signs/scale). grid[0]=all 0x08, scale=0, signs=0
+    # db = 1.0 * (0.5 + 0) * 0.25 = 0.125; 0.125 * 8 = 1.0
+    block = np.frombuffer(np.float16(1.0).tobytes() + bytes(64), dtype=np.uint8).copy()
+    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 16).numpy().flatten(), np.ones(256, dtype=np.float32))
+
+  def test_dequantization_iq2_xs_hardcoded(self):
+    # IQ2_XS: d + 32 uint16 qs + 8 scale bytes. qs=0 -> grid[0]=all 0x08, signs=0; scales=0
+    block = np.frombuffer(np.float16(1.0).tobytes() + bytes(64) + bytes(8), dtype=np.uint8).copy()
+    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 17).numpy().flatten(), np.ones(256, dtype=np.float32))
+
+  def test_dequantization_iq1_s_hardcoded(self):
+    # IQ1_S: d + qs[32] + qh[16]. qs=qh=0 -> grid[0]=all -1, scale=1, delta=+0.125 -> -0.875
+    block = np.frombuffer(np.float16(1.0).tobytes() + bytes(48), dtype=np.uint8).copy()
+    expected = np.full(256, -0.875, dtype=np.float32)
+    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 19).numpy().flatten(), expected)
+
+  def test_dequantization_iq1_m_hardcoded(self):
+    # IQ1_M: qs[32] + qh[16] + scales[8]. f16 1.0=0x3C00 packed in high nibbles; qs=qh=0 -> -0.875
+    scales = bytes([0x00, 0x00, 0x00, 0x00, 0x00, 0xC0, 0x00, 0x30])
+    block = np.frombuffer(bytes(48) + scales, dtype=np.uint8).copy()
+    expected = np.full(256, -0.875, dtype=np.float32)
+    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 29).numpy().flatten(), expected)
+
+  def test_dequantization_iq4_nl_hardcoded(self):
+    # IQ4_NL: 2-byte fp16 scale + 16 packed bytes. low nibbles first, then high
+    lut = list(_ggml.kvalues_iq4nl)
+    block = np.frombuffer(np.float16(1.0).tobytes() + bytes(range(16)), dtype=np.uint8).copy()
+    expected = np.array(lut + [lut[0]]*16, dtype=np.float32)
+    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, 20).numpy().flatten(), expected)
+
   def test_dequantization_mxfp4_hardcoded(self):
     # MXFP4: 1 byte shared exponent E + 16 packed bytes (32 x 4-bit values)
     # nibble: bit3=sign, bit2:1=exp, bit0=mant; E=128 gives scale=1.0
@@ -52,13 +118,20 @@ class TestGGUF(unittest.TestCase):
   def test_dequantization_q5_0(self): self._test_dequantization(GGMLQuantizationType.Q5_0)
   def test_dequantization_q5_1(self): self._test_dequantization(GGMLQuantizationType.Q5_1)
   def test_dequantization_q8_0(self): self._test_dequantization(GGMLQuantizationType.Q8_0)
+  def test_dequantization_q2_k(self): self._test_dequantization(GGMLQuantizationType.Q2_K)
+  def test_dequantization_q3_k(self): self._test_dequantization(GGMLQuantizationType.Q3_K)
   def test_dequantization_q4_k(self): self._test_dequantization(GGMLQuantizationType.Q4_K)
   def test_dequantization_q5_k(self): self._test_dequantization(GGMLQuantizationType.Q5_K)
   def test_dequantization_q6_k(self): self._test_dequantization(GGMLQuantizationType.Q6_K)
+  def test_dequantization_iq2_xxs(self): self._test_dequantization(GGMLQuantizationType.IQ2_XXS)
+  def test_dequantization_iq2_xs(self): self._test_dequantization(GGMLQuantizationType.IQ2_XS)
   def test_dequantization_iq3_xxs(self): self._test_dequantization(GGMLQuantizationType.IQ3_XXS)
+  def test_dequantization_iq1_s(self): self._test_dequantization(GGMLQuantizationType.IQ1_S)
+  def test_dequantization_iq4_nl(self): self._test_dequantization(GGMLQuantizationType.IQ4_NL)
   def test_dequantization_iq3_s(self): self._test_dequantization(GGMLQuantizationType.IQ3_S)
   def test_dequantization_iq2_s(self): self._test_dequantization(GGMLQuantizationType.IQ2_S)
   def test_dequantization_iq4_xs(self): self._test_dequantization(GGMLQuantizationType.IQ4_XS)
+  def test_dequantization_iq1_m(self): self._test_dequantization(GGMLQuantizationType.IQ1_M)
   def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4)
   @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, "Backend must support bfloat16")
   def test_dequantization_bf16(self): self._test_dequantization(GGMLQuantizationType.BF16)
@@ -203,12 +276,23 @@ class TestGGUFGEMV(unittest.TestCase):
       q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
       scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
       if qtype in (GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0,
-                   GGMLQuantizationType.IQ3_XXS,
-                   GGMLQuantizationType.IQ2_S,
+                   GGMLQuantizationType.IQ2_XXS, GGMLQuantizationType.IQ2_XS,
+                   GGMLQuantizationType.IQ3_XXS, GGMLQuantizationType.IQ4_NL,
+                   GGMLQuantizationType.IQ1_S, GGMLQuantizationType.IQ2_S,
                    GGMLQuantizationType.IQ3_S, GGMLQuantizationType.IQ4_XS): q_data[:, :2] = scales[:, :2]  # d at offset 0
       elif qtype in (GGMLQuantizationType.Q5_1, GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K):
         q_data[:, :4] = scales[:, :4]  # d, m/dmin at offset 0
-      elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]               # d at end
+      elif qtype == GGMLQuantizationType.Q2_K: q_data[:, -4:] = scales[:, :4]  # d, dmin at end
+      elif qtype in (GGMLQuantizationType.Q6_K, GGMLQuantizationType.Q3_K): q_data[:, -2:] = scales[:, :2]  # d at end
+      elif qtype == GGMLQuantizationType.IQ1_M:
+        s = np.float16(rng.standard_normal(n_blocks)).view(np.uint16)
+        sc = q_data[:, -8:].copy().view(np.uint16).reshape(n_blocks, 4)
+        sc &= np.uint16(0x0FFF)
+        sc[:, 0] |= (s & np.uint16(0x000F)) << 12
+        sc[:, 1] |= (s & np.uint16(0x00F0)) << 8
+        sc[:, 2] |= (s & np.uint16(0x0F00)) << 4
+        sc[:, 3] |= (s & np.uint16(0xF000))
+        q_data[:, -8:] = sc.reshape(n_blocks, -1).view(np.uint8)
       elif qtype == GGMLQuantizationType.MXFP4: q_data[:, 0] = rng.integers(120, 136, size=n_blocks, dtype=np.uint8) # constrain byte0
       q_data = q_data.flatten()
     ref = dequantize(q_data, qtype).reshape(rows, cols)
@@ -235,13 +319,20 @@ class TestGGUFGEMV(unittest.TestCase):
   def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)
   def test_gguf_gemv_q5_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_0)
   def test_gguf_gemv_q5_1(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_1)
+  def test_gguf_gemv_q2_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q2_K)
+  def test_gguf_gemv_q3_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q3_K)
   def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
   def test_gguf_gemv_q5_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_K)
   def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
+  def test_gguf_gemv_iq2_xxs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_XXS)
+  def test_gguf_gemv_iq2_xs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_XS)
   def test_gguf_gemv_iq3_xxs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_XXS)
+  def test_gguf_gemv_iq1_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ1_S)
+  def test_gguf_gemv_iq4_nl(self): self._test_gguf_gemv(GGMLQuantizationType.IQ4_NL)
   def test_gguf_gemv_iq3_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_S)
   def test_gguf_gemv_iq2_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_S)
   def test_gguf_gemv_iq4_xs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ4_XS)
+  def test_gguf_gemv_iq1_m(self): self._test_gguf_gemv(GGMLQuantizationType.IQ1_M)
   def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4)
   @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, "Backend must support bfloat16")
   def test_gguf_gemv_bf16(self): self._test_gguf_gemv(GGMLQuantizationType.BF16)

tinygrad/codegen/__init__.py

@@ -126,7 +126,7 @@ def do_devectorize(b:UOp):
   if not all(x.shape == b.shape or x.base.is_invalid for x in b.src): return None
   src = []
   for idx_c in itertools.product(*[[UOp.const(i) for i in range(x)] for x in b.shape]):
-    src.append(b.replace(dtype=None, src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src)))
+    src.append(b.replace(src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src)))
   return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
 
 def do_stack_wmma(u:UOp):
@@ -391,7 +391,15 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
 
   if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
-  if SPEC: type_verify(sink, spec_program)
+  if SPEC:
+    import os
+    if os.environ.get("DBGTV"):
+      try: type_verify(sink, spec_program)
+      except RuntimeError:
+        from tinygrad.uop.render import print_uops
+        print_uops(list(sink.toposort()))
+        raise
+    else: type_verify(sink, spec_program)
 
   # return the rewritten sink
   return sink

tinygrad/codegen/decomp/dtype.py

@@ -140,8 +140,8 @@ def f2f_store(st, idx, val, fr:DType, to:DType):
 pm_long_decomp: PatternMatcher = PatternMatcher([
   # the decomp's own bottom-up rewrite can mint bare consts mid-flight: word splitting commits them at the long sibling's dtype
   (UPat(GroupOp.All, name='x'), lambda x: commit_weak_consts(x, next((s.dtype for s in x.src if s.dtype in l2i_dt), None))),
-  (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), src=(UPat.var("sz"),), name="x"), lambda x,sz:
-   UOp(x.op, src=(sz*2,), arg=replace(x.arg, dtype=l2i_dt[x.dtype]), tag=x.tag)),
+  (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), name="x"), lambda x:
+   UOp(x.op, arg=replace(x.arg, dtype=l2i_dt[x.dtype], size=None if x.arg.size is None else x.arg.size*2), tag=x.tag)),
   (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x:
    reindex(x, x.tag[0]).replace(tag=None) if x.tag is not None else None),
   (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val:

tinygrad/codegen/gpudims.py

@@ -73,7 +73,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
   subs = {}
   for r in s_topo:
     # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
-    if r.op is Ops.STORE and (idx := r.src[0]).src[0].addrspace == AddrSpace.GLOBAL:
+    if r.op is Ops.STORE and len((idx := r.src[0]).src) and idx.src[0].addrspace == AddrSpace.GLOBAL:
       missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
       if len(missing_locals):
         assert len(idx.src) == 2, "index has 2 sources"

tinygrad/codegen/late/coalesce.py

@@ -1,7 +1,8 @@
 import itertools, functools
 from collections import defaultdict
+from dataclasses import replace
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType
-from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg, graph_rewrite
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, graph_rewrite
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate, sym
 from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape
 from tinygrad.renderer import Renderer
@@ -85,7 +86,8 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
   if len(cands) == 0: return None
   # and tiebreak with indexing complexity (ie. number of nodes)
   h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice))
-  buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),))
+  # the image dims are stored in the param's arg, the size stays the flat buffer len
+  buf = buf.replace(arg=replace(buf.arg, image=(h, w)))
   shapes[buf.arg.slot] = (h, w)
   if valid.op is not Ops.CONST or valid.val is not True:
     return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid))

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.CAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
+PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.BITCAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
 
 class LinearScanRegallocContext:
   # returns the uop that defines the virtual register

tinygrad/dtype.py

@@ -120,8 +120,6 @@ class DTypes:
   uint32: Final[DType] = DType.new(6, 32, "unsigned int", 'I')
   int64: Final[DType] = DType.new(7, 64, "long", 'q')
   uint64: Final[DType] = DType.new(8, 64, "unsigned long", 'Q')
-  _uint128: Final[DType] = DType.new(8, 128, "uint128", None)
-  _uint256: Final[DType] = DType.new(8, 256, "uint256", None)
   weakfloat: Final[DType] = DType.new(9, 800, "weakfloat", None)
   fp8e4m3: Final[DType] = DType.new(10, 8, "float8_e4m3", None)
   fp8e5m2: Final[DType] = DType.new(11, 8, "float8_e5m2", None)

tinygrad/engine/jit.py

@@ -67,7 +67,7 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
   if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear")
 
   # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
-  linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True)
+  linear = linear.substitute({u: UOp.param(i, u.dtype, u.max_numel(), u.device) for i,u in enumerate(input_uops)}, walk=True)
   linear = memory_plan_rewrite(linear, held_bufs)
   linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
   if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)

tinygrad/engine/realize.py

@@ -2,8 +2,8 @@ from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
 import random, itertools, math, weakref, array, decimal
 from dataclasses import dataclass, replace, field
-from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple, tqdm, dedup
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, to_tuple, tqdm, dedup
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, HCQ2, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.dtype import dtypes
@@ -305,17 +305,17 @@ pm_exec = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])
 
-if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV # noqa: E402 # down here, hcq2 imports realize
+from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV # noqa: E402 # down here, hcq2 imports realize
 
 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
   linear = lower_and_compile(linear)
   linear = graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
-  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
+  if HCQ2: linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
   return linear
 
-def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear
+def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if HCQ2 else linear
 
 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
   inputs = list(input_uops)

tinygrad/helpers.py

@@ -240,7 +240,7 @@ TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
 RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0)
-VALIDATE_WITH_CPU = ContextVar("VALIDATE_WITH_CPU", 0)
+VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 0)
 # TODO: this is broken for some indexing
 DISABLE_FAST_IDIV = ContextVar("DISABLE_FAST_IDIV", 1)
 FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0)

tinygrad/llm/cli.py

@@ -145,6 +145,7 @@ def main():
   parser.add_argument("--serve", nargs='?', type=int, const=8000, metavar="PORT", help="Run OpenAI compatible API (optional port, default 8000)")
   parser.add_argument("--warmup", action="store_true", help="warmup the JIT")
   parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
+  parser.add_argument("--no_chat_template", action="store_true", help="Don't use the model's chat template, always use the fallback template")
   args = parser.parse_args()
 
   # load the model
@@ -160,7 +161,7 @@ def main():
 
   # use the model's chat template if jinja2 is available (enables model-specific formatting)
   template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
-  if (ct := kv.get('tokenizer.chat_template')) is not None:
+  if not args.no_chat_template and (ct := kv.get('tokenizer.chat_template')) is not None:
     try:
       import jinja2
       env = jinja2.Environment()

tinygrad/llm/gguf.py

@@ -18,7 +18,9 @@ _GGML_NATIVE = {0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8, 25: dtype
 
 # quant types {ggml_type: (number of elements, number of bytes)}
 _GGML_QUANT = {2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
-               12:(256,144), 13:(256,176), 14:(256,210), 18:(256,98), 21:(256,110), 22:(256,82), 23:(256,136), 39:(32,17), 41:(128,18)}
+               10:(256,84), 11:(256,110), 12:(256,144), 13:(256,176), 14:(256,210),
+               16:(256,66), 17:(256,74), 18:(256,98), 19:(256,50), 20:(32,18), 21:(256,110), 22:(256,82), 23:(256,136),
+               29:(256,56), 39:(32,17), 41:(128,18)}
 
 def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   """
@@ -27,8 +29,9 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   Supported native types: float32 (id: 0), float16 (id: 1), int8 (id: 24),
   int16 (id: 25), int32 (id: 26), int64 (id: 27), float64 (id: 28), bfloat16 (id: 30)
   Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q5_0 (id: 6),
-  Q5_1 (id: 7), Q8_0 (id: 8), Q4_K (id: 12), Q5_K (id: 13),
-  Q6_K (id: 14), IQ3_XXS (id: 18), IQ3_S (id: 21), IQ2_S (id: 22), IQ4_XS (id: 23), MXFP4 (id: 39), Q1_0 (id: 41)
+  Q5_1 (id: 7), Q8_0 (id: 8), Q2_K (id: 10), Q3_K (id: 11), Q4_K (id: 12), Q5_K (id: 13),
+  Q6_K (id: 14), IQ2_XXS (id: 16), IQ2_XS (id: 17), IQ3_XXS (id: 18), IQ1_S (id: 19),
+  IQ4_NL (id: 20), IQ3_S (id: 21), IQ2_S (id: 22), IQ4_XS (id: 23), IQ1_M (id: 29), MXFP4 (id: 39), Q1_0 (id: 41)
   """
   # https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356
 
@@ -54,6 +57,19 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       q = q_to_uint8(blocks[:,qh_off+4:], 4).bitcast(dtypes.int8) + qh * 16
       return q * d + (blocks[:,2:4].bitcast(dtypes.float16).cast(dtypes.float32) if ggml_type == 7 else -16 * d)
     if ggml_type == 8: return blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32) * blocks[:,2:].bitcast(dtypes.int8)
+    # Q2_K: 256 elements per 84-byte block (scales:16, qs:64, d:2, dmin:2)
+    if ggml_type == 10:
+      d, dmin = (blocks[:,i:i+2].bitcast(dtypes.float16).cast(dtypes.float32).unsqueeze(-1) for i in [80, 82])
+      sc = blocks[:, :16]
+      q = q_to_uint8(blocks[:, 16:80].reshape((-1, 2, 32)), 2).reshape((-1, 16, 16))
+      return (d * sc.bitwise_and(0xF).unsqueeze(-1) * q - dmin * sc.rshift(4).unsqueeze(-1)).flatten(-2)
+    # Q3_K: 256 elements per 110-byte block (hmask:32, qs:64, scales:12, d:2)
+    if ggml_type == 11:
+      d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32).unsqueeze(-1)
+      sc = q_to_uint8(blocks[:,96:104], 4).bitwise_or(q_to_uint8(blocks[:,104:108], 2).lshift(4)).bitcast(dtypes.int8) - 32
+      q = q_to_uint8(blocks[:,32:96].reshape((-1, 2, 32)), 2).reshape((-1, 16, 16))
+      qh = q_to_uint8(blocks[:,:32], 1).reshape((-1, 16, 16))
+      return (d * sc.unsqueeze(-1) * (q.bitcast(dtypes.int8) - qh.bitwise_xor(1).lshift(2).bitcast(dtypes.int8))).flatten(-2)
      # Q4_K: 256 elements per 144-byte block (d:2, dmin:2, scales:12, qs:128)
      # Q5_K: 256 elements per 176-byte block (d:2, dmin:2, scales:12, qh:32, qs:128)
     if ggml_type in (12, 13):
@@ -79,6 +95,41 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       signs = (q_to_uint8(even_signs[sign_idx].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 8, 4, 8))
       grid = _ggml_iq_grid(t.device, _ggml.iq3xxs_grid, (256, 4))[blocks[:, 2:66]].reshape((-1, 8, 4, 8))
       return (db * grid * signs).flatten(-3)
+    # IQ2_XXS: 256 elements per 66-byte block (d:2, qs:64). 8 groups of 32: 4 grid bytes + packed signs/scale.
+    if ggml_type == 16:
+      d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
+      qs_u32 = blocks[:, 2:].bitcast(dtypes.uint32).reshape((-1, 8, 2))
+      db = d * (qs_u32[:, :, 1].rshift(28).cast(dtypes.float32) + 0.5).reshape((-1, 8, 1, 1)) * 0.25
+      sign_idx = qs_u32[:, :, 1].unsqueeze(-1).rshift(Tensor.const((0, 7, 14, 21), dtypes.uint32))
+      sign_idx = sign_idx.bitwise_and(0x7F).reshape((-1, 32)).cast(dtypes.int32)
+      even_signs = Tensor([i | (0x80 if i.bit_count() % 2 else 0) for i in range(128)], dtype=dtypes.uint8, device=t.device)
+      signs = (q_to_uint8(even_signs[sign_idx].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 8, 4, 8))
+      grid = _ggml_iq_grid(t.device, _ggml.iq2xxs_grid, (256, 8))[blocks[:, 2:].reshape((-1, 8, 8))[:, :, :4]].reshape((-1, 8, 4, 8))
+      return (db * grid * signs).flatten(-3)
+    # IQ2_XS: 256 elements per 74-byte block (d:2, qs:64 as uint16, scales:8)
+    if ggml_type == 17:
+      d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
+      db = d * (q_to_uint8(blocks[:, 66:74].reshape((-1, 8, 1)), 4).reshape((-1, 16)).cast(dtypes.float32) + 0.5).reshape((-1, 16, 1, 1)) * 0.25
+      qs = blocks[:, 2:66].bitcast(dtypes.uint16)
+      sign_idx = qs.rshift(9).cast(dtypes.int32)
+      even_signs = Tensor([i | (0x80 if i.bit_count() % 2 else 0) for i in range(128)], dtype=dtypes.uint8, device=t.device)
+      signs = (q_to_uint8(even_signs[sign_idx].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 16, 2, 8))
+      grid = _ggml_iq_grid(t.device, _ggml.iq2xs_grid, (512, 8))[qs.bitwise_and(511)].reshape((-1, 16, 2, 8))
+      return (db * grid * signs).flatten(-3)
+    # IQ1_S: 256 elements per 50-byte block (d:2, qs:32, qh:16). grid bytes are int8 {-1,0,1}.
+    if ggml_type == 19:
+      d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
+      qh = blocks[:, 34:50].bitcast(dtypes.uint16)
+      dl = d * (qh.rshift(12).bitwise_and(7).cast(dtypes.float32) * 2 + 1).reshape((-1, 8, 1, 1))
+      delta = (qh.bitwise_and(0x8000) == 0).where(0.125, -0.125).reshape((-1, 8, 1, 1))
+      qh_hi = qh.unsqueeze(-1).rshift(Tensor.const((0, 3, 6, 9), dtypes.uint16)).bitwise_and(7).lshift(8)
+      q = blocks[:, 2:34].cast(dtypes.uint16) + qh_hi.reshape((-1, 32))
+      grid = _ggml_iq_grid(t.device, _ggml.iq1s_grid, (2048, 8))[q].reshape((-1, 8, 4, 8))
+      grid = (grid > 127).where(grid - 256, grid)
+      return (dl * (grid + delta)).flatten(-3)
+    if ggml_type == 20:
+      d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32)
+      return d * Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device)[q_to_uint8(blocks[:, 2:], 4)]
     if ggml_type == 21:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
       scales = (1 + 2 * q_to_uint8(blocks[:, 106:110].reshape((-1, 4, 1)), 4).reshape((-1, 8))).cast(dtypes.float32).reshape((-1, 8, 1, 1))
@@ -102,6 +153,20 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       scales = (scales_l.bitwise_or(scales_h.lshift(4)).bitcast(dtypes.int8) - 32).cast(dtypes.float32).reshape((-1, 8, 1))
       q = (qs:=blocks[:, 8:].reshape((-1, 8, 16))).bitwise_and(0xF).cat(qs.rshift(4), dim=2)
       return (d * scales * iq4_xs_lut[q]).flatten(-2)
+    # IQ1_M: 256 elements per 56-byte block (qs:32, qh:16, scales:8). f16 scale packed in high nibbles.
+    if ggml_type == 29:
+      sc16 = blocks[:, 48:56].bitcast(dtypes.uint16)
+      d = sc16.bitwise_and(0xF000).rshift(Tensor.const((12, 8, 4, 0), dtypes.uint16))
+      d = d[:, 0:1].bitwise_or(d[:, 1:2]).bitwise_or(d[:, 2:3]).bitwise_or(d[:, 3:4])
+      d = d.bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1, 1))
+      scales = sc16.unsqueeze(-1).rshift(Tensor.const((0, 3, 6, 9), dtypes.uint16)).bitwise_and(7)
+      dl = d * (scales.cast(dtypes.float32) * 2 + 1).reshape((-1, 8, 2, 1, 1))
+      qh_n = Tensor.stack(blocks[:, 32:48].bitwise_and(0x0F), blocks[:, 32:48].rshift(4), dim=-1).reshape((-1, 32))
+      q = blocks[:, :32].cast(dtypes.uint16) + qh_n.bitwise_and(7).cast(dtypes.uint16).lshift(8)
+      delta = (qh_n.bitwise_and(0x08) == 0).where(0.125, -0.125).reshape((-1, 8, 2, 2, 1))
+      grid = _ggml_iq_grid(t.device, _ggml.iq1s_grid, (2048, 8))[q].reshape((-1, 8, 2, 2, 8))
+      grid = (grid > 127).where(grid - 256, grid)
+      return (dl * (grid + delta)).flatten(-4)
     if ggml_type == 39:
       e = blocks[:, 0].cast(dtypes.uint32)
       small_bits = Tensor([0x00200000, 0x00400000], dtype=dtypes.uint32, device=t.device)[e.clip(0, 1).cast(dtypes.int32)] # e = 0 or e = 1 case

tinygrad/llm/kernels/amd.py

@@ -116,7 +116,7 @@ def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
   if lanes is None: return ptr.load(arg="nontemporal")
   buf, coords = ptr.src[0], ptr.src[1:]
   idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0))
-  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load(dtype=ptr.dtype)
+  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load()
 
 def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
 def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()

tinygrad/mixin/gradient.py

@@ -31,6 +31,10 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
     return (None,) + k.arg.grad_fxn(on_dev(ctx, 0), k)
   assert fxn.op is Ops.TUPLE, f"expected TUPLE body for gradient, got {fxn.op}"
   params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
+  # grads are collected at the flat param storage: reshape to each arg's view (max view shrunk to symbolic)
+  def shaped_grad(grad:UOp, i:int) -> UOp:
+    a = args[i]
+    return grad.view_as(a.shard_shape, a.axis) if a.axis is not None and isinstance(a.device, tuple) else grad.view_as(a._shape)
   grad_args = ctx.src
   root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else
     g if g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
@@ -39,7 +43,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(fxn.src)} if k.arg.precompile else {}
   fwd_outs = tuple(k.gettuple(i) for i in range(len(fxn.src))) if k.arg.precompile else ()
   # collect needed gradient bodies, compact unused params, create a single backward CALL
-  grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
+  grad_bodies = [(i, shaped_grad(grads[p], i)) for i in needed if (p:=params.get(i)) is not None and p in grads]
   bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
   bwd_body = renumber_invalid_outputs(bwd_body)
   bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))

tinygrad/mixin/op.py

@@ -1067,7 +1067,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
       reshape[i] = expand[i] = size[i]
       if mode == "linear":
         arr = type(self).arange(size[i])
-        num, den = (arr*(in_sz-1), size[i]-1) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2)
+        num, den = (arr*(in_sz-1), max(size[i]-1, 1)) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2)
         num = num.clip(0, (in_sz-1)*den)
         low, high, perc = [y.reshape(reshape).expand(expand) for y in (num//den, (num+den-1)//den, (num % den).cast(dtypes.float32)/den)]
         x = x.gather(i, low).lerp(x.gather(i, high), perc)

tinygrad/nn/__init__.py

@@ -99,7 +99,7 @@ class Conv2d:
     if isinstance(padding, str):
       if padding.lower() != 'same': raise ValueError(f"Invalid padding string {padding!r}, only 'same' is supported")
       if stride != 1: raise ValueError("padding='same' is not supported for strided convolutions")
-      pad = [(d*(k-1)//2, d*(k-1) - d*(k-1)//2) for d,k in zip(make_tuple(dilation, len(self.kernel_size)), self.kernel_size[::-1])]
+      pad = [(d*(k-1)//2, d*(k-1) - d*(k-1)//2) for d,k in zip(make_tuple(dilation, len(self.kernel_size))[::-1], self.kernel_size[::-1])]
       padding = tuple(flatten(pad))
     self.stride, self.dilation, self.groups, self.padding = stride, dilation, groups, padding
     scale = 1 / math.sqrt(in_channels * prod(self.kernel_size))

tinygrad/nn/onnx.py

@@ -556,8 +556,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     return tuple(t if cond else e for t,e in zip(then_out.values(), else_out.values()))
 
   def Identity(x:Tensor): return x
-  def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:list[float]|None=None,
-              value_int:int|None=None, value_ints:list[int]|None=None, value_string:str|None=None, value_strings:list[str]|None=None):
+  def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:tuple[float, ...]|None=None,
+              value_int:int|None=None, value_ints:tuple[int, ...]|None=None, value_string:str|None=None, value_strings:tuple[str, ...]|None=None):
     if value is not None: return value
     if value_float is not None: return Tensor(value_float, dtype=dtypes.float32)
     if value_floats is not None: return Tensor(list(value_floats), dtype=dtypes.float32)
@@ -594,7 +594,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   # ***** Unary Ops (math) *****
   def Not(x:Tensor): return x.logical_not()
-  def Clip(x: Tensor, min:Tensor|None=None, max:Tensor|None=None): return x if min is None and max is None else x.clip(min, max)  # noqa: A002 # pylint: disable=redefined-builtin
+  def Clip(x: Tensor, min:Tensor|float|None=None, max:Tensor|float|None=None): return x if min is None and max is None else x.clip(min, max)  # noqa: A002 # pylint: disable=redefined-builtin
   def IsInf(x:Tensor, detect_negative:int=1, detect_positive:int=1): return x.isinf(bool(detect_positive), bool(detect_negative))
 
   # ***** Unary Ops (activation) *****
@@ -643,26 +643,26 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def Min(*data_0:Tensor): return functools.reduce(Tensor.minimum, data_0)
   def Sum(*data_0:Tensor): return functools.reduce(Tensor.add, data_0)
   def Mean(*data_0:Tensor): return Sum(*data_0) / len(data_0)
-  def ReduceMax(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceMax(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     return data.max(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceMin(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceMin(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     return data.min(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceSum(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     return data.sum(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceMean(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceMean(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     return data.mean(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceSumSquare(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceSumSquare(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     return ReduceSum(data.square(), axes, keepdims, noop_with_empty_axes)
-  def ReduceProd(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceProd(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     return data.prod(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceL1(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceL1(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     return ReduceSum(data.abs(), axes, keepdims, noop_with_empty_axes)
-  def ReduceL2(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceL2(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     dtype = dtypes.float if data.dtype in (dtypes.float16, dtypes.bfloat16) else data.dtype
     return ReduceSum(data.cast(dtype).square(), axes, keepdims, noop_with_empty_axes).sqrt().cast(data.dtype)
-  def ReduceLogSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceLogSum(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     return ReduceSum(data, axes, keepdims, noop_with_empty_axes).log()
-  def ReduceLogSumExp(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceLogSumExp(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
     return ReduceSum(data.exp(), axes, keepdims, noop_with_empty_axes).log()
   def ArgMax(x:Tensor, axis:int=0, keepdims:int=1, select_last_index:int=0):
     if select_last_index: return ((int(x.shape[axis])-1) - x.flip(axis).argmax(axis, keepdim=keepdims)).cast(dtypes.int64)
@@ -671,32 +671,32 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     return ArgMax(-x, axis=axis, keepdims=keepdims, select_last_index=select_last_index)
 
   # ***** Movement Ops *****
-  def Reshape(data:Tensor, shape:list[int], allowzero:int=0):
+  def Reshape(data:Tensor, shape:Sequence[int], allowzero:int=0):
     return data.reshape([x if x != 0 else (0 if allowzero else data.shape[i]) for i,x in enumerate(shape)])
   def Flatten(x:Tensor, axis:int=1): return x.reshape(prod(x.shape[0:axis]), -1)
   def Expand(x:Tensor, shape:list[int]): return x.expand(_broadcast_shape(x.shape, tuple(shape)))
   def Shrink(x:Tensor, bias:float=0.0, lambd:float=0.5): return (x < -lambd)*(x+bias) + (x > lambd)*(x-bias)
-  def Transpose(x:Tensor, perm:list[int]|None=None): return x.permute(order=perm or list(range(x.ndim)[::-1]))
+  def Transpose(x:Tensor, perm:tuple[int, ...]|None=None): return x.permute(order=perm or list(range(x.ndim)[::-1]))
 
-  def Squeeze(data:Tensor, axes:list[int]|None=None):
+  def Squeeze(data:Tensor, axes:Sequence[int]|None=None):
     return data.squeeze() if axes is None else functools.reduce(lambda d, dim: d.squeeze(dim), sorted(axes, reverse=True), data)
-  def Unsqueeze(data:Tensor, axes:list[int]): return functools.reduce(lambda d, dim: d.unsqueeze(dim), sorted(axes), data)
+  def Unsqueeze(data:Tensor, axes:Sequence[int]): return functools.reduce(lambda d, dim: d.unsqueeze(dim), sorted(axes), data)
 
   def Tile(x:Tensor, repeats:list[int]): return x.repeat(repeats)
   def Concat(*xs:Tensor, axis:int): return Tensor.cat(*xs, dim=axis)
-  def Slice(data:Tensor, starts:list[int], ends:list[int], axes:list[int]|None=None, steps:list[int]|None=None):
+  def Slice(data:Tensor, starts:Sequence[int], ends:Sequence[int], axes:Sequence[int]|None=None, steps:list[int]|None=None):
     axes = axes or list(range(data.ndim))
     steps = steps or [1] * data.ndim
     slices = [slice(None)] * data.ndim
     for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i])
     return data[tuple(slices)]
 
-  def Split(data:Tensor, split:list[int]|None=None, num_outputs:int=0, axis:int=0):
+  def Split(data:Tensor, split:Sequence[int]|None=None, num_outputs:int=0, axis:int=0):
     sz = int(data.shape[axis])
     if split is None: split = [sz // num_outputs + (1 if i < sz % num_outputs else 0) for i in range(num_outputs)]
     return data.split(split, axis)
 
-  def Pad(x:Tensor, pads:list[int], constant_value:ConstType|None=None, axes:list[int]|None=None,
+  def Pad(x:Tensor, pads:Sequence[int], constant_value:ConstType|None=None, axes:list[int]|None=None,
           mode:Literal["constant", "reflect", "edge", "wrap"]="constant", value=0):
     value = _resolve_const(value if constant_value is None else constant_value)
     axes = axes or list(range(x.ndim))
@@ -704,7 +704,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     for i,axis in enumerate(axes): real_pads[axis%x.ndim], real_pads[axis%x.ndim+x.ndim] = pads[i], pads[i+len(axes)]
     return x.pad(padding=_onnx_pads_to_tiny_pads(real_pads), mode={"edge":"replicate", "wrap":"circular"}.get(mode, mode), value=value)
 
-  def CenterCropPad(t:Tensor, shape:list[int], axes:list[int]|None=None):
+  def CenterCropPad(t:Tensor, shape:list[int], axes:tuple[int, ...]|None=None):
     shrink_arg:list[None|tuple[sint,sint]] = [None] * t.ndim
     pad_arg:list[None|tuple[sint,sint]] = [None] * t.ndim
     for s, x in zip(shape, axes or range(t.ndim)):
@@ -714,26 +714,26 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     return t.shrink(tuple(shrink_arg)).pad(tuple(pad_arg))
 
   # ***** Processing Ops *****
-  def AveragePool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, count_include_pad:int=0,
-                  dilations:list[int]|int=1, pads:list[int]|int=0, strides:list[int]|int=1):
+  def AveragePool(X: Tensor, kernel_shape:tuple[int, ...], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, count_include_pad:int=0,
+                  dilations:tuple[int, ...]|int=1, pads:tuple[int, ...]|int=0, strides:tuple[int, ...]|int=1):
     pool_pads = _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad)
     return X.avg_pool2d(tuple(kernel_shape), strides, dilations, pool_pads, ceil_mode=ceil_mode, count_include_pad=count_include_pad)
 
-  def MaxPool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, dilations:list[int]|int=1, pads:list[int]|int=0,
-              storage_order:int=0, strides:list[int]|int=1):
+  def MaxPool(X: Tensor, kernel_shape:tuple[int, ...], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, dilations:tuple[int, ...]|int=1,
+              pads:tuple[int, ...]|int=0, storage_order:int=0, strides:tuple[int, ...]|int=1):
     pool_pads = _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad)
     out = X.max_pool2d(tuple(kernel_shape), strides, dilations, pool_pads, ceil_mode=ceil_mode, return_indices=True)
     ret, idx = cast(tuple[Tensor, Tensor], out)
     return ret, idx.transpose(-2, -1).cast(dtypes.int64) if storage_order else idx.cast(dtypes.int64)
 
-  def Conv(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
-          kernel_shape:list[int]|None=None, pads:list[int]|int=0, strides:list[int]|int=1):
+  def Conv(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:tuple[int, ...]|int=1, group:int=1,
+          kernel_shape:tuple[int, ...]|None=None, pads:tuple[int, ...]|int=0, strides:tuple[int, ...]|int=1):
     return X.conv2d(W, B, stride=strides, groups=group, dilation=dilations,
                     padding=_resolve_pool_pads(X, pads, kernel_shape or W.shape[2:], dilations, strides, auto_pad))
 
-  def ConvTranspose(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
-                    kernel_shape:list[int]|None=None, pads:list[int]|None=None, output_shape:list[int]|None=None, output_padding:list[int]|int=0,
-                    strides:list[int]|int=1):
+  def ConvTranspose(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:tuple[int, ...]|int=1, group:int=1,
+                    kernel_shape:tuple[int, ...]|None=None, pads:Sequence[int]|None=None, output_shape:Sequence[int]|None=None,
+                    output_padding:tuple[int, ...]|int=0, strides:tuple[int, ...]|int=1):
     input_shape_, kernel_shape_ = X.shape[2:], (kernel_shape or W.shape[2:])
     strides_, dilations_, output_padding_ = (make_tuple(x, len(input_shape_)) for x in (strides, dilations, output_padding))
     if output_shape is not None: # we pad according to output_shape
@@ -747,10 +747,10 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     pads = _onnx_pads_to_tiny_pads(pads)
     return X.conv_transpose2d(W, B, group, strides_, dilations_, pads, output_padding_)
 
-  def MaxUnpool(xT: Tensor, xI: Tensor, outshape: list[int]|None=None, kernel_shape:list[int]|None=None, pads:list[int]|int=0,
-                strides:list[int]|int=1):
+  def MaxUnpool(xT: Tensor, xI: Tensor, outshape: list[int]|None=None, kernel_shape:Sequence[int]|None=None, pads:tuple[int, ...]|int=0,
+                strides:tuple[int, ...]|int=1):
     if kernel_shape is None: kernel_shape = []
-    pads_: int | tuple[int, ...] = tuple(pads) if isinstance(pads, list) else pads
+    pads_: int | tuple[int, ...] = pads if isinstance(pads, int) else _onnx_pads_to_tiny_pads(pads)
     return Tensor.max_unpool2d(xT, xI, tuple(kernel_shape), strides, 1, pads_, outshape if outshape is None else tuple(outshape))
 
   def GlobalAveragePool(X:Tensor): return X.mean(axis=tuple(range(2, X.ndim)), keepdim=True)
@@ -761,7 +761,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     if C is not None: ret = ret + beta * (C if broadcast == 0 else C.reshape([-1 if i < len(C.shape) else 1 for i in range(ret.ndim)][::-1]))
     return ret
 
-  def Einsum(*Inputs:list[Tensor], equation:str): return Tensor.einsum(equation, *Inputs)
+  def Einsum(*Inputs:Tensor, equation:str): return Tensor.einsum(equation, *Inputs)
 
   def CumSum(X:Tensor, axis:int|list[int], exclusive:int=0, reverse:int=0):
     axis = X._resolve_dim(_resolve_const(axis))
@@ -774,8 +774,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     k_ = _resolve_const(k)
     return x.triu(k_) if upper else x.tril(k_)
 
-  def Resize(X:Tensor, roi:list[float]|None=None, scales:list[float]|None=None, sizes:list[int]|None=None, antialias:int=0,
-        axes:list[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0,
+  def Resize(X:Tensor, roi:list[float]|None=None, scales:Sequence[float]|None=None, sizes:list[int]|None=None, antialias:int=0,
+        axes:Sequence[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0,
         extrapolation_value:float=0.0, keep_aspect_ratio_policy:str='stretch', mode:str='nearest', nearest_mode:str='round_prefer_floor'):
     def _apply_transformation(input_sz, output_sz, scale_dim, mode):
       index = Tensor.arange(output_sz)
@@ -876,7 +876,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
         gathered_values = [X.gather(i, idx) for idx in expanded_indices]
         X = sum(v * c for v, c in zip(gathered_values, expanded_coeffs))
     return X.permute(*argsort(perm)) if perm else X
-  def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode)  # deprecated
+  def Upsample(X:Tensor, scales:Sequence[float], mode:str): return Resize(X=X, scales=scales, mode=mode)  # deprecated
 
   def TopK(X:Tensor, K:int|list[int], axis:int=-1, largest:int=1, sorted:int=1):  # noqa: A002 # pylint: disable=redefined-builtin
     val, idx = X.topk(_resolve_const(K), axis, bool(largest), bool(sorted))
@@ -937,9 +937,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     if segment_embedding is not None: embedding_sum = embedding_sum + embedding(segment_ids, segment_embedding.shape[0], segment_embedding)
     out = embedding_sum.layernorm(eps=epsilon) * gamma + beta
     return out, None, embedding_sum
-  def MeanVarianceNormalization(x:Tensor, axis:list[int]|None=None):
-    if axis is None: axis = [0,2,3]
-    return (x - x.mean(axis, keepdim=True)) / (x.std(axis, keepdim=True, correction=0) + 1e-9)
+  def MeanVarianceNormalization(x:Tensor, axes:Sequence[int]=(0,2,3)):
+    return (x - x.mean(axes, keepdim=True)) / (x.std(axes, keepdim=True, correction=0) + 1e-9)
 
   def LpNormalization(x:Tensor, axis:int=-1, p:int=2):
     return x / (x.abs().sum(axis, keepdim=True) if p == 1 else x.square().sum(axis, keepdim=True).sqrt())
@@ -1001,7 +1000,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None,
                         attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None, do_rotary:int=0, mask_filter_value:float=-10000.0,
-                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
+                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:Sequence[int]|None=None,
                         rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
     assert not do_rotary and not attention_bias, "TODO"
     if qkv_hidden_sizes is None: qkv_hidden_sizes = [int(weights.shape[1] // 3)] * 3
@@ -1246,8 +1245,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     ret = _qlinearop_float(GlobalAveragePool, [X], [x_zero_point], [x_scale], y_scale, y_zero_point)
     return ret.permute(0, *range(2, ret.ndim), 1) if channels_last else ret  # NCHW -> NHWC
 
-  def ConvInteger(x: Tensor, w: Tensor, x_zero_point:Tensor = Tensor(0), w_zero_point:Tensor = Tensor(0), B: Tensor | None = None, **opts) -> Tensor:
-    return _op_integer(Conv, [x,w], [x_zero_point,w_zero_point], **{"B":B, **opts})
+  def ConvInteger(x: Tensor, w: Tensor, x_zero_point:Tensor = Tensor(0), w_zero_point:Tensor = Tensor(0), **opts) -> Tensor:
+    return _op_integer(Conv, [x,w], [x_zero_point,w_zero_point], **opts)
 
   def MatMulInteger(A: Tensor, B: Tensor, a_zero_point: Tensor = Tensor(0), b_zero_point: Tensor = Tensor(0)) -> Tensor:
     return _op_integer(Tensor.matmul, [A,B], [a_zero_point,b_zero_point])

tinygrad/renderer/__init__.py

@@ -9,8 +9,8 @@ from tinygrad.device import Compiler
 
 # an access takes its dtype from the buffer it indexes, so accessing at another dtype restates the storage on the buffer that owns it
 def with_storage(x:UOp, dt:DType) -> UOp:
-  if x.op in {Ops.PARAM, Ops.BUFFER}: return x.replace(dtype=None, arg=replace(x.arg, dtype=dt))
-  return x.replace(dtype=None, src=(with_storage(x.src[0], dt),)+x.src[1:])
+  if x.op in {Ops.PARAM, Ops.BUFFER}: return x.replace(arg=replace(x.arg, dtype=dt))
+  return x.replace(src=(with_storage(x.src[0], dt),)+x.src[1:])
 
 @dataclass(frozen=True)
 class Estimates:

tinygrad/renderer/isa/__init__.py

@@ -27,7 +27,8 @@ class IselContext:
     return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,))
 
 def greg(u:UOp):
-  if u.op in {Ops.NOOP, Ops.AFTER} and u.src: return greg(u.src[0])
+  # a BITCAST that no instruction lowered is bit identity, it is an alias of the register its src is in
+  if u.op in {Ops.NOOP, Ops.AFTER, Ops.BITCAST} and u.src: return greg(u.src[0])
   if isinstance(u.tag, tuple): return u.tag[0]
   return u.tag
 

tinygrad/renderer/isa/x86.py

@@ -2,9 +2,10 @@
 # allow semicolons to put multiple ops on one line
 import sys, struct, functools
 from typing import cast
+from dataclasses import replace
 from tinygrad.dtype import dtypes, DType, truncate, AddrSpace
 from tinygrad.uop import FastEnum, auto, Ops, GroupOp
-from tinygrad.uop.ops import UOp, UPat, PatternMatcher
+from tinygrad.uop.ops import UOp, UPat, PatternMatcher, promo_dtype
 from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg
 from tinygrad.helpers import getenv, NUM_CPU_THREADS, unwrap, Target
 
@@ -33,14 +34,11 @@ class X86Ops(FastEnum):
   # bitcasts
   VMOVD = auto(); VMOVQ = auto(); VMOVDm = auto(); VMOVQm = auto()
   # comparisons
-  VUCOMISS = auto(); VUCOMISD = auto()
   VCMPSS = auto(); VCMPSD = auto(); VCMPPS = auto(); VCMPPD = auto()
-  VPCMPGTB = auto(); VPCMPGTW = auto(); VPCMPGTD = auto(); VPCMPGTQ = auto()
-  VPCMPEQB = auto(); VPCMPEQW = auto(); VPCMPEQD = auto(); VPCMPEQQ = auto()
   SETNE = auto(); SETE = auto(); SETL = auto(); SETB = auto()
   # where
   CMOVNE = auto(); CMOVE = auto(); CMOVL = auto(); CMOVB = auto()
-  VPBLENDVB = auto(); VBLENDVPS = auto(); VBLENDVPD = auto()
+  VBLENDVPS = auto(); VBLENDVPD = auto()
   # jumps
   JNE = auto(); JE = auto(); JL = auto(); JB = auto(); JGE = auto(); JMP = auto()
   # vectorize / gep
@@ -90,12 +88,11 @@ class X86GroupOp:
                 X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VADDPS, X86Ops.VADDPD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VSUBPS, X86Ops.VSUBPD,
                 X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VMULPS, X86Ops.VMULPD, X86Ops.VDIVSS, X86Ops.VDIVSD, X86Ops.VDIVPS, X86Ops.VDIVPD,
                 X86Ops.VPADDB, X86Ops.VPADDW, X86Ops.VPADDD, X86Ops.VPADDQ, X86Ops.VPSUBB, X86Ops.VPSUBW, X86Ops.VPSUBD, X86Ops.VPSUBQ,
-                X86Ops.VPCMPEQB, X86Ops.VPCMPEQW, X86Ops.VPCMPEQD, X86Ops.VPCMPEQQ, X86Ops.VPBLENDVB, X86Ops.VBLENDVPS, X86Ops.VBLENDVPD,
-                X86Ops.VPCMPGTB, X86Ops.VPCMPGTW, X86Ops.VPCMPGTD, X86Ops.VPCMPGTQ, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VCMPPS, X86Ops.VCMPPD,
+                X86Ops.VBLENDVPS, X86Ops.VBLENDVPD, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VCMPPS, X86Ops.VCMPPD,
                 X86Ops.VPMULLW, X86Ops.VPMULLD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD, X86Ops.VINSERTPS,
                 X86Ops.VPINSRB, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.VPINSRQ, X86Ops.VPAND, X86Ops.VPOR, X86Ops.VPXOR, X86Ops.VPSLLVD,
                 X86Ops.VPSLLVQ, X86Ops.VPSRLVD, X86Ops.VPSRLVQ, X86Ops.VPSRAVD, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB,
-                X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.VUCOMISS, X86Ops.VUCOMISD, X86Ops.IDIV, X86Ops.DIV}
+                X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.IDIV, X86Ops.DIV}
 
   # X86Ops that can write to memory
   WriteMem = {X86Ops.MOVm, X86Ops.MOVi, X86Ops.VMOVSSm, X86Ops.VMOVSDm, X86Ops.VMOVUPSm, X86Ops.VMOVDm, X86Ops.VMOVQm,
@@ -110,7 +107,7 @@ class X86GroupOp:
   # X86Ops that write flags or can modify flags to undefined values
   WriteFlags = {X86Ops.CMP, X86Ops.CMPi, X86Ops.ADD, X86Ops.ADDi, X86Ops.SUB, X86Ops.SUBi, X86Ops.IMUL, X86Ops.IMULi, X86Ops.IDIV, X86Ops.DIV,
                 X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi,
-                X86Ops.OR, X86Ops.ORi, X86Ops.VUCOMISS, X86Ops.VUCOMISD}
+                X86Ops.OR, X86Ops.ORi}
 
   # X86Ops whose first src is the rm field
   Rm1st = ReadMem1st | (ReadMem2nd & TwoAddress) | {X86Ops.VPSRLDQ}
@@ -118,8 +115,6 @@ class X86GroupOp:
   # X86Ops whose second src is the rm field
   Rm2nd = ReadMem2nd
 
-  All = set(X86Ops)
-
 # ***** X86 legalization *****
 
 extra_matcher = PatternMatcher([
@@ -145,14 +140,14 @@ extra_matcher = PatternMatcher([
   # float16 alus are done in float32
   (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op,
    src=tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
-  (UPat(GroupOp.Comparison, src=(UPat.var("a", dtypes.float16), UPat.var("b")), name="x"),
-   lambda x,a,b: UOp(x.op, src=(a.cast(dtypes.float32), b.cast(dtypes.float32))).cast(x.dtype)),
+  (UPat(GroupOp.Comparison, src=[UPat(dtype=dtypes.float16), UPat()], name="x"),
+   lambda x: UOp(x.op, src=tuple(s.cast(dtypes.float32) for s in x.src)).cast(x.dtype)),
   # no cmpne for packed ints, y != x => !(y==x)
   (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"),
    lambda y,x,cmp: UOp(Ops.CMPEQ, src=(y,x))^True if y.max_numel() > 1 else None),
-  # float WHERE needs a mask unless its comparison already has a float operand
+  # a float WHERE blends at the width of its value, so it needs a comparison at that width to make the mask
   (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats+(dtypes.weakfloat,)), UPat.var("b")).named("w"),
-   lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and not dtypes.is_float(m.src[0].dtype) else None),
+   lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and promo_dtype(m.src) is not w.dtype else None),
   # rewrite -x -> 0 - x
   (UPat(Ops.NEG, name="x"), lambda x: UOp(Ops.SUB, src=(x.const_like(0),) + x.src)),
   # TODO: add support for mod, requires support for accessing the 2nd+ reg of a multi output instruction
@@ -166,36 +161,35 @@ 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.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)))
-  return ptr.load(dtype=x.dtype)
+  local_idx = local.index(UOp.cconst(0, dtypes.int32))
+  # the gate picks one of two places to access, the AFTER orders the access after the scratch store
+  sel = gate.where(addr, local_idx)
+  return UOp(Ops.AFTER, (sel, (local_idx if x.max_numel() == 1 else local).store(alt))).load()
 
 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.cconst(0, dtypes.int32), dtype=dtypes.uint64))
-  return UOp(Ops.AFTER, addr.dtype, (sel,)).store(val)
+  sel = gate.where(addr, local.index(UOp.cconst(0, dtypes.int32)))
+  return UOp(Ops.AFTER, (sel,)).store(val)
+
+# a gate the flags can be picked with, or the bool compared to zero that replaces one they can't: only an integer
+# comparison sets the flags, see cmp. NOTE: the 0 is int so the bool zero-extends and compares as int (a byte compare renders
+# different kernels)
+def flag_gate(m:UOp) -> UOp|None:
+  return None if m.op in GroupOp.Comparison and m.src[0].dtype not in dtypes.floats else m.ne(UOp.cconst(0, dtypes.int))
 
 # legalize the new style graph for isel. NOTE: this runs after the spec is verified, some of these rewrites violate it
 pre_isel_matcher = PatternMatcher([
-  # noop casts: zero extending scalar 32bit int, same-width signed/unsigned, narrowing scalar int
-  (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.NOOP, arg=None) if y.max_numel() == 1 else None),
+  # widening a scalar uint32 is free, the 32bit write that produced it already zeroed the upper half
+  (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.BITCAST) if y.max_numel() == 1 else None),
+  # a cast between same width ints keeps every bit (a bool is a 0/1 byte), the honest op is a bitcast
   (UPat.var("y", dtypes.ints+(dtypes.bool,)).cast(dtypes.ints, name="x"),
-   lambda y,x: x.replace(op=Ops.NOOP, arg=None) if x.dtype.itemsize == y.dtype.itemsize and y.max_numel() == 1 else None),
-  (UPat.var("y", dtypes.ints).cast(dtypes.ints, name="x"),
-   lambda y,x: x.replace(op=Ops.NOOP, arg=None) if x.dtype.itemsize < y.dtype.itemsize and y.max_numel() == 1 else None),
-  # bitcasts between scalar floats and ints are real, rest are noops
-  (UPat.var("y").bitcast().named("x"), lambda y,x: None if y.dtype in dtypes.floats and x.dtype in dtypes.ints or \
-   y.dtype in dtypes.ints and x.dtype in dtypes.floats else x.replace(op=Ops.NOOP, arg=None)),
+   lambda y,x: x.replace(op=Ops.BITCAST) if x.dtype.itemsize == y.dtype.itemsize else None),
   # gated load/store become a conditional move on the address, the load/store are unconditional
   (UPat((Ops.INDEX, Ops.SHRINK), name="addr").load(UPat.var("alt"), UPat.var("gate"), name="x"), gated_load),
   (UPat((Ops.INDEX, Ops.SHRINK), name="addr").store(UPat.var("val"), UPat.var("gate")), gated_store),
-  # TODO: remove this once we allow all flag producing ops in cmove
-  # 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.cconst(0, dtypes.int)).where(a,b) if m.op not in GroupOp.Comparison else None),
+  # a conditional backedge picks with the flags, and so does the cmove, which is legalized in isel
+  (UPat(Ops.END, src=(UPat(), UPat(), UPat.var("m", dtypes.bool)), name="x"),
+   lambda m,x: x.replace(src=x.src[:2]+(g,)) if (g:=flag_gate(m)) is not None else None),
 ])
 
 # ***** X86 registers *****
@@ -217,27 +211,29 @@ CALLEE_SAVED = (RBX, RBP, GPR[12], GPR[13], GPR[14], GPR[15]) + ((RSI, RDI) + XM
 
 reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"}, "rdx": {4:"edx", 2:"dx", 1:"dl"}, "rbx": {4:"ebx", 2:"bx", 1:"bl"},
         "rsp": {4:"esp", 2:"sp", 1:"spl"}, "rbp": {4:"ebp", 2:"bp", 1:"bpl"}, "rsi": {4:"esi", 2:"si", 1:"sil"}, "rdi": {4:"edi", 2:"di", 1:"dil"},
-        **{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}, **{f"xmm{i}": {64:f"zmm{i}", 32:f"ymm{i}"} for i in range(16)}}
+        **{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}}
 
 # ***** 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].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 def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, arg=(X86Ops.DEFINE, dt), tag=None if reg is None else (reg,))
 def imm(dt:DType, v:int) -> UOp: return UOp.cconst(truncate[dt](v), dt).rtag()
 def to_imm(c:UOp) -> UOp|None:
   if not (c.op is Ops.CAST and (v:=c.src[0]).op is Ops.CONST): return None
   if c.dtype in dtypes.int64s: return imm(dtypes.int32, v.val) if not v.overflows(dtypes.int32) else None
   if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, v.val)
   return None
+# the flag path, which only an integer comparison can take: an x86 float compare sets carry, zero and parity together when an
+# operand is NaN, so a NaN reads as "below" and as "equal", and it clears sign and overflow, so nothing reads as "less"
 def cmp(x:UOp) -> UOp:
-  if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void)
-  if x.src[0].dtype is dtypes.float64: return x.ins(X86Ops.VUCOMISD, dtype=dtypes.void)
+  if x.src[0].dtype in dtypes.floats: raise RuntimeError(f"no flag compare for {x.src[0].dtype}, a float gate must be a mask")
   return x.ins(X86Ops.CMP, dtype=dtypes.void) if (i:=to_imm(x.src[1])) is None else x.ins(X86Ops.CMPi, dtype=dtypes.void, src=(x.src[0], i))
-def vcmp(x:UOp) -> UOp:
-  v = imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
-  if x.dtype is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, src=x.src + (v,))
-  return x.ins(X86Ops.VCMPSD if x.max_numel() == 1 else X86Ops.VCMPPD, src=x.src + (v,))
+# a float comparison consumed as a mask writes all ones in the width of its operands, the instruction states that width. a bool comparison can't
+def mask(x:UOp) -> UOp:
+  dt, v = x.src[0].dtype, imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
+  if dt is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, dtype=dt, src=x.src + (v,))
+  return x.ins(X86Ops.VCMPSD if x.max_numel() == 1 else X86Ops.VCMPPD, dtype=dt, src=x.src + (v,))
 
 # vinsertps xmm2, xmm0, xmm1, imm
 # inserts any 32 bit element in xmm1 into any position in xmm0 according to immm, result is written to xmm2
@@ -250,9 +246,9 @@ def vinsertps(x:UOp) -> UOp:
 
 # vpinsq xmm2, xmm0, rax, imm
 # inserts element in rax into any position in xmm0, result is written to xmm2 according to imm
-def vpins(x:UOp) -> UOp:
+def vpins(x:UOp, srcs:tuple[UOp, ...]) -> UOp:
   op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.itemsize]
-  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, x.src[i], imm(dtypes.uint8, i))), range(len(x.src)), def_reg(x.dtype))
+  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, srcs[i], imm(dtypes.uint8, i))), range(len(srcs)), def_reg(x.dtype))
 
 # we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg
 def idiv(ctx:IselContext, x:UOp) -> UOp:
@@ -262,8 +258,8 @@ def idiv(ctx:IselContext, x:UOp) -> UOp:
   elif x.dtype in dtypes.uints: ext = [x.ins(X86Ops.MOVi, src=(imm(min(dtypes.uint32, x.dtype), 0),), tag=(RDX,))]
   else: ext = [x.ins(X86Ops.SARi, src=(x.src[0], imm(dtypes.uint8, x.dtype.itemsize * 8 - 1)), tag=(RDX,))]
   # for 8bit need to zero/sign extend al to ah
-  if x.dtype is dtypes.uint8: dividend = UOp(Ops.INS, arg=X86Ops.MOVZX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,))
-  elif x.dtype is dtypes.int8: dividend = UOp(Ops.INS, arg=X86Ops.MOVSX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,))
+  if x.dtype is dtypes.uint8: dividend = UOp(Ops.INS, arg=(X86Ops.MOVZX, dtypes.int16), src=(x.src[0],), tag=(RAX,))
+  elif x.dtype is dtypes.int8: dividend = UOp(Ops.INS, arg=(X86Ops.MOVSX, dtypes.int16), src=(x.src[0],), tag=(RAX,))
   else: dividend = x.ins(X86Ops.MOV, src=(x.src[0],), tag=(RAX,))
   # divisor can't be in rax or rdx
   divisor = x.ins(X86Ops.MOV, src=(x.src[1],), tag=tuple(r for r in WGPR if r not in (RAX, RDX)))
@@ -282,7 +278,7 @@ def shift(x:UOp, op:X86Ops) -> UOp:
 # it is materialized as an immediate so the address stays correct if the base register is ever spilled and refilled
 def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]:
   def _disp(v:int) -> UOp: return imm(dtypes.int32 if abs(v) > dtypes.int8.max else dtypes.int8, v)
-  def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v
+  def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v.cast(dtypes.uint32) if v.dtype.itemsize < 4 else v
   if x.op not in {Ops.INDEX, Ops.SHRINK}: return (x, UOp(Ops.NOOP), _disp(0), imm(dtypes.uint8, x.dtype.itemsize))
   base, idx = x.src[0], x.src[1]
   # buffers are indexed by element, everything else (the stack pointer) by byte
@@ -293,15 +289,22 @@ def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]:
   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)
 
+# the value of a place is its address, computed into a 64bit register
+def lea(x:UOp) -> UOp: return x.ins(X86Ops.LEA, dtype=dtypes.uint64, src=fold_address(x)) if x.op in {Ops.INDEX, Ops.SHRINK} else x
+def where_places(m:UOp, a:UOp, b:UOp) -> UOp|None:
+  p = a if a.op in {Ops.INDEX, Ops.SHRINK} else b if b.op in {Ops.INDEX, Ops.SHRINK} else None
+  return m.where(lea(a), lea(b)) if p is not None and not _is_vec_xmm(p.src[0]) else None
+
 def abi(ctx:IselContext, x:UOp) -> UOp|None:
   if isinstance(x.tag, tuple): return None
   i = ctx.func_args.index(x)
-  # buffer params hold addresses, their value moves as a 64bit int
+  # a buffer param holds an address, the abi passes it as a 64bit int and a ParamArg owns the dtype its param is passed as
   dt = dtypes.uint64 if x.op is Ops.PARAM and x.arg.addrspace is AddrSpace.GLOBAL else x.dtype
+  arg = replace(x.arg, dtype=dt) if x.op is Ops.PARAM else x.arg
   # the shape srcs of a PARAM are not values, tag them so they aren't materialized into registers
-  def _reg_arg(r:Register) -> tuple[UOp, ...]: return (x.replace(dtype=dt, src=tuple(s.rtag() for s in x.src), tag=(r,)),)
+  def _reg_arg(r:Register) -> tuple[UOp, ...]: return (x.replace(arg=arg, src=tuple(s.rtag() for s in x.src), tag=(r,)),)
   def _stack_arg(disp:int):
-    return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=X86Ops.FRAME_INDEX, dtype=dtypes.int32, tag=disp), imm(dtypes.uint8, 8))
+    return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=(X86Ops.FRAME_INDEX, dtypes.int32), tag=disp), imm(dtypes.uint8, 8))
   if sys.platform == "win32": src = _reg_arg((RCX, RDX, GPR[8], GPR[9])[i]) if i < 4 else _stack_arg((i-3)*8+32)
   else: src = _reg_arg((RDI, RSI, RDX, RCX, GPR[8], GPR[9])[i]) if i < 6 else _stack_arg((i-5)*8)
   # this move "cleanses" the abi register constraint
@@ -312,7 +315,7 @@ GPR_DEST_OPS = {X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ,
 XMM_OPS = {op for op in X86Ops if op.name.startswith('V')} - GPR_DEST_OPS
 
 def _is_vec_xmm(y: UOp) -> bool:
-  return (y.op is Ops.INS and y.arg in XMM_OPS) or (y.op not in (Ops.BUFFER, Ops.PARAM, Ops.AFTER, Ops.INS) and y.max_numel() > 1)
+  return (y.op is Ops.INS and y.arg[0] in XMM_OPS) or (y.op not in (Ops.BUFFER, Ops.PARAM, Ops.AFTER, Ops.INS) and y.max_numel() > 1)
 
 def _xmm_sz(x: UOp) -> X86Ops:
   bits = x.max_numel() * x.dtype.itemsize
@@ -328,10 +331,10 @@ def _xmm_sz_m(x: UOp) -> X86Ops:
 
 def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None:
   # register placeholders with real registers
-  if x.arg is X86Ops.DEFINE and x.tag is not None: return None
-  if x.arg is X86Ops.LOOP_CMP: return None
+  if x.op is Ops.INS and x.arg[0] is X86Ops.DEFINE and x.tag is not None: return None
+  if x.op is Ops.INS and x.arg[0] is X86Ops.LOOP_CMP: return None
   # this is an immediate
-  if x.arg is X86Ops.FRAME_INDEX: return None
+  if x.op is Ops.INS and x.arg[0] is X86Ops.FRAME_INDEX: return None
   # no register definition
   if x.dtype is dtypes.void: return None
   # already allocated vregs
@@ -340,18 +343,16 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None:
   defs = []
   if isinstance(x.tag, tuple): defs = [ctx.vreg(x.tag)]
   elif x.op is Ops.BUFFER: defs = [ctx.vreg(WGPR)]
-  elif x.dtype in dtypes.floats or (x.op is Ops.INS and x.arg in XMM_OPS) or x.max_numel() > 1: defs = [ctx.vreg(XMM)]
+  elif x.dtype in dtypes.floats or (x.op is Ops.INS and x.arg[0] in XMM_OPS) or x.max_numel() > 1: defs = [ctx.vreg(XMM)]
   elif x.dtype in dtypes.ints+(dtypes.bool,): defs = [ctx.vreg(WGPR)]
   # TODO: add this once the scheduler can track register pressure
-  # if x.arg in X86GroupOp.WriteFlags: defs.append(ctx.vreg(RFLAGS))
+  # if x.arg[0] in X86GroupOp.WriteFlags: defs.append(ctx.vreg(RFLAGS))
   # the size src of a BUFFER is not a value, tag it so it isn't materialized into a register
   if x.op is Ops.BUFFER: return x.replace(src=tuple(s.rtag() for s in x.src), tag=tuple(defs))
   return x.replace(tag=tuple(defs))
 
 isel_matcher = PatternMatcher([
   # **** Op -> Op ****
-  # 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").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),
@@ -363,25 +364,28 @@ isel_matcher = PatternMatcher([
   # so regalloc builds the prologue/epilogue naturally
   (UPat(Ops.SINK, name="x"), lambda x:
    x.replace(src=(x.ins(X86Ops.RET, src=x.src + tuple(def_reg(dtypes.uint64 if r in GPR else dtypes.float64, r) for r in CALLEE_SAVED)),)) \
-    if not x.src or x.src[0].arg is not X86Ops.RET else None),
+    if not x.src or x.src[0].op is not Ops.INS or x.src[0].arg[0] is not X86Ops.RET else None),
   # function abi constraints
   (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi),
+  # a gate that picks between two places moves an address, not a value: lea each place arm so the select is a 64bit address
+  (UPat.var("m").where(UPat.var("a"), UPat.var("b")), where_places),
   # constants that can't be immediates, move them to registers
   (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),
-  (UPat.var("m").where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda m,a,b:
-   a.ins(X86Ops.VBLENDVPS, src=(b, a, m.replace(dtype=m.src[0].dtype)))),
-  (UPat.var("m").where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda m,a,b:
-   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
+  # a float comparison has no flag form, the blend consumes it as a mask. mask and value are the same width
+  (UPat(GroupOp.Comparison, src=(UPat(dtype=dtypes.float32), UPat()), name="m").where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda m,a,b:
+   a.ins(X86Ops.VBLENDVPS, src=(b, a, mask(m)))),
+  (UPat(GroupOp.Comparison, src=(UPat(dtype=dtypes.float64), UPat()), name="m").where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda m,a,b:
+   a.ins(X86Ops.VBLENDVPD, src=(b, a, mask(m)))),
+  # in this case we have a mask producing comparison whose user expects a bool, the mask is all ones or all zeros so the bool is mask & 1
   (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.cconst(1, dt))).f(Ops.NOOP, dtype=dtypes.bool)),
+   UOp(Ops.AND, src=(mask(x).bitcast(dt:=to_int(y.dtype)), UOp.cconst(1, dt))).bitcast(dtypes.bool)),
   # conditional moves that use flags
+  # TODO: remove this once we allow all flag producing ops in cmove
+  # the blends took every float gate a mask can serve, so a gate that is still not an integer comparison becomes one here
+  (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")), lambda m,a,b: g.where(a, b) if (g:=flag_gate(m)) is not None else None),
   (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)))),
   (UPat(Ops.CMPLT, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVB, src=(b, a, cmp(m)))),
@@ -397,16 +401,6 @@ isel_matcher = PatternMatcher([
   (UPat(Ops.CMPLT, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETL, src=(cmp(x),))),
   (UPat(Ops.CMPEQ, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETE, src=(cmp(x),))),
   (UPat(Ops.CMPNE, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETNE, src=(cmp(x),))),
-  # comparisons that produce masks (these aren't bool dtype)
-  (UPat(GroupOp.Comparison, src=(UPat(dtype=(dtypes.float32, dtypes.float64)), UPat()), name="x"), vcmp),
-  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int8s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQB)),
-  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int16s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQW)),
-  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int32s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQD)),
-  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int64s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQQ)),
-  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int8s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTB, src=(b, a))),
-  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int16s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTW, src=(b, a))),
-  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int32s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTD, src=(b, a))),
-  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int64s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTQ, src=(b, a))),
   # float unary
   (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPS)),
   (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPD)),
@@ -415,10 +409,9 @@ isel_matcher = PatternMatcher([
   (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x:
    x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPD, src=(y, imm(dtypes.uint8, 3)))),
   # for float16 we route the srcs through gprs, this is suboptimal for values in xmms, in that case we want vpunpcklwd
-  (UPat(Ops.STACK, dtypes.float16, name="x"), lambda x:
-   vpins(x.replace(src=tuple(s.bitcast(dtypes.int16) for s in x.src)))),
+  (UPat(Ops.STACK, dtypes.float16, name="x"), lambda x: vpins(x, tuple(s.bitcast(dtypes.int16) for s in x.src))),
   (UPat(Ops.STACK, dtypes.float32, name="x"), vinsertps),
-  (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), vpins),
+  (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), lambda x: vpins(x, x.src)),
   # INDEX on a vector register value extracts a single element
   (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),
@@ -500,11 +493,15 @@ isel_matcher = PatternMatcher([
   (UPat.var("y", dtypes.float64).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSD2SS, src=(y, y))),
   (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))),
   (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))),
-  (UPat(dtype=dtypes.uints+(dtypes.bool,)).cast(dtypes.ints, name="x"), lambda x:
+  # movzx only extends from an 8 or 16bit src, there is no 32 to 64bit form
+  (UPat(dtype=(dtypes.uint8, dtypes.uint16, dtypes.bool)).cast(dtypes.ints, name="x"), lambda x:
    x.ins(X86Ops.MOVZX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None),
   (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD) if x.max_numel() == 1 else None),
   (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x:
    x.ins(X86Ops.MOVSX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None),
+  # the remaining scalar int casts have no extending instruction, they are a MOV sized by the destination:
+  # narrowing writes the low bits, and widening is only from uint32 whose upper half a 32bit write already zeroed
+  (UPat(dtype=dtypes.ints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOV) if x.max_numel() == 1 else None),
   (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBW)),
   (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBD)),
   (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBQ)),
@@ -519,14 +516,13 @@ isel_matcher = PatternMatcher([
   (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXDQ)),
   # bitcasts between scalar floats and ints
   (UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, 0)))),
-  (UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), vpins),
+  (UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), lambda x: vpins(x, x.src)),
   (UPat(dtype=dtypes.int32s).bitcast(dtypes.float32).named("x"), lambda x: x.ins(X86Ops.VMOVD)),
   (UPat(dtype=dtypes.int64s).bitcast(dtypes.float64).named("x"), lambda x: x.ins(X86Ops.VMOVQ)),
   (UPat(dtype=dtypes.float32).bitcast(dtypes.int32s).named("x"), lambda x: x.ins(X86Ops.VMOVDm)),
   (UPat(dtype=dtypes.float64).bitcast(dtypes.int64s).named("x"), lambda x: x.ins(X86Ops.VMOVQm)),
   # index on a buffer (or the stack pointer) computes an address, addresses are 64bit values
-  (UPat((Ops.INDEX, Ops.SHRINK), name="x"),
-   lambda x: x.ins(X86Ops.LEA, dtype=dtypes.uint64, src=fold_address(x)) if not _is_vec_xmm(x.src[0]) else None),
+  (UPat((Ops.INDEX, Ops.SHRINK), name="x"), lambda x: lea(x) if not _is_vec_xmm(x.src[0]) else None),
   # TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q
   # copy, load, store
   # NOTE: copy here violates the spec, it only happens post register allocation when a reg to reg move needs to be inserted
@@ -555,7 +551,7 @@ isel_matcher = PatternMatcher([
 # so we rematerialize. This is different from rematerialization you might want to do in regalloc because it is not optional,
 # regalloc shouldn't rematerialize if a src of the instruction is dead, but here you need to as there's no fallback load from stack
 def flag_rematerialize(ctx:PreRegAllocContext, x:UOp):
-  flag_def = x if x.arg in X86GroupOp.WriteFlags or x.op in (Ops.RANGE, Ops.END) else x.src[-1] if x.arg in X86GroupOp.ReadFlags else None
+  flag_def = x if x.op in (Ops.RANGE, Ops.END) or x.arg[0] in X86GroupOp.WriteFlags else x.src[-1] if x.arg[0] in X86GroupOp.ReadFlags else None
   if flag_def is None: return None
   if ctx.lock is not None and ctx.lock is not flag_def: ctx.clobbered.add(ctx.lock)
   ctx.lock = flag_def
@@ -571,19 +567,19 @@ pre_regalloc_matcher = PatternMatcher([
 # TODO: control flow should be overhauled so that this isn't necessary
 def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
   loop_label = "_".join(str(i) for i in x.arg[:-1])
-  label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_{loop_label}")
+  label = UOp(Ops.INS, arg=(X86Ops.LABEL, dtypes.void), tag=f".LOOP_{loop_label}")
   # loop, cmp on backedge all we need is a jmp tag
   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.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}")
+    cmp = UOp(Ops.INS, arg=(X86Ops.CMPi if x.src[0].op is Ops.CAST else X86Ops.CMP, dtypes.void), src=(acc, x.src[0]))
+    jump_out = UOp(Ops.INS, arg=(X86Ops.JGE, dtypes.void), src=(cmp,), tag=f".LOOP_OUT_{loop_label}")
     ctx.loop_label[acc] = loop_label
     return (acc, [acc, label, cmp, jump_out])
 
 def lower_end(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
-  end_label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")
-  jmp = UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}")
+  end_label = UOp(Ops.INS, arg=(X86Ops.LABEL, dtypes.void), tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")
+  jmp = UOp(Ops.INS, arg=(X86Ops.JMP, dtypes.void), tag=f".LOOP_{ctx.loop_label[x.src[1]]}")
   inc = x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),))
   return (inc, [inc, jmp, end_label])
 
@@ -595,16 +591,16 @@ 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:=UOp.cconst(ctx.stack_size + x.tag, x.dtype), [nx])),
+  (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=UOp.cconst(ctx.stack_size + x.tag, x.dtype), [nx]) if x.arg[0] is X86Ops.FRAME_INDEX else None),
   # 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),
+  (UPat(Ops.INS, name="x"), lambda ctx,x: lower_loop(ctx, x) if x.arg[0] is X86Ops.LOOP_CMP else None),
   # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
   (UPat(Ops.RANGE, name="x"), lower_range),
   # rewrite END to ACC + 1 -> JUMP -> LABEL, also add the out of loop JUMP to the src so this becomes the jump target
   (UPat(Ops.END, name="x"), lower_end),
   # rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move
   (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]),
-   [ctx.ren.copy(x.src[0], greg(x)), nx] if greg(x) != greg(x.src[0]) else [nx]) if x.arg in X86GroupOp.TwoAddress else None),
+   [ctx.ren.copy(x.src[0], greg(x)), nx] if greg(x) != greg(x.src[0]) else [nx]) if x.arg[0] in X86GroupOp.TwoAddress else None),
 ])
 
 # ***** X86 instruction encoding *****
@@ -629,16 +625,15 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     r, _x, b = reg >> 3, idx >> 3, rm >> 3
     if sel: # VEX bytes
       vvvv = cast(Register, greg(vvvv_uop)).index if vvvv_uop is not None else 0
-      l = (max(reg_sz, rm_sz) > 16) & 0b1
-      if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp])
-      else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp])
+      if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | pp])
+      else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | pp])
     else: # optional PREFIX and REX bytes
       # PREFIX byte signaling 16 bit variant of instruction
       if sz == 2: inst += bytes([0x66])
       # bit signaling 64 bit variant of instruction
       w = sz == 8
       # legacy 8bit opcode is 1 less than 16-64bit variants
-      demote = (rm_sz == 1 or reg_sz == 1) and x.arg not in X86GroupOp.ReadFlags | {X86Ops.LEA}
+      demote = (rm_sz == 1 or reg_sz == 1) and x.arg[0] not in X86GroupOp.ReadFlags | {X86Ops.LEA}
       # REX byte is required when 64 bit or an extended reg is used (index 8 - 15) or lower 8 bits of (rsp, rbp, rsi, rdi) are accessed
       if w | r | _x | b | (reg_sz == 1 & reg >> 2) | (rm_sz == 1 & rm >> 2) | (demote and disp_uop is None and rm >= 4):
         inst += bytes([0b0100 << 4 | w << 3 | r << 2 | _x << 1 | b])
@@ -680,22 +675,22 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
   # get the encoding structure of the uop
   # when a uop writes to memory it takes the form of a store, dtype is void, no definition
   address:tuple[UOp|None, ...]
-  if x.arg in X86GroupOp.WriteMem:
+  if x.arg[0] 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.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 x.arg[0] 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.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:
+  if x.arg[0] in X86GroupOp.Rm2nd:
     if len(x.src) > 4: address, rest = x.src[1:5], x.src[:1] + x.src[5:]
     else: address, rest = (x.src[1], None, None, None), x.src[:1] + x.src[2:]
-    # cmp/vucomiss reg, rm don't define a new register
+    # cmp reg, rm doesn't define a new register
     return _encode(x, *address, *rest) if x.dtype is not dtypes.void else _encode(rest[0], *address)
 
   return None
@@ -760,18 +755,12 @@ encodings = {
   # packed int binary
   X86Ops.VPSLLVD: lambda x: encode(x, 0x47, pp=1, sel=2), X86Ops.VPSLLVQ: lambda x: encode(x, 0x47, pp=1, sel=2, we=1),
   X86Ops.VPSRLVD: lambda x: encode(x, 0x45, pp=1, sel=2), X86Ops.VPSRLVQ: lambda x: encode(x, 0x45, pp=1, sel=2, we=1),
-  X86Ops.VPCMPGTB: lambda x: encode(x, 0x64, pp=1, sel=1), X86Ops.VPCMPGTW: lambda x: encode(x, 0x65, pp=1, sel=1),
-  X86Ops.VPCMPGTD: lambda x: encode(x, 0x66, pp=1, sel=1), X86Ops.VPCMPGTQ: lambda x: encode(x, 0x37, pp=1, sel=2),
-  X86Ops.VPCMPEQB: lambda x: encode(x, 0x74, pp=1, sel=1), X86Ops.VPCMPEQW: lambda x: encode(x, 0x75, pp=1, sel=1),
-  X86Ops.VPCMPEQD: lambda x: encode(x, 0x76, pp=1, sel=1), X86Ops.VPCMPEQQ: lambda x: encode(x, 0x29, pp=1, sel=2),
   X86Ops.VPMULLW: lambda x: encode(x, 0xD5, pp=1, sel=1), X86Ops.VPMULLD: lambda x: encode(x, 0x40, pp=1, sel=2),
   X86Ops.VPADDB: lambda x: encode(x, 0xFC, pp=1, sel=1), X86Ops.VPADDW: lambda x: encode(x, 0xFD, pp=1, sel=1),
   X86Ops.VPADDD: lambda x: encode(x, 0xFE, pp=1, sel=1), X86Ops.VPADDQ: lambda x: encode(x, 0xD4, pp=1, sel=1),
   X86Ops.VPSUBB: lambda x: encode(x, 0xF8, pp=1, sel=1), X86Ops.VPSUBW: lambda x: encode(x, 0xF9, pp=1, sel=1),
   X86Ops.VPSUBD: lambda x: encode(x, 0xFA, pp=1, sel=1), X86Ops.VPSUBQ: lambda x: encode(x, 0xFB, pp=1, sel=1),
   X86Ops.VPSRAVD: lambda x: encode(x, 0x46, pp=1, sel=2),
-  # float cmp
-  X86Ops.VUCOMISS: lambda x: encode(x, 0x2E, pp=0, sel=1), X86Ops.VUCOMISD: lambda x: encode(x, 0x2E, pp=1, sel=1),
   # scalar / packed float binary
   X86Ops.VADDSS: lambda x: encode(x, 0x58, pp=2, sel=1), X86Ops.VADDPS: lambda x: encode(x, 0x58, pp=0, sel=1),
   X86Ops.VADDSD: lambda x: encode(x, 0x58, pp=3, sel=1), X86Ops.VADDPD: lambda x: encode(x, 0x58, pp=1, sel=1),
@@ -787,7 +776,6 @@ encodings = {
   X86Ops.CMOVB: lambda x: encode(x, 0x0F42), X86Ops.CMOVL: lambda x: encode(x, 0x0F4C),
   X86Ops.CMOVE: lambda x: encode(x, 0x0F44), X86Ops.CMOVNE: lambda x: encode(x, 0x0F45),
   X86Ops.VBLENDVPS: lambda x: encode(x, 0x4A, pp=1, sel=3), X86Ops.VBLENDVPD: lambda x: encode(x, 0x4B, pp=1, sel=3),
-  X86Ops.VPBLENDVB: lambda x: encode(x, 0x4C, pp=1, sel=3),
   # shuffles
   X86Ops.VPSRLDQ: lambda x: encode(x, 0x73, reg=3, pp=1, sel=1),
   X86Ops.VPINSRB: lambda x: encode(x, 0x20, pp=1, sel=3), X86Ops.VPINSRW: lambda x: encode(x, 0xC4, pp=1, sel=1),
@@ -823,28 +811,28 @@ class X86Renderer(ISARenderer):
     super().__init__(target)
     from tinygrad.runtime.support.compiler_cpu import X86Compiler
     self.compiler = X86Compiler()
-  def is_two_address(self, x:UOp) -> bool: return x.arg in X86GroupOp.TwoAddress
+  def is_two_address(self, x:UOp) -> bool: return x.op is Ops.INS and x.arg[0] in X86GroupOp.TwoAddress
   def stack_pointer(self) -> UOp: return def_reg(dtypes.uint64, RSP)
   # the value of a BUFFER is its address, it moves through registers and the stack as a 64bit int
   def copy(self, x:UOp, reg:Register):
-    dt = dtypes.uint64 if x.op is Ops.BUFFER else x.dtype
-    ret = isel_matcher.rewrite(UOp(Ops.COPY, dt, (x,), tag=reg))
+    if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
+    ret = isel_matcher.rewrite(UOp(Ops.COPY, (x,), tag=reg))
     assert ret is not None, f"failed to copy {x}"
     return ret
 
   def spill(self, disp:UOp, x:UOp) -> UOp:
-    if x.op is Ops.BUFFER: x = x.replace(dtype=dtypes.uint64)
+    if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
     is_xmm = isinstance(x.tag, tuple) and x.tag[0].cons[0].size == 16
     op = X86Ops.VMOVUPSm if is_xmm else X86Ops.MOVm
-    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=op, tag=x.tag)
+    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
 
   def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
     is_xmm = reg.cons[0].size == 16
     dt = dtypes.uint64 if x.op is Ops.BUFFER else x.dtype
-    return UOp(Ops.INS, dt, fold_address(self.stack_pointer().index(disp)), X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, (reg,))
+    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)), arg=(X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, dt), tag=(reg,))
 
   def asm_str(self, uops:list[UOp], function_name:str) -> str:
-    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_op(x:UOp) -> str: return f"    {(o[7:-1] if (o:=str(x.arg[0]))[-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.src[0].val) if s.op is Ops.CAST else reg_strs[o].get(s.dtype.itemsize, o) if \
@@ -852,17 +840,17 @@ class X86Renderer(ISARenderer):
       def _mem_adress(base:UOp, idx:UOp, disp:UOp, sz:UOp) -> list[str]:
         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:])
-      elif len(x.src) > 4 and x.arg in X86GroupOp.Rm2nd: ret = _format((x, x.src[0])) + _mem_adress(*x.src[1:5]) + _format(x.src[5:])
+      if len(x.src) > 4 and x.arg[0] in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:])
+      elif len(x.src) > 3 and x.arg[0] in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:])
+      elif len(x.src) > 4 and x.arg[0] in X86GroupOp.Rm2nd: ret = _format((x, x.src[0])) + _mem_adress(*x.src[1:5]) + _format(x.src[5:])
       else: ret = _format((x,) + x.src)
       return ", ".join(ret)
 
     asm = [f".{function_name}:"]
     for u in uops:
-      if u.op is not Ops.INS or u.arg is X86Ops.DEFINE: continue
-      if u.arg is X86Ops.LABEL: asm.append(f"{str(u.tag)}:")
-      elif u.arg is X86Ops.RET: asm.append(_format_op(u))
+      if u.op is not Ops.INS or u.arg[0] is X86Ops.DEFINE: continue
+      if u.arg[0] is X86Ops.LABEL: asm.append(f"{str(u.tag)}:")
+      elif u.arg[0] is X86Ops.RET: asm.append(_format_op(u))
       else: asm.append(_format_op(u) + " " + _format_operands(u))
     return "\n".join(asm)
 
@@ -871,15 +859,15 @@ class X86Renderer(ISARenderer):
     jumps: dict[UOp, int] = {}
     binary = bytearray()
     for u in uops:
-      if u.op is not Ops.INS or u.arg is X86Ops.DEFINE: continue
-      if u.arg is X86Ops.LOOP_CMP: continue
-      if u.arg is X86Ops.LABEL:
+      if u.op is not Ops.INS or u.arg[0] is X86Ops.DEFINE: continue
+      if u.arg[0] is X86Ops.LOOP_CMP: continue
+      if u.arg[0] is X86Ops.LABEL:
         targets[u.tag] = len(binary)
         continue
-      if u.arg not in encodings or (l:=encodings[u.arg](u)) is None:
-        raise RuntimeError(f"failed to encode {u.arg} with {u.dtype} srcs {[x.dtype for x in u.src]}")
+      if u.arg[0] not in encodings or (l:=encodings[u.arg[0]](u)) is None:
+        raise RuntimeError(f"failed to encode {u.arg[0]} with {u.dtype} srcs {[x.dtype for x in u.src]}")
       binary.extend(l)
-      if u.arg in (X86Ops.JL, X86Ops.JB, X86Ops.JE, X86Ops.JNE, X86Ops.JGE, X86Ops.JMP): jumps[u] = len(binary)
+      if u.arg[0] in (X86Ops.JL, X86Ops.JB, X86Ops.JE, X86Ops.JNE, X86Ops.JGE, X86Ops.JMP): jumps[u] = len(binary)
     # fixup jump targets now that encoding size is known
     for u in uops:
       if (t:=jumps.get(u)) is not None: binary[t-4:t] = (targets[u.tag] - t).to_bytes(4, 'little', signed=True)

tinygrad/renderer/nir.py

@@ -125,7 +125,7 @@ class NIRRenderer(Renderer):
     (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
     # a bool is one bit in NIR but a byte in memory, so every access to a bool buffer goes through a uint8 view of it
     (UPat(Ops.LOAD, dtypes.bool, name="x"),
-     lambda x: x.replace(dtype=None, src=(with_storage(x.src[0], dtypes.uint8),)+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())
+     lambda x: x.replace(src=(with_storage(x.src[0], dtypes.uint8),)+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())
       +x.src[2:]).cast(dtypes.bool)),
     (UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
      lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), x.src[1].cast(dtypes.uint8))+x.src[2:])),

tinygrad/renderer/ptx.py

@@ -47,7 +47,7 @@ ptx_matcher = PatternMatcher([
     lambda x: (UOp(x.op, src=tuple(vv.cast(dtypes.float32) for vv in x.src), arg=x.arg).cast(dtypes.half))),
   # a bool is a predicate register in PTX but a byte in memory, so a bool buffer is accessed through a uint8 view of it
   (UPat(Ops.LOAD, dtypes.bool, src=(UPat(name="idx"),), name="x", allow_any_len=True),
-   lambda x,idx: x.replace(dtype=None, src=(with_storage(idx, dtypes.uint8),) + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ())
+   lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8),) + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ())
      + x.src[2:]).cast(dtypes.bool) if idx.addrspace != AddrSpace.REG else None),
   (UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
    lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), x.src[1].cast(dtypes.uint8))+x.src[2:]) if idx.addrspace != AddrSpace.REG else None),

tinygrad/runtime/ops_amd.py

@@ -8,7 +8,7 @@ from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filte
 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 VIZ, ceildiv, unwrap, pluralize
+from tinygrad.helpers import VIZ, HCQ2, ceildiv, unwrap, pluralize
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
@@ -1153,4 +1153,4 @@ class AMDDevice(HCQCompiled):
 
   def hw_copy_queues(self): return [(f"SDMA:{i}", functools.partial(unwrap(self.hw_copy_queue_t), queue_idx=i)) for i in self.sdma_queues]
 
-if getenv("HCQ2"): from extra.hcq2.ops_amd2 import * # noqa: F401, F403 # pylint: disable=unused-import
+if HCQ2: from extra.hcq2.ops_amd2 import * # noqa: F401, F403 # pylint: disable=unused-import

tinygrad/runtime/ops_cpu.py

@@ -26,24 +26,24 @@ MAX_ARGS, CMD_SIZE, RING_SLOTS, FUNCS = 63, 64, (16 << 10), (() if WIN else ('cl
 
 def signal_prog():
   val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
-  return UOp.param(0, dtypes.uint32, (1,))[0].store(val.cast(dtypes.uint32))
+  return UOp.param(0, dtypes.uint32, 1)[0].store(val.cast(dtypes.uint32))
 
 def wait_prog():
   val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
-  return (v:=UOp.param(0, dtypes.uint32, (1,), volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32))
+  return (v:=UOp.param(0, dtypes.uint32, 1, volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32))
 
 def timestamp_prog():
   if WIN: val = UOp.const(0, dtypes.uint64)
   else:
-    fn, ts = UOp.param(1, dtypes.uint64, (1,)), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG)
+    fn, ts = UOp.param(1, dtypes.uint64, 1), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG)
     call = fn[0].load().call(UOp.const(6 if OSX else 1, dtypes.int), ts[0], ret_dtype=dtypes.void) # clock_gettime(CLOCK_MONOTONIC, &ts)
     val = ts.after(call)[0].load() * 1_000_000_000 + ts.after(call)[1].load()
-  return UOp.param(0, dtypes.uint64, (1,))[0].store(val)
+  return UOp.param(0, dtypes.uint64, 1)[0].store(val)
 
 def worker_prog():
-  ring = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True)
-  wait, done = UOp.param(1, dtypes.uint64, (1,), volatile=True), UOp.param(2, dtypes.uint64, (1,), volatile=True)
-  sem, cur = UOp.param(3, dtypes.uint64, (1,)), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last
+  ring = UOp.param(0, dtypes.uint64, RING_SLOTS * CMD_SIZE, volatile=True)
+  wait, done = UOp.param(1, dtypes.uint64, 1, volatile=True), UOp.param(2, dtypes.uint64, 1, volatile=True)
+  sem, cur = UOp.param(3, dtypes.uint64, 1), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last
 
   # spin on windows, sem_wait to sleep on posix
   if WIN: ready = (v:=wait.after(lw:=UOp.loop(1), cur)[0].load()).end(lw, v <= cur)

tinygrad/runtime/ops_qcom.py

@@ -20,7 +20,7 @@ BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2
 def dcache_flush():
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
   from tinygrad.codegen import to_program
-  buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU)
+  buf, n = UOp.param(0, dtypes.uint8, 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");', dtypes.void))
   sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg=('__asm__ volatile("dsb sy" ::: "memory");', dtypes.void)),

tinygrad/runtime/support/hcq.py

@@ -1,24 +1,16 @@
 from __future__ import annotations
 from typing import cast, Callable, Type, TypeVar, Generic, Any
-import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, itertools
+import contextlib, decimal, statistics, time, ctypes, array, os, collections, itertools
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
-from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, unwrap
+from tinygrad.helpers import DEV, PROFILE, getenv, from_mv, cpu_profile, ProfileRangeEvent, unwrap
 from tinygrad.helpers import suppress_finalizing, pluralize, TracingKey
 from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
 from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
-from tinygrad.runtime.support.memory import BumpAllocator
+from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
 from tinygrad.renderer import Renderer
 
-class MMIOInterface:
-  def __init__(self, addr:int, nbytes:int, fmt='B'): self.mv, self.addr, self.nbytes, self.fmt = to_mv(addr, nbytes).cast(fmt), addr, nbytes, fmt
-  def __len__(self): return self.nbytes // struct.calcsize(self.fmt)
-  def __getitem__(self, k): return (self.mv[k] if self.fmt == 'B' else self.mv[k].tolist()) if isinstance(k, slice) else self.mv[k]
-  def __setitem__(self, k, v): self.mv[k] = v
-  def view(self, offset:int=0, size:int|None=None, fmt=None) -> MMIOInterface:
-    return MMIOInterface(self.addr+offset, (self.nbytes - offset) if size is None else size, fmt=fmt or self.fmt)
-
 class FileIOInterface:
   """
   Hardware Abstraction Layer for HCQ devices. The class provides a unified interface for interacting with hardware devices.

tinygrad/runtime/support/hcq2.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import cast, TypeVar, Generic, Any, Sequence, Iterable
+from typing import cast, TypeVar, Generic, Any, Sequence, Iterable, TYPE_CHECKING
 import struct, functools, time, collections, itertools, decimal, statistics
 from dataclasses import replace, dataclass, field
 from tinygrad.helpers import suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap, PROFILE
@@ -9,11 +9,11 @@ from tinygrad.device import ProfileDeviceEvent, ProfileGraphEntry, ProfileGraphE
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, rewrite_group, GroupOp
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.dtype import dtypes, truncate, DType
-from tinygrad.runtime.support.hcq import MMIOInterface, HCQBuffer
-from tinygrad.runtime.support.memory import BumpAllocator
+from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop
-from tinygrad.engine.realize import pm_flatten_linear, lower_and_compile
+from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear,lower_and_compile
+
+if TYPE_CHECKING: from tinygrad.runtime.support.hcq import HCQBuffer # TODO: remove that
 
 # *****************
 # 0. helpers
@@ -139,7 +139,7 @@ def _get_enqueue_devs(call:UOp) -> Any|None:
 
 def copy_with_kernel(call:UOp, dst:UOp, src:UOp) -> UOp|None:
   if (devs:=_get_enqueue_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None
-  d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2))
+  d, s = (UOp.param(i, dst.dtype, n:=dst.max_numel(), device=devs) for i in range(2))
   ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1)
   return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src))
 
@@ -168,8 +168,8 @@ class BatchCtx:
   slots:dict[str, int] = field(default_factory=lambda: collections.defaultdict(lambda: next(UOp.unique_num)))
 
 def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]:
-  return [[b if (b:=_lane(a, lane)).op is Ops.PARAM or (b.op is Ops.MSELECT and b.src[0].op is Ops.PARAM) else b.buffer
-           for a in get_call_arg_uops(call)] for lane in range(len(devices))]
+  def dep_buf(b:UOp) -> Any: return base if (base:=(b.src[0] if b.op is Ops.MSELECT else b).base).op is Ops.PARAM else b.buffer
+  return [[dep_buf(_lane(a, lane)) for a in get_call_arg_uops(call)] for lane in range(len(devices))]
 
 def _wait_ins(ctx:BatchCtx, bufs_by_lane:list[list[Any]], write, devices:tuple[str, ...], queue:str, tag:int) -> list[UOp]:
   deps:list[Dep] = []
@@ -462,7 +462,7 @@ def hcq_lower(linear:UOp, pm_encode:PatternMatcher) -> UOp:
   linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
 
   # and compile it
-  return lower_and_compile(graph_rewrite(linear, pm_replace_params, walk=True, name="replace params"))
+  with Context(EMULATED_DTYPES=""): return lower_and_compile(graph_rewrite(linear, pm_replace_params, walk=True, name="replace params"))
 
 @rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
 def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:

tinygrad/runtime/support/memory.py

@@ -1,6 +1,15 @@
-import collections, functools, dataclasses, enum
+from __future__ import annotations
+import collections, functools, dataclasses, enum, struct
 from typing import Any, ClassVar
-from tinygrad.helpers import round_up, getenv
+from tinygrad.helpers import round_up, getenv, to_mv
+
+class MMIOInterface:
+  def __init__(self, addr:int, nbytes:int, fmt='B'): self.mv, self.addr, self.nbytes, self.fmt = to_mv(addr, nbytes).cast(fmt), addr, nbytes, fmt
+  def __len__(self): return self.nbytes // struct.calcsize(self.fmt)
+  def __getitem__(self, k): return (self.mv[k] if self.fmt == 'B' else self.mv[k].tolist()) if isinstance(k, slice) else self.mv[k]
+  def __setitem__(self, k, v): self.mv[k] = v
+  def view(self, offset:int=0, size:int|None=None, fmt=None) -> MMIOInterface:
+    return MMIOInterface(self.addr+offset, (self.nbytes - offset) if size is None else size, fmt=fmt or self.fmt)
 
 class BumpAllocator:
   def __init__(self, size:int, base:int=0, wrap:bool=True): self.size, self.ptr, self.base, self.wrap = size, 0, base, wrap

tinygrad/schedule/__init__.py

@@ -94,7 +94,7 @@ pm_post_sched_cache = PatternMatcher([
   # only resolve buffer PARAMs (slot>=0); ALU/shape vars use slot=-1 and must not be swapped for call args
   (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg.slot] if x.arg.slot >= 0 else None),
   # create new BUFFERs
-  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
+  (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b:
    create_new_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
 ])
 

tinygrad/schedule/multi.py

@@ -266,7 +266,7 @@ def store_dest_multi(root:UOp, multi:UOp):
 
 def passthrough_multi(root:UOp, multi:UOp):
   new_src = (multi.src[0],)+tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src[1:])
-  return UOp(root.op, root.dtype, src=new_src, arg=root.arg).unshard(multi.arg, multi.src[1:])
+  return UOp(root.op, src=new_src, arg=root.arg).unshard(multi.arg, multi.src[1:])
 
 def rewrite_into_function(call:UOp):
   if call.arg.precompile: return None
@@ -280,13 +280,8 @@ def rewrite_into_function(call:UOp):
                            for i, s in enumerate(new_body.src)])
   return call.replace(src=(new_body,)+new_args)
 
-def param_to_multi(p:UOp):
-  if p.axis is None: return None
-  return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.multiple_of, p.arg.name, p.arg.addrspace).unshard(p.axis)
-
 # NOTE: this is the same pattern as unrolled ranges
 multi_pm = PatternMatcher([
-  (UPat(Ops.PARAM, name="p"), param_to_multi),
   (UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.UNSHARD])), alu_multi),
   (UPat(Ops.REDUCE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), reduce_multi),
   (UPat(Ops.RESHAPE, src=(UPat(Ops.UNSHARD, name="multi"), UPat()), name="root"), reshape_multi),

tinygrad/schedule/prepare.py

@@ -1,7 +1,7 @@
 import itertools
 from tinygrad.dtype import dtypes, to_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp
-from tinygrad.uop.ops import graph_rewrite, rewrite_group, shape_to_shape_arg, ParamArg, identity_element
+from tinygrad.uop.ops import graph_rewrite, rewrite_group, ParamArg, identity_element
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.helpers import prod, getenv, all_int, DEBUG, SPLIT_REDUCEOP, OPENPILOT_HACKS, FLOAT16, argsort
 from tinygrad.schedule.indexing import apply_movement_op
@@ -100,10 +100,19 @@ def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None:
     if [x.arg.slot for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg.slot for x in params]}")
     if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")
 
+  # params have a flat storage size in the arg, the logical shape is a view (RESHAPE/SHRINK/UNSHARD) on top of it.
+  # substitute args by their flat max-shaped storage view so the movement views on the params stay valid
+  def flat_storage(a:UOp) -> tuple[int, UOp]:  # returns (size, view of a as flat max-shaped storage)
+    shp = a.max_shard_shape if a.axis is not None and isinstance(a.device, tuple) else a.max_shape
+    return (n:=prod(shp)), a if a.shape == (n,) else a.pad_to(shp).reshape((n,))
   dict_map = {x:args[x.arg.slot] for x in params}
   for i, (p, a) in enumerate(dict_map.items()):
-    if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
-    if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
+    if p.arg.size is not None:
+      n, flat = flat_storage(a)
+      if p.arg.size != n: raise TypeError(f"arg {i} shape mismatch: expected size {p.arg.size}, got {a.shape}")
+      dict_map[p] = flat
+    elif a.shape != ():
+      raise TypeError(f"arg {i} shape mismatch: expected scalar, got {a.shape}")
     if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
   return c.src[0].substitute(dict_map, walk=True)
 
@@ -192,9 +201,9 @@ def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
     # if there's already a buffer, we just use it
     return existing_buf.flatten().store(input_src)
   # create the output buffer
-  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
+  buf = UOp(Ops.BUFFER, arg=ParamArg(next(ctx), copy.dtype, size=prod(input_src.max_shape), device=copy.device))
   # reshape back to input
-  return buf.after(buf.store(input_src)).reshape(copy.shape)
+  return buf.reshape(input_src.max_shape).after(buf.store(input_src)).reshape(copy.shape)
 
 pm_copy_to_store = PatternMatcher([
   (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),

tinygrad/schedule/rangeify.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field, replace
 from typing import cast
 import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import prod, dedup, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
@@ -162,6 +162,12 @@ pm_no_indexing_calls = PatternMatcher([
   (UPat(Ops.CALL, name="u"), no_indexing_calls),
 ])
 
+# the kernel graph is what gets executed: no shape views left in it, the storage of a value is just the storage
+pm_no_views = PatternMatcher([
+  (UPat((Ops.RESHAPE, Ops.SHRINK), name="v", src=(UPat((Ops.AFTER, Ops.PARAM, Ops.UNSHARD, Ops.MSTACK, Ops.BUFFER)),), allow_any_len=True), lambda v:
+   v.src[0]),
+])
+
 DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8, "CPU": 31} # TODO: get from device?
 @dataclass
 class LimitBufsContext:
@@ -223,7 +229,7 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
 
   # NOTE: the local BUFFER needs to be disambiguated here
   if x.arg.addrspace == AddrSpace.GLOBAL:
-    buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(next(ctx), x.dtype, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
+    buf = UOp(Ops.BUFFER, arg=ParamArg(next(ctx), x.dtype, size=size, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
     do_store = buf.index(idx).store(x.src[0]).end(*rngs)
     return buf.after(do_store)
 
@@ -282,8 +288,7 @@ class LocalAddBufferContext:
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
   # Variables (ALU buffers with a value range) are scalar symbolic values, not real buffers: they become ALU params with no slot
   if buf.is_variable: return buf.replace(op=Ops.PARAM)
-  param = UOp(Ops.PARAM, src=(UOp.const(prod(buf.max_shape)),),
-              arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device))
+  param = UOp(Ops.PARAM, arg=ParamArg(ctx.dg, buf.dtype, prod(buf.max_shape), addrspace=buf.addrspace, device=buf.device))
   ret = param.reshape(buf.max_shape)
   # if the buffer has symbolic shape, shrink the max-sized view to the actual shape
   if buf.max_shape != buf.shape: ret = ret.shrink(tuple((0, s) for s in buf.shape))
@@ -383,6 +388,7 @@ def get_kernel_graph(tsink:UOp) -> UOp:
   tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_param_range_tags, ctx=itertools.count(paramarg_start), bottom_up=True, name="stage to store")
   tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
   tsink = graph_rewrite(tsink, pm_no_indexing_calls, name="remove indexing from call args")
+  tsink = graph_rewrite(tsink, pm_no_views, name="remove views from the kernel graph")
 
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
   if SPEC:

tinygrad/tensor.py

@@ -198,9 +198,7 @@ def finalize_after(ctx:AllocCtx, x:UOp):
 
 def replace_input_buffer(ctx:AllocCtx, b:UOp):
   ctx.replacements.append(b)
-  if b.is_bound_var or b.is_variable: return b.param_like(len(ctx.replacements)-1)
-  return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
-                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
+  return b.param_like(len(ctx.replacements)-1)
 
 def replace_input_view(ctx:AllocCtx, b:UOp): return replace_input_buffer(ctx, b) if b in ctx.views else None
 
@@ -211,7 +209,7 @@ pm_finalize_call = PatternMatcher([
 
 pm_replace_buf = PatternMatcher([
   # replace BUFFER with PARAM for cache key normalization
-  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
+  (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b:
    replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
   # replace SHRINK with PARAM
   (UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), name="b", allow_any_len=True), replace_input_view),
@@ -381,7 +379,7 @@ class Tensor(RandMixin):
   # ***** data handlers ****
 
   def as_param(self, slot:int):
-    return Tensor(UOp.param(slot, self.dtype, self.uop.shard_shape, self.device, axis=self.uop.axis))
+    return Tensor(self.uop.param_like(slot))
 
   def call(self, *lst:Tensor, fxn:Tensor|UOp, grad_fxn:Callable|None=None) -> Tensor:
     fret = fxn._uop.call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn)

tinygrad/uop/ops.py

@@ -23,17 +23,21 @@ class AxisType(Enum):
 class ParamArg:
   slot: int
   dtype: DType
+  # number of elements in the buffer. always a concrete int (never symbolic), None for scalars (shape ())
+  size: int|None = None
   vmin_vmax: tuple[PyConst, PyConst]|None = None
   multiple_of: int|None = None
   name: str|None = None
   addrspace: AddrSpace|None = AddrSpace.GLOBAL
-  axis: int|None = None
   device: str|tuple[str, ...]|None = None
   volatile: bool = False
+  # (h, w) if this is an image2d buffer, then size == h*w*4
+  image: tuple[int, int]|None = None
   def __repr__(self):
-    fields = (("vmin_vmax", None), ("multiple_of", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None),
-              ("volatile", False))
-    args = [repr(self.slot), repr(self.dtype)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
+    fields = (("vmin_vmax", None), ("multiple_of", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("device", None),
+              ("volatile", False), ("image", None))
+    args = [repr(self.slot), repr(self.dtype)] + ([repr(self.size)] if self.size is not None else []) + \
+      [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
     return f"ParamArg({', '.join(args)})"
 axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.WEAK: "L",
                 AxisType.LOOP: "L", AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
@@ -111,25 +115,22 @@ def promo_dtype(src:tuple[UOp,...]) -> DType:
   dts = [x.dtype for x in src]
   return dts[0] if all_same(dts) else least_upper_dtype(*dts)
 
-def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
-  # here are the dtype production rules, eventually this will go in UOp as a recursive property
+def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType:
+  # here are the dtype production rules, this is the only place a UOp dtype comes from
   match op:
     case Ops.STORE | Ops.LINEAR | Ops.SINK | Ops.PROGRAM | Ops.SOURCE | \
-         Ops.END | Ops.BARRIER | Ops.GROUP | Ops.IF | Ops.ENDIF | \
+         Ops.END | Ops.BARRIER | Ops.GROUP | Ops.IF | Ops.ENDIF | Ops.NOOP | \
          Ops.TUPLE | Ops.FUNCTION | Ops.CUSTOM_FUNCTION | Ops.REWRITE_ERROR | Ops.PYLITERAL:
       # always void
       return dtypes.void
     case Ops.CALL:
-      # a CALL of an opaque body is void, a CALL of an address can return a value
-      return dtypes.void if src[0].dtype is dtypes.void else None
+      # a CALL of an opaque body (CallInfo arg) is void, a CALL of an address states its return dtype in the arg
+      return arg if isinstance(arg, DType) else dtypes.void
     case Ops.CUSTOM | Ops.CUSTOMI:
       assert isinstance(arg, tuple) and len(arg) == 2 and isinstance(arg[1], DType), f"CUSTOM/CUSTOMI arg must be (str, DType), got {arg}"
       return arg[1]
     case Ops.INS:
-      return None
-    case Ops.NOOP:
-      # NOOP can be void or carry any dtype (e.g. x.f(Ops.NOOP) or substitute base with NOOP)
-      return None
+      return arg[1] if isinstance(arg, tuple) else dtypes.void
     case Ops.INDEX:
       # an image access is always float, no matter the storage dtype
       # TODO: should there be a CAST so src[0].dtype just work?
@@ -190,14 +191,11 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
 
 class UOpMetaClass(type):
   ucache:dict[tuple, weakref.ReferenceType[UOp]] = {}
-  def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
+  def __call__(cls, op:Ops, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
                metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None):
-    if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void
-    # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py
-    if SPEC == 2 and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
-      raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}")
-    if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret
-    UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key))
+    # NOTE: type(arg) is in the key because True == 1 and False == 0 as dict keys, and a CONST's dtype is the type of its arg
+    if (wret:=UOpMetaClass.ucache.get(key:=(op, src, arg, tag, type(arg)), None)) is not None and (ret:=wret()) is not None: return ret
+    UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key[:4]))
     if metadata is not None: all_metadata[created] = metadata
     # NOTE: this value is set by pickle when pickling a realized tensor
     if _buffer is not None:
@@ -238,24 +236,24 @@ from tinygrad.mixin.rand import RandMixin
 @dataclass(eq=False, slots=True)
 class UOp(RandMixin, metaclass=UOpMetaClass):
   op:Ops
-  dtype:DType = dtypes.void
   src:tuple[UOp, ...] = tuple()
   arg:Any = None
   tag:Any = None
+  @functools.cached_property
+  def dtype(self) -> DType: return dtype_from_uop(self.op, self.src, self.arg)
   def __del__(self):
     # NOTE: getattr because this object may be partially constructed (e.g. if __init__ raised, like the BEAM timeout SIGALRM)
     if Ops is not None and getattr(self, 'op', None) is Ops.BUFFER and (buffer:=buffers.get(self)) is not None: buffer.ref(-1)
-    try: del UOpMetaClass.ucache[(self.op, self.dtype, self.src, self.arg, self.tag)]
+    try: del UOpMetaClass.ucache[(self.op, self.src, self.arg, self.tag, type(self.arg))]
     except (AttributeError, KeyError): pass
   def __reduce__(self):
-    args = [self.op, self.dtype, self.src, self.arg, self.tag, self.metadata]
+    args = [self.op, self.src, self.arg, self.tag, self.metadata]
     if self.op is Ops.BUFFER and self.realized is not None: args.append(self.realized)
     return UOp, tuple(args)
   def replace(self, **kwargs) -> UOp:
-    new_args = (kwargs.pop("op", self.op), kwargs.pop("dtype", self.dtype), kwargs.pop("src", self.src),
-                kwargs.pop("arg", self.arg), kwargs.pop("tag", self.tag))
+    new_args = (kwargs.pop("op", self.op), kwargs.pop("src", self.src), kwargs.pop("arg", self.arg), kwargs.pop("tag", self.tag))
     assert len(kwargs) == 0, f"unused kwargs in replace {list(kwargs)}"
-    if (self.op, self.dtype, self.src, self.arg, self.tag) == new_args: return self
+    if (self.op, self.src, self.arg, self.tag) == new_args: return self
     return UOp(*new_args)
   def rtag(self, tag=True): return self.replace(tag=tag)
   @property
@@ -277,7 +275,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return repr(self.arg)
   def tagstr(self): return f", tag={self.tag}" if self.tag is not None else ""
 
-  def f(self, op, **kwargs): return UOp(op, dtype=kwargs.pop("dtype", self.dtype), src=(self,), **kwargs)
+  def f(self, op, **kwargs): return UOp(op, src=(self,), **kwargs)
 
   @functools.cached_property
   def backward_slice(self:UOp) -> dict[UOp, None]:
@@ -381,9 +379,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       case Ops.GETADDR: return ()
       case Ops.RANGE | Ops.SPECIAL: return ()
       case Ops.BINARY: return (len(self.arg),)
-      case Ops.BUFFER:
-        if len(self.src): return self.src[0].as_shape
-        return ()
+      case Ops.BUFFER | Ops.PARAM:
+        # PARAM/BUFFER don't have a shape input, they have a size in the arg: int gives shape (size,), None gives ()
+        if (img:=self.arg.image) is not None: return (img[0], img[1], 4)
+        return () if self.arg.size is None else (self.arg.size,)
       case Ops.CUSTOM | Ops.CUSTOMI:
         if self.dtype is dtypes.void: return None
         input_shapes = [x._shape for x in self.src if x._shape is not None]
@@ -394,10 +393,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
         # STAGE adds the existing shape to the front, opposite of INDEX
         return tuple([int(r.vmax+1) for r in self.src[1:]])+self.src[0].shape
 
-      # param has shape as the only arg
-      case Ops.PARAM:
-        return self.src[0].as_shape
-
       # wmma output shape = accumulator shape (src[2])
       case Ops.WMMA:
         wmma_b = _broadcast_shape(self.src[0].shape[:-1], self.src[1].shape[:-1], self.src[2].shape[:-1])
@@ -547,7 +542,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   @recursive_property
   def trace_num(self):
     num = next(ucount)
-    uop_fields[num] = (self.op, self.dtype, tuple(s.trace_num for s in self.src), self.arg, self.tag)+((self.metadata,) if TRACEMETA>=2 else ())
+    uop_fields[num] = (self.op, tuple(s.trace_num for s in self.src), self.arg, self.tag)+((self.metadata,) if TRACEMETA>=2 else ())
     return num
 
   # *** uop syntactic sugar ***
@@ -606,7 +601,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   @property
   def without_after(self) -> UOp: return self.src[0] if self.op is Ops.AFTER else self
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
-  def ins(self, arg, **kwargs): return UOp(Ops.INS, kwargs.pop("dtype", self.dtype), kwargs.pop("src", self.src), arg, kwargs.pop("tag", self.tag))
+  def ins(self, arg, **kwargs): return UOp(Ops.INS, kwargs.pop("src", self.src), (arg, kwargs.pop("dtype", self.dtype)), kwargs.pop("tag", self.tag))
   def contract(self, *rngs:UOp):
     assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast"
     return UOp.stack(*[self.substitute(dict(zip(rngs, [r.const_like(i) for r,i in zip(rngs, idx)])))
@@ -708,7 +703,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if self.op is Ops.GETTUPLE:
       in_tuple = self.src[0].src[0] if self.src[0].op is Ops.FUNCTION else self.src[0]
       return in_tuple.src[self.arg].axis if in_tuple.op is Ops.TUPLE else None
-    if self.op is Ops.PARAM: return self.arg.axis
+    if self.op is Ops.PARAM: return None
     # NOTE: they all have to share an axis, we always choose [-1]. src axes are right-aligned into the output shape
     if self.op in GroupOp.ALU.union({Ops.STACK}):
       return axes[-1] if (axes := dedup([x.axis+len(self.shape)-len(x.shape) for x in self.src if x.axis is not None])) else None
@@ -804,7 +799,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       case Ops.PERMUTE | Ops.FLIP: src_args = []
       case Ops.STACK:
         srcs = (self,)+tuple(arg)
-        dtype = cast(DType, dtype_from_uop(Ops.STACK, srcs, None))
+        dtype = dtype_from_uop(Ops.STACK, srcs, None)
         return UOp(Ops.STACK, src=tuple(u if u.base.is_invalid else UOp.const(u.val, dtype) if u.op is Ops.CONST else u.cast(dtype) for u in srcs))
       case _: raise RuntimeError(f"{op} is not a MovementOp")
     usrcs = [shape_to_shape_arg(arg) for arg in src_args]
@@ -821,8 +816,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   @staticmethod
   def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
     if dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {dtype}")
+    assert isinstance(size, int), f"new_buffer size must be a concrete int, got {size}"
     slot = next(UOp.unique_num) if num is None else num
-    return UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(slot, dtype, device=device))
+    return UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, size=size, device=device))
   @staticmethod
   def from_buffer(opaque:Buffer, device:str|tuple[str, ...]|None=None):
     if (uop:=UOp.new_buffer(device or opaque.device, opaque.size, opaque.dtype, num=-id(opaque))) not in buffers: buffers[uop] = opaque.ref(1)
@@ -978,7 +974,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     # a Variable is a 0-d BUFFER in the ALU addrspace; binding it is storing a CONST into it
     # param=True creates the kernel-side form directly: an ALU PARAM (what the BUFFER becomes inside kernels)
     arg = ParamArg(-1, dtype, name=name, vmin_vmax=(min_val, max_val), multiple_of=multiple_of, addrspace=AddrSpace.ALU)
-    return UOp(Ops.PARAM if param else Ops.BUFFER, src=(shape_to_shape_arg(()),), arg=arg)
+    return UOp(Ops.PARAM if param else Ops.BUFFER, arg=arg)
   @property
   def is_variable(self) -> bool:
     # a Variable is a 0-d BUFFER in the ALU addrspace that carries a value range (it becomes a PARAM inside kernels)
@@ -1148,11 +1144,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     dtype = strong_dtype(dtype)  # storage is never weak: a placeholder commits the width of what's put in it
     if slot is None: slot = next(UOp.unique_num)
     if addrspace is AddrSpace.GLOBAL:
-      ret = UOp(Ops.PARAM, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace, device=device,volatile=volatile))
+      ret = UOp(Ops.PARAM, arg=ParamArg(slot, dtype, size=prod(shape), addrspace=addrspace, device=device, volatile=volatile))
     else:
       assert addrspace in (AddrSpace.LOCAL, AddrSpace.REG)
       assert device is None, "LOCAL and REG placeholders cannot have a device"
-      ret = UOp(Ops.BUFFER, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace))
+      ret = UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, size=prod(shape), addrspace=addrspace))
     if tag is not None: ret = ret.rtag(tag)
     if len(shape) > 1: ret = ret.reshape(shape)
     return ret
@@ -1166,20 +1162,35 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   # TODO: this should replace placeholder
   @staticmethod
-  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None,
-            multiple_of:int|None=None, name=None, addrspace=AddrSpace.GLOBAL, axis:int|None=None, volatile:bool=False):
+  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|sint|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None,
+            multiple_of:int|None=None, name=None, addrspace=AddrSpace.GLOBAL, volatile:bool=False):
+    """create a PARAM: a single sint or 1-d shape gives a flat param of that size, a None shape gives a scalar param.
+    the arg only stores the concrete max size (never symbolic): a multi-dim shape is a RESHAPE on top of the flat param,
+    a symbolic shape is a max-size param shrunk to the real shape"""
     if dtype in dtypes.weaks: raise RuntimeError(f"cannot create param for weak dtype {dtype}")
-    if shape is not None and axis is not None and isinstance(device, tuple):
-      shape = tuple(s*len(device) if i == axis else s for i,s in enumerate(shape))
-    src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
-    return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
+    if isinstance(shape, (int, UOp)): shape = (shape,)
+    if shape is None or len(shape) == 0:
+      return UOp(Ops.PARAM, arg=ParamArg(slot, dtype, None, vmin_vmax, multiple_of, name, addrspace, device, volatile))
+    max_shape = to_max_shape(shape)
+    ret = UOp(Ops.PARAM, arg=ParamArg(slot, dtype, prod(max_shape), vmin_vmax, multiple_of, name, addrspace, device, volatile))
+    return ret.view_as(shape)
   def param_like(self, slot:int):
     # Variables become ALU params in the call body; the stored value (if bound) stays in the call args
     if self.is_bound_var or self.is_variable:
       b = self.src[0] if self.op is Ops.AFTER else self
-      return UOp(Ops.PARAM, src=b.src, arg=replace(b.arg, slot=slot, name=f"p{slot}"))
+      return UOp(Ops.PARAM, arg=replace(b.arg, slot=slot, name=f"p{slot}"))
     addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
-    return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)
+    # multi-device values become a per-shard sized param wrapped in UNSHARD: the sharding lives in the graph, not the arg
+    if self.axis is not None and isinstance(self.device, tuple):
+      return UOp(Ops.PARAM, arg=ParamArg(slot, self.dtype, prod(to_max_shape(self.shard_shape)),
+                                         addrspace=addrspace, device=self.device)).view_as(self.shard_shape, self.axis)
+    return UOp.param(slot, self.dtype, self._shape, self.device, addrspace=addrspace)
+  def view_as(self:UOp, shape:tuple[sint, ...], axis:int|None=None) -> UOp:
+    """view flat storage as the given (possibly symbolic) shape, optionally sharded on axis, the UNSHARD gives back the multiplied shape"""
+    max_shape = to_max_shape(shape)
+    ret = self.reshape(max_shape) if len(shape) > 1 else self
+    if tuple(max_shape) != tuple(shape): ret = ret.shrink_to(shape)
+    return ret if axis is None else ret.unshard(axis)
 
   @staticmethod
   def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)
@@ -1188,7 +1199,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.CUSTOM_FUNCTION}
   def call(self, *srcs:UOp, ret_dtype:DType|None=None, grad_fxn:Callable|None=None,
            name:str|None=None, precompile:bool=False, precompile_backward:bool=False, aux:Any=None) -> UOp:
-    if ret_dtype is not None: return UOp(Ops.CALL, ret_dtype, src=(self,)+srcs)
+    if ret_dtype is not None: return UOp(Ops.CALL, src=(self,)+srcs, arg=ret_dtype)
     # calls are launched per device, so an open DEVICE range is allowed to cross the call boundary
     assert all(r.arg[-1] is AxisType.DEVICE for r in self.ranges), \
       f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
@@ -1675,7 +1686,7 @@ class RewriteContext:
       else:
         # rebuild node with rewritten srcs
         new_src = tuple(self.replace.get(x, x) for x in n.src)
-        new_n = UOp(n.op, _rebuild_dtype(n, new_src), new_src, n.arg, n.tag) if new_src != n.src else n
+        new_n = UOp(n.op, new_src, n.arg, n.tag) if new_src != n.src else n
         # top-down: try pm on rebuilt node, use result as-is (no re-traversal)
         if self.pm is not None and (rewritten:=self.pm_rewrite(new_n)) is not None: new_n = rewritten
         self.replace[n] = new_n
@@ -1734,7 +1745,7 @@ class RewriteContext:
               continue
           else:
             # if srcs changed from rewrites, construct a new UOp with the new srcs
-            new_src_n = UOp(new_n.op, _rebuild_dtype(new_n, new_src), new_src, new_n.arg, new_n.tag)
+            new_src_n = UOp(new_n.op, new_src, new_n.arg, new_n.tag)
           # trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n
           stack.append((n, 2, new_src_n))
           stack.append((new_src_n, 0, new_src_n))
@@ -1754,11 +1765,6 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N
   rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls)
   return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)
 
-def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType:
-  # TODO: delete this once the dtype field is removed, every rebuild will re-derive
-  if all(a.dtype is b.dtype for a,b in zip(n.src, new_src)): return n.dtype
-  return dtype_from_uop(n.op, new_src, n.arg) or n.dtype
-
 def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(x, dtype)
 def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape)
 

tinygrad/uop/render.py

@@ -11,7 +11,7 @@ def pretty_print(x:UOp, cache=None, d=0)->str:
   if cache is None: dfs(x, cache:={})
   if (cx:=cache.setdefault(x, [0,0,False]))[2]: return f"{' '*d}x{cx[0]}"
   cx[2], srcs = True, (''.join(f'\n{pretty_print(s, cache, d+2)},' for s in x.src))
-  return f"{' '*d}{f'x{cx[0]}:=' * (cx[1]>1)}{type(x).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=({srcs}))"
+  return f"{' '*d}{f'x{cx[0]}:=' * (cx[1]>1)}{type(x).__name__}({x.op}, arg={x.argstr()}{x.tagstr()}, src=({srcs}))"
 
 # ***** uop helpers *****
 
@@ -84,7 +84,7 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.val})"),
   (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})" if x.dtype != x.src[0].dtype else None),
   (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].val}, {repr(x.arg)})"),
-  (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x:
+  (UPat(Ops.BUFFER, src=(), name="x"), lambda x:
     f"UOp.new_buffer({repr(x.arg.device)}, {x.max_numel()}, {x.dtype}, {x.arg.slot})"
     if isinstance(x.arg, ParamArg) and x.addrspace is AddrSpace.GLOBAL else None),
   (UPat(Ops.COPY, src=(UPat(name="x"),), name="copy"), lambda ctx,x,copy: f"{ctx[x]}.copy_to_device({repr(copy.arg)})"),
@@ -94,10 +94,6 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
     "UOp.range("+', '.join([str(c.val)] + [repr(y) for y in x.arg])+
       (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+")"),
-  # TODO: index shouldn't mismatch dtype
-  (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
-   f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
-    f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else None),
   # TODO: movement ops simplify stuff, this can break SPEC=2
   #(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
   # NOTE: CMPNE doesn't work cause there's no __rne__
@@ -117,7 +113,7 @@ pm_pyrender_extra = PatternMatcher([
 
 # NOTE: you can remove pm_pyrender_extra and it'll still be correct
 pm_pyrender = pm_pyrender_extra+PatternMatcher([
-  (UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
+  (UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
 ])
 
 def _render_with_splits(lst:list[UOp], pm:PatternMatcher, to_render:set[UOp], split_depth:int=100) -> dict[str, str]:

tinygrad/uop/spec.py

@@ -87,10 +87,9 @@ spec_shared = PatternMatcher([
   # a loop-ended END requires a trailing bool condition for the backedge (loop again while true)
   (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void), UPat(dtype=dtypes.bool))), lambda: True),
 
-  # PARAM
-  (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
-  (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x:
-   isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.REG, AddrSpace.LOCAL)),
+  # PARAM/BUFFER have a size in the arg, no shape input
+  (UPat(Ops.PARAM, src=(), name="x"), lambda x: isinstance(x.arg, ParamArg)),
+  (UPat(Ops.BUFFER, src=(), name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.REG, AddrSpace.LOCAL)),
 
   # GROUP of stores (or groups, or NOOPs)
   (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.INS, Ops.END))), lambda: True),
@@ -141,12 +140,12 @@ spec_tensor = PatternMatcher([
    lambda u: dtypes.is_float(u.dtype) or u.src[0].base.is_invalid),
 
   # BUFFER
-  (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf:
-   (isinstance(buf.dtype, DType) and matches_dtype(buf.src[0], dtypes.weakint) and is_device(buf.arg.device))
+  (UPat(Ops.BUFFER, src=(), name="buf"), lambda buf:
+   (isinstance(buf.dtype, DType) and isinstance(buf.arg.size, int) and is_device(buf.arg.device))
    if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None),
 
   # a Variable is a 0-d ALU BUFFER with a value range and no device
-  (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf: buf.arg.device is None if buf.is_variable else None),
+  (UPat(Ops.BUFFER, src=(), name="buf"), lambda buf: buf.arg.device is None if buf.is_variable else None),
 
   # custom function
   (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),
@@ -229,7 +228,8 @@ spec_program = PatternMatcher([
 ])+spec_shared
 
 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.GETADDR, dtypes.uint64, src=(UPat((Ops.BUFFER, Ops.PARAM, Ops.SHRINK, Ops.BITCAST, Ops.MSTACK)).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),
 ])+spec_shared
 
@@ -260,11 +260,10 @@ spec_kernel_graph = PatternMatcher([
   (UPat(Ops.STACK, name="s"), lambda s: all(x.op in (Ops.CONST, Ops.PARAM) or x.is_variable or x.is_bound_var for x in s.src) or None),
   # linear for more kernels (TODO: we should enter non sink calls)
   #(UPat(Ops.LINEAR), lambda: True),
-  # param is outside buffer, buffer is local buffer
-  (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
+  # param is outside buffer, buffer is local buffer. params have a size in the arg, no shape input
+  (UPat(Ops.PARAM, src=(), name="x"), lambda x: isinstance(x.arg, ParamArg)),
   (UPat(Ops.BUFFER, name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.GLOBAL, AddrSpace.ALU)),
-  # RESHAPE/BITCAST are NOOPs in the kernel graph (do we need them?)
-  (UPat((Ops.RESHAPE, Ops.BITCAST)), lambda: True),
+  (UPat(Ops.BITCAST), lambda: True),
   # mstack/mselect
   (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),
   (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),

tinygrad/uop/symbolic.py

@@ -150,7 +150,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
   (UPat(GroupOp.ALU-{Ops.THREEFRY}, src=bare_const, name="a"), fold_const_alu),
   (UPat(GroupOp.ALU-{Ops.THREEFRY}, src=casted_const, name="a"), fold_const_alu),
   (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=[casted_const, bare_const], name="a"), lambda a:
-   a.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in a.src))
+   a.replace(src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in a.src))
    if (dt:=promo_dtype(a.src)) not in dtypes.weaks else None),
   # bool MUL is AND, ADD/MAX is OR. prevents other rules to rewrite bool ADD/MUL incorrectly
   (UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), lambda x,y: x&y),

tinygrad/uop/weak.py

@@ -24,7 +24,7 @@ def derived_dtypes(u:UOp, src:tuple[UOp, ...]) -> tuple[DType, DType]|None:
 def commit_srcs_at(u:UOp, dt:DType) -> UOp|None:
   # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too
   dts = derived_dtypes(u, u.src)
-  ret = u.replace(dtype=None, src=tuple(UOp.const(dt.const(s.val)) if s.op is Ops.CONST and s.dtype in dtypes.weaks and dts is not None else
+  ret = u.replace(src=tuple(UOp.const(dt.const(s.val)) if s.op is Ops.CONST and s.dtype in dtypes.weaks and dts is not None else
                                         commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
   return None if ret is u else ret
 
@@ -61,9 +61,9 @@ def lower_weak_node(u:UOp) -> UOp|None:
   if u.op in _lower_weak_ops and src != u.src and not any(s.dtype in dtypes.weaks and s.op is not Ops.CONST for s in src[start:]):
     dt = strong_dtype(least_upper_dtype(default_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
                       else unwrap(dtype_from_uop(u.op, src, u.arg)))
-    return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid or s.dtype in dtypes.weaks else commit_weak(s, dt)
+    return u.replace(src=src[:start]+tuple(s if s.base.is_invalid or s.dtype in dtypes.weaks else commit_weak(s, dt)
                                                        for s in src[start:])).cast(u.dtype)
-  return None if src == u.src else u.replace(dtype=None, src=src)
+  return None if src == u.src else u.replace(src=src)
 
 pm_lower_weak = PatternMatcher([
   # a gated long index into a small buffer narrows; its out-of-gate value is discarded
@@ -74,7 +74,7 @@ pm_lower_weak = PatternMatcher([
   (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),
    lambda u,x: x.cast(default_dtype(u.src[0])).cast(default_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
   (UPat((Ops.PARAM, Ops.BUFFER), dtype=dtypes.weakint, name="u"),
-    lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=default_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
+    lambda u: u.replace(arg=replace(u.arg, dtype=default_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
   (UPat(GroupOp.All, name="u"), lower_weak_node),
 ])
 

tinygrad/viz/serve.py

@@ -171,9 +171,9 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
 
 def _reconstruct(data:VizData, a:int, depth:int|None=None):
   if depth is None and a in data.all_uops: return data.all_uops[a]
-  op, dtype, src, arg, *rest = data.trace.uop_fields[a]
-  if depth is not None and depth <= 0: return UOp(op, dtype, (), arg, *rest)
-  ret = UOp(op, dtype, tuple(_reconstruct(data, s, None if depth is None else depth-1) for s in src), arg, *rest)
+  op, src, arg, *rest = data.trace.uop_fields[a]
+  if depth is not None and depth <= 0: return UOp(op, (), arg, *rest)
+  ret = UOp(op, tuple(_reconstruct(data, s, None if depth is None else depth-1) for s in src), arg, *rest)
   if depth is None: data.all_uops[a] = ret
   return ret