direct lds

.github/workflows/benchmark.yml

@@ -515,12 +515,14 @@ jobs:
       run: |
         echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: openpilot compile3 big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo PICKLE_OOB=1 PYTHONPATH="." TC_OPT=2 GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/10926f2c0911821ca0e72439c1c3bf3ec11f0a08789aa14b7ee8f25379b2afa4 openpilot.pkl
-    - name: openpilot load_pickle big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_load_pickle PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=25 python3 examples/openpilot/load_pickle.py openpilot.pkl
-    - name: openpilot run_pickle big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
+    - name: reset chestnut
+      run: python3 extra/usbgpu/debug.py -rn
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot load_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+    - name: openpilot run_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py
     - name: Test copy speeds
       run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 

.github/workflows/sticky_comment.py

@@ -1,34 +0,0 @@
-#!/usr/bin/env python3
-# Sticky PR comment via the REST API: find an existing comment containing MARKER and PATCH it, or POST a new one.
-# Works on GitHub and Gitea (stdlib only, replaces marocchino/sticky-pull-request-comment which needs GraphQL).
-# Env vars: GITHUB_TOKEN, GITHUB_API_URL, GITHUB_REPOSITORY (set by the runner), PR_NUMBER, MARKER, and BODY_FILE or MESSAGE.
-import json, os, sys, urllib.request
-
-api, repo = os.environ["GITHUB_API_URL"], os.environ["GITHUB_REPOSITORY"]
-pr, marker = os.environ["PR_NUMBER"], os.environ["MARKER"]
-body = open(os.environ["BODY_FILE"]).read() if os.environ.get("BODY_FILE") else os.environ["MESSAGE"]
-
-if not body.strip():
-  print("comment body is empty, not posting")
-  sys.exit(0)
-
-def req(url, method="GET", payload=None):
-  r = urllib.request.Request(url, data=None if payload is None else json.dumps(payload).encode(), method=method,
-    headers={"Authorization": f"token {os.environ['GITHUB_TOKEN']}", "Accept": "application/json", "Content-Type": "application/json"})
-  return json.load(urllib.request.urlopen(r))
-
-# find the latest sticky comment (paginate, 100 comments per page)
-existing, page = None, 1
-while True:
-  comments = req(f"{api}/repos/{repo}/issues/{pr}/comments?per_page=100&page={page}")
-  stickies = [c for c in comments if marker in (c.get("body") or "")]
-  if stickies: existing = stickies[-1]
-  if not comments or len(comments) < 100: break
-  page += 1
-
-if existing is not None and existing["body"] == body:
-  print("comment is already up to date")
-  sys.exit(0)
-url = f"{api}/repos/{repo}/issues/comments/{existing['id']}" if existing is not None else f"{api}/repos/{repo}/issues/{pr}/comments"
-resp = req(url, 'PATCH' if existing is not None else 'POST', {'body': body})
-print(f"{'updated' if existing is not None else 'created'} comment {resp['id']}")

.github/workflows/szdiff.yml

@@ -26,14 +26,14 @@ jobs:
       - name: Check whether branch is up-to-date
         id: brstat
         run: |
-          # fetch master from the base repo (tinygrad/tinygrad on GitHub, the mirror on Gitea), not the PR head remote
-          git fetch "${{ github.event.pull_request.base.repo.clone_url }}" master
+          git remote add tinygrad https://github.com/tinygrad/tinygrad
+          git fetch tinygrad master
           echo "${{ github.event.pull_request.head.sha }}"
-          git rev-list --left-right --count FETCH_HEAD...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}'
-          count=$(git rev-list --left-right --count FETCH_HEAD...${{ github.event.pull_request.head.sha }} | awk '{print $1}')
+          git rev-list --left-right --count  tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}'
+          count=$(git rev-list --left-right --count  tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print $1}')
           if [ $count -gt 0 ]
           then
-            echo "Current branch is behind ${{ github.event.pull_request.base.repo.full_name }} master branch!"
+            echo "Current branch is behind tinygrad master branch!"
             echo "stat=true" >> "$GITHUB_OUTPUT"
           else
             echo "stat=false" >> "$GITHUB_OUTPUT"
@@ -75,13 +75,13 @@ jobs:
           python sz.py "$BASE" "$PR" > loc_content.txt
       - name: Comment Code Line Diff
         continue-on-error: false
-        env:
+        uses: marocchino/sticky-pull-request-comment@v3
+        with:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          PR_NUMBER: ${{ github.event.pull_request.number }}
-          MARKER: "### Changes"
-          BODY_FILE: loc_content.txt
-        # note: run the script from the base checkout, never from the PR checkout
-        run: python3 "$GITHUB_WORKSPACE/base/.github/workflows/sticky_comment.py"
+          ignore_empty: true
+          skip_unchanged: true
+          recreate: true
+          path: loc_content.txt
 
   rebase:
     name: Core Library Line Difference
@@ -91,14 +91,12 @@ jobs:
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'true'
     steps:
-      # pull_request_target: a plain checkout gets the base repo, so no PR code is executed
-      - uses: actions/checkout@v6
       - name: Comment Rebase
         continue-on-error: false
-        env:
+        uses: marocchino/sticky-pull-request-comment@v3
+        with:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          PR_NUMBER: ${{ github.event.pull_request.number }}
-          MARKER: "line count difference bot is disabled"
-          MESSAGE: |
-            This branch currently is behind ${{ github.event.pull_request.base.repo.full_name }} master. The line count difference bot is disabled.
-        run: python3 .github/workflows/sticky_comment.py
+          skip_unchanged: true
+          recreate: true
+          message: |
+            This branch currently is behind tinygrad/master. The line count difference bot is disabled.

.github/workflows/test.yml

@@ -353,8 +353,6 @@ jobs:
         run: DEV=NULL NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
       - name: Test llama 3 training
         run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
-      - name: Test gpt-oss training
-        run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=32 BS=2 SEQLEN=128 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 MXFP8=1 VOCAB_SIZE=32000 LAYERS=2 EXPERTS=4 MODEL=gptoss PYTHONPATH=. python3 examples/mlperf/model_train.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -643,8 +641,12 @@ jobs:
       with:
         key: unittest-macos
         deps: testing_unit
+        amd: 'true'
+        ocelot: 'true'
     - name: Run unit tests
       run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
     - name: Test tensor core ops (fake)
       run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
     - name: Test tensor core ops (real)
@@ -655,25 +657,6 @@ jobs:
       run: DEV=METAL python3 -m pytest test/device/test_metal.py
     #- name: Fuzz Test linearizer
     #  run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  unittestmacosmock:
-    name: MacOS (unit, mock)
-    runs-on: macos-26
-    timeout-minutes: 20
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: unittest-macos-mock
-        deps: testing_unit
-        amd: 'true'
-        ocelot: 'true'
-    - name: Run NULL backend tests
-      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
     - name: Run pytest (amd)
       env:
         DEV: MOCKKFD+AMD

examples/gpt2.py

@@ -22,6 +22,10 @@ class Attention:
     self.head_dim = dim // n_heads
 
   def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor:
+    if mask is not None or start_pos.val == 0:
+      # no symbolic shape qkv when consuming prompts
+      start_pos = start_pos.val
+
     if HALF: x = x.half()
     xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
     xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
@@ -34,8 +38,12 @@ class Attention:
     # update the cache
     self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
 
-    keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
-    values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+    if start_pos > 0:
+      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
+      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+    else:
+      keys = xk
+      values = xv
 
     xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
     return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, self.dim))
@@ -78,14 +86,15 @@ class Transformer:
       seqlen = tokens.shape[1]
       tok_emb = self.wte(tokens)
 
-    # start_pos is a bound Variable, so everything below it stays symbolic
-    pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))
+    # not symbolic when consuming the prompt
+    selected_pos = (0, seqlen) if start_pos.val == 0 else (start_pos, start_pos+1)
+    pos_emb = self.wpe(self.allpos.shrink((None, selected_pos)))
 
     h = tok_emb + pos_emb
 
     if HALF: h = h.half()
 
-    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None
+    mask = Tensor.full((1, 1, seqlen, start_pos.val+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos.val+1) if seqlen > 1 else None
 
     for hi in self.h: h = hi(h, start_pos, mask)
 

examples/mlperf/model_train.py

@@ -1711,10 +1711,9 @@ def train_gptoss():
     wandb.init(config=config, **wandb_args, project="MLPerf-gpt-oss")
 
   model_params = GPT_OSS_20B
-  model_params['vocab_size'] = getenv("VOCAB_SIZE", 128256)
+  model_params['vocab_size'] = 128256
   real_vocab_size = model_params['vocab_size']
   if (layers:=getenv("LAYERS")) != 0: model_params['n_layers'] = layers
-  if (experts:=getenv("EXPERTS")) != 0: model_params['n_experts'] = experts
   print(f"model parameters: {model_params}")
 
   model = GPTOSS(**model_params, max_context=SEQLEN)
@@ -1749,10 +1748,7 @@ def train_gptoss():
 
   from extra.gemm.cdna_asm_gemm import _mx_block_scale
   model_state = get_state_dict(model)
-  def _scale_key(n):
-    if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_scale.{b[1]}") in model_state: return c
-    return f"{n}_scale"
-  fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE}
+  fp8_scale_names = {n: f"{n}_scale" for n, t in model_state.items() if t.dtype == FP8_DTYPE}
   fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()]
   for wname, sname in fp8_scale_names.items():
     w, scale = model_state[wname], model_state[sname]

examples/mlperf/models/flat_llama.py

@@ -25,7 +25,6 @@ FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
 SPLIT_W13 = getenv("SPLIT_W13", 0)
 COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
 MXFP8 = getenv("MXFP8", 0)
-MXFP4 = getenv("MXFP4", 0)
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
@@ -45,11 +44,6 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
       from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
       if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
     return (x @ w.T,)
-  if MXFP4:
-    assert x is not None, "MXFP4 matmul requires an unquantized input"
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
-    if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),)
-    return (x @ w.T,)
   assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
   if MXFP8:
     from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
@@ -85,7 +79,7 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
 
 def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
                          next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor):
-  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
     x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
     out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
@@ -98,7 +92,7 @@ def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, ep
 
 def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
                              next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
-  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
     x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
     out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
@@ -114,7 +108,7 @@ def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
                              amax_x2:Tensor, next_amax_x2:Tensor,
                              grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
                              grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
-  if FUSED_SILU_W13 and not MXFP4:
+  if FUSED_SILU_W13:
     from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
     x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
                                                  next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
@@ -185,9 +179,6 @@ class FlatTransformer:
       from extra.gemm.cdna_asm_gemm import quantize_mxfp8
       w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
       return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
-    if MXFP4:
-      # FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
-      return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
     amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
     scale = FP8_MAX / (amax + 1e-8)
     inv_scale = (amax + 1e-8) / FP8_MAX

examples/mlperf/models/gpt_oss.py

@@ -13,14 +13,10 @@ from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
 from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
-from extra.gemm.moe_gemm import grouped_mx_gemm
-from extra.gemm.moe_routing import route, dispatch, combine
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
-INIT_STD = 0.02
-ASM_GEMM = getenv("ASM_GEMM", 0)
-
+INIT_STD = 0.008
 
 def _quant_dequant_fwd(x:Tensor) -> Tensor:
   # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
@@ -63,34 +59,10 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
 
 def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
   l_shape = x.shape[:-1]
-  if ASM_GEMM:
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
-    x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0]
-    wq, ws = w_q, w_scale
-    if (pad := (-K) % 256):
-      x2 = x2.pad(((0, 0), (0, pad)))
-      wq = wq.pad(((0, 0), (0, pad)))
-      ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
-    if (npad := (-N) % 256):
-      wq = wq.pad(((0, npad), (0, 0)))
-      ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
-    x_q, x_e8, x_si = quantize_mxfp8(x2)
-    if x_si is not None and can_use_asm_gemm(x_q, wq.T):
-      out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True)
-      return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
   x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
   w_phys = dequant_weight(w_q, w_scale)
   return (x_phys @ w_phys.T).cast(dtypes.bfloat16)
 
-def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor:
-  if (r := (-t.shape[axis]) % mult) == 0: return t
-  pads = [(0, 0)] * t.ndim
-  pads[axis] = (0, r)
-  return t.pad(tuple(pads))
-
-def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1)
-def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2)
-
 def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
   x_glu, x_linear = x[..., ::2], x[..., 1::2]
   x_glu = x_glu.clamp(max_=limit)
@@ -127,9 +99,9 @@ class GPTOSS:
     self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
     self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
     self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
-    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True)
+    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim)
     self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
-    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True)
+    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std)
     self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
 
     # output
@@ -139,15 +111,10 @@ class GPTOSS:
     self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
     self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)
 
-  def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False):
-    def _one(*s:int):
-      w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std)
-      w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w)
-      return w_q, w_e8.is_param_(False)
-    if moe:
-      qs = [_one(*shape[1:]) for _ in range(shape[0])]
-      return [q[0] for q in qs], [q[1] for q in qs]
-    return _one(*shape)
+  def _quant_weight(self, *shape:int, std:float=INIT_STD):
+    w = Tensor.zeros(*shape) if getenv("ZEROS") else Tensor.normal(*shape, mean=0.0, std=std)
+    w_q, w_e8, _ = quantize_mxfp8(w)
+    return w_q, w_e8.is_param_(False)
 
   def _attn_mask(self, seqlen:int, dtype) -> Tensor:
     i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
@@ -206,32 +173,17 @@ class GPTOSS:
                    w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
     x_normed, rrms = rmsnorm(x, self.norm_eps)
     inp = x_normed * ffn_norm
+
     logits = inp.float() @ gate.float().T + gate_bias.float()
-    dim, inter = self.dim, self.intermediate_size
+    thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
+    weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
 
-    if getenv("GROUPED_MOE", 0):
-      bsz, seqlen = x.shape[:2]
-      inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts)
-      r = route(logits, self.experts_per_tok, self.n_experts)
-      onehot = r.rows_e.one_hot(self.n_experts).float()
-      xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r)
-      h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16)
-      y = swiglu(h, self.swiglu_limit)
-      z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
-          + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
-      out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
-    else:
-      thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
-      weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
-
-      out = None
-      for e in range(self.n_experts):
-        gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous()
-        dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous()
-        gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e]
-        y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous()
-        contrib = weights[..., e:e+1].cast(y.dtype) * y
-        out = contrib if out is None else out + contrib
+    out = None
+    for e in range(self.n_experts):
+      gate_up = matmul_mx(inp, w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e]
+      y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), w_down[e], w_down_scale[e]) + w_down_bias[e]).contiguous()
+      contrib = weights[..., e:e+1].cast(y.dtype) * y
+      out = contrib if out is None else out + contrib
     return out, [x_normed, rrms]
 
   @function(precompile=True, precompile_backward=True)

examples/mlperf/optim.py

@@ -2,7 +2,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.nn.optim import Optimizer, OptimizerGroup
 from tinygrad.helpers import FUSE_OPTIM, getenv
-from tinygrad.uop.ops import UOp, Ops, AxisType
+from tinygrad.uop.ops import UOp, Ops
 
 STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
 MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
@@ -42,8 +42,8 @@ class GradAccClipAdamW(Optimizer):
       self.master_params = None
 
   def _zero_shard(self, t:Tensor) -> Tensor:
-    if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone()
+    if not self.zero or (t.shape[0] % len(self.device)) != 0: return t
+    return Tensor(t.uop._shard(0, len(self.device)).unshard(0)).clone()
 
   def _zero_gather(self, t:Tensor) -> Tensor:
     if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
@@ -123,9 +123,6 @@ class GradAccClipAdamW(Optimizer):
     return out.shard_like(t) if offloaded else out
 
 class GradAccClipAdamWGroup(OptimizerGroup):
-  def __init__(self, *optimizers:GradAccClipAdamW):
-    super().__init__(*optimizers)
-    for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr
   def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
     offset = 0
     to_realize = []

examples/openpilot/compile3.py

@@ -1,4 +1,4 @@
-import os, sys, pickle, time, re, tempfile, struct, shutil, io
+import os, sys, pickle, time, re
 import numpy as np
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 
@@ -9,39 +9,6 @@ from tinygrad.nn.onnx import OnnxRunner
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
-PICKLE_OOB = getenv("PICKLE_OOB")
-
-def dump_pickle(obj, f):
-  if PICKLE_OOB:
-    # allows pickling when buffers don't fit in (CPU) RAM
-    # from openpilot/selfdrive/modeld/helpers.py
-    with tempfile.TemporaryFile(dir=".") as tmp:
-      def buffer_callback(pb: pickle.PickleBuffer):
-        m = pb.raw()
-        tmp.write(struct.pack('<q', m.nbytes))
-        tmp.write(m)
-        pb.release() # keep peak ram at ~1 buffer
-      stream = io.BytesIO()
-      pickle.Pickler(stream, protocol=5, buffer_callback=buffer_callback).dump(obj)
-      opcodes = stream.getvalue()
-      f.write(struct.pack('<q', len(opcodes)))
-      f.write(opcodes)
-      tmp.seek(0)
-      shutil.copyfileobj(tmp, f)
-  else: pickle.dump(obj, f)
-
-def load_pickle(f):
-  if PICKLE_OOB:
-    # allows unpickling when buffers don't fit in (CPU) RAM
-    # from openpilot/selfdrive/modeld/helpers.py
-    opcodes = f.read(struct.unpack('<q', f.read(8))[0])
-    def buffers():
-      while (h := f.read(8)):
-        pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
-        f.readinto(pb)
-        yield pb
-    return pickle.load(io.BytesIO(opcodes), buffers=buffers())
-  else: return pickle.load(f)
 
 def compile(onnx_file):
   run_onnx = OnnxRunner(onnx_file)
@@ -98,7 +65,8 @@ def compile(onnx_file):
   if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
     assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
 
-  with open(OUTPUT, "wb") as f: dump_pickle(run_onnx_jit, f)
+  with open(OUTPUT, "wb") as f:
+    pickle.dump(run_onnx_jit, f)
   mdl_sz = os.path.getsize(onnx_file)
   pkl_sz = os.path.getsize(OUTPUT)
   print(f"mdl size is {mdl_sz/1e6:.2f}M")
@@ -168,7 +136,7 @@ def bench(run, inputs):
 
 if __name__ == "__main__":
   if getenv("RUN_PICKLE"):
-    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
     inputs = {name: Tensor(Tensor.randn(*view.shape, dtype=dtype).numpy(), device=device)
               for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
     test_vs_compile(pickle_loaded, inputs)
@@ -176,7 +144,7 @@ if __name__ == "__main__":
     onnx_file = fetch(OPENPILOT_MODEL)
     inputs, outputs = compile(onnx_file)
 
-    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
 
     test_vs_compile(pickle_loaded, inputs, outputs)
     if getenv("SELFTEST"):

examples/openpilot/load_pickle.py

@@ -1,6 +1,5 @@
-import sys
+import sys, pickle
 from extra.bench_log import WallTimeEvent, BenchEvent
-from examples.openpilot.compile3 import load_pickle
 from tinygrad.helpers import getenv
 
 PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
@@ -8,7 +7,7 @@ PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
 load_times = []
 
 for _ in range(10):
-  with WallTimeEvent(BenchEvent.STEP) as wte: load_pickle(open(PKL, 'rb'))
+  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
   load_times.append(wte.time)
   print(f"pickle load: {wte.time:6.2f} s")
 

extra/export_model.py

@@ -264,7 +264,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
         if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and \
            any(s.op is Ops.PARAM and s.addrspace is AddrSpace.ALU for s in dim.src) and any(s.op is Ops.CONST for s in dim.src):
           name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src)
-          global_size[j] = f"_{name.expr}[0] + {val.val}"
+          global_size[j] = f"_{name.expr}[0] + {val.arg}"
 
   prg = ""
   if target == "clang":

extra/gemm/cdna_asm_gemm.py

@@ -1,9 +1,8 @@
-import atexit, functools, math, pathlib
+import atexit, functools, pathlib
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
+from tinygrad.helpers import getenv, all_same, DEBUG
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
 
@@ -108,23 +107,6 @@ def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:U
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                                UOp(Ops.BINARY, arg=lib)))
 
-# ** MXFP4 GEMM custom kernel
-
-@functools.cache
-def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp, tile_m:int, tile_n:int) -> UOp:
-  from extra.gemm.gemm_mxfp4 import build_kernel
-  M, half_k = math.prod(A.shape[:-1]), A.shape[-1]
-  N, half_k_b = math.prod(B.shape[:-1]), B.shape[-1]
-  K = half_k * 2
-  assert half_k == half_k_b and math.prod(C.shape[:-1]) == M and C.shape[-1] == N
-  threads = UOp.special(256, "lidx0")
-  groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
-  lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
-  sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
-                  arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
-  insts = build_kernel(M, N, K, tile_m, tile_n)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
-
 def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
   # 1x32 block scaling along the last axis
   *batch, K = x.shape
@@ -137,50 +119,6 @@ def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
   packed = mx_pack(e8) if len(batch) == 1 and scale_K % 4 == 0 else None
   return x_clamped.cast(FP8_DTYPE), e8, packed
 
-def _mxfp4_shuffle_weight(x:Tensor) -> Tensor:
-  # shuffle_weight(x, layout=(16, 16)) on the packed uint8 buffer.
-  if x.ndim == 3:
-    ndev, rows, half_k = x.shape
-    return x.reshape(ndev, rows//16, 16, half_k//32, 2, 16).permute(0, 1, 3, 4, 2, 5).reshape(ndev, rows, half_k).contiguous()
-  rows, half_k = x.shape
-  return x.reshape(rows//16, 16, half_k//32, 2, 16).permute(0, 2, 3, 1, 4).reshape(rows, half_k).contiguous()
-
-def _mxfp4_shuffle_scales(x:Tensor) -> Tensor:
-  # e8m0_shuffle: each 256x8 scale tile is arranged for the raw MFMA scale loads.
-  if x.ndim == 3:
-    ndev, rows, scale_k = x.shape
-    return x.reshape(ndev, rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 1, 4, 6, 3, 5, 2).reshape(ndev, rows, scale_k).contiguous()
-  rows, scale_k = x.shape
-  return x.reshape(rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 3, 5, 2, 4, 1).reshape(rows, scale_k).contiguous()
-
-def quantize_mxfp4(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
-  # OCP MXFP4: 1x32 blocks, E2M1 values packed low-nibble first, and E8M0 scales.
-  *batch, K = x.shape
-  rows = math.prod(batch)
-  assert x.ndim >= 2 and K % 256 == 0 and rows % 32 == 0, \
-    f"mxfp4 quantization needs rows%32 and K%256, got {x.shape}"
-  xb = x.float().reshape(*batch, K//32, 32)
-  amax = xb.abs().max(axis=-1)
-
-  # even scale rounding: round the fp32 significand before choosing 2^(floor(log2)-2).
-  amax_rounded = ((amax.bitcast(dtypes.uint32) + 0x200000) & 0xFF800000).bitcast(dtypes.float32)
-  scale_exp = (amax_rounded.maximum(2**-126).log2().floor() - 2).clamp(-127, 127)
-  e8 = (scale_exp + 127).cast(dtypes.uint8)
-  scaled = xb * (-scale_exp).exp2().reshape(*batch, K//32, 1)
-
-  mag = scaled.abs()
-  code = sum(x.cast(dtypes.uint8) for x in
-             (mag > .25, mag >= .75, mag > 1.25, mag >= 1.75, mag > 2.5, mag >= 3.5, mag > 5.0))
-  code = code | ((scaled < 0).cast(dtypes.uint8) << 3)
-  code = code.reshape(*batch, K)
-  packed = code[..., 0::2] | (code[..., 1::2] << 4)
-  if isinstance(x.device, tuple) and x.uop.axis == x.ndim-2 and x.shape[x.uop.axis] == len(x.device):
-    axis = x.uop.axis
-    order = (axis, *range(axis), *range(axis+1, e8.ndim))
-    e8_local = e8.permute(order)
-    return packed, e8, _mxfp4_shuffle_scales(e8_local.reshape(e8_local.shape[0], -1, K//32))
-  return packed, e8, _mxfp4_shuffle_scales(e8.reshape(rows, K//32))
-
 def mx_pack(e8:Tensor) -> Tensor:
   rows, scale_K = e8.shape
   return e8.reshape(rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(rows, scale_K // 4).permute(1, 0).contiguous()
@@ -236,10 +174,10 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   m = UOp.range(M, 1)
   n = UOp.range(N, 2)
   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)
+  mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
+         B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
   red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
-  store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
+  store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n)).store(red).end(m, n)
   return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
 
 # ** bf16 A @ B.T kernel in C
@@ -306,6 +244,7 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
   if reduce_out: out = out.sum(0)
   return out.squeeze(0) if out.ndim == 3 else out
 
+
 # ** backward gemm, might use the asm gemm
 
 def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=False, has_w_post:bool=False):
@@ -402,29 +341,13 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F
   if wp is not None: grad_b = grad_b / wp.reshape(-1, 1)
   return (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])
 
-# ** mxfp4 gemm backward
-
-def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp):
-  # The raw kernel consumes quantized buffers, while the final two inputs retain the BF16 operands for STE gradients.
-  inputs = kernel.src[1:]  # (out, a_q, b_q, scale_a, scale_b, a, w)
-  assert len(inputs) == 7
-  a, w = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device)
-  g = Tensor(gradient, device=a.device)[:a.shape[0]].cast(dtypes.bfloat16)
-  grad_a = asm_gemm(g, w, mxfp4=True)
-  a_flat, g_flat = a.reshape(-1, a.shape[-1]), g.reshape(-1, g.shape[-1])
-  grad_w = asm_gemm(g_flat.T, a_flat, mxfp4=True)
-  return (None, None, None, None, None, grad_a.uop, grad_w.uop)
-
 # ** main gemm function
 
 def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None,
              next_grad_amax_state:Tensor|None=None,
              w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None,
-             a_pretranspose:Tensor|None=None, mxfp4:bool=False) -> Tensor:
+             a_pretranspose:Tensor|None=None) -> Tensor:
   assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
-  if mxfp4:
-    assert not mx and mx_scales is None, "mxfp4 owns quantization; mx/mx_scales are for mxfp8"
-    assert a.dtype == dtypes.bfloat16, f"cannot quantize {a.dtype} to mxfp4"
   counters["used"] += 1
   unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
   if unfold_batch:
@@ -432,7 +355,7 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
     a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
   squeeze = a.ndim == 2
   if squeeze: a = a.unsqueeze(0)
-  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE or mxfp4 else a.dtype
+  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE else a.dtype
 
   batch, M, K = a.shape
   N = b.shape[1]
@@ -455,21 +378,7 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
   renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
   dname, arch = dname.split(":")[0], renderer.target.arch
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    if mxfp4:
-      tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0)
-      fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n)
-      w = b.T
-      if k_sharded:
-        ndev = len(a.device)
-        a_q, _, scale_a = quantize_mxfp4(a.reshape(batch, M, ndev, K))
-        b_q, _, scale_b = quantize_mxfp4(w.reshape(w.shape[0], ndev, K))
-        b_q = _mxfp4_shuffle_weight(b_q.permute(1, 0, 2))
-      else:
-        a_q, _, scale_a = quantize_mxfp4(a.reshape(batch*M, K))
-        b_q, _, scale_b = quantize_mxfp4(w)
-        a_q, b_q = a_q.reshape(batch, M, K//2).contiguous(), _mxfp4_shuffle_weight(b_q)
-      out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0]
-    elif mx:
+    if mx:
       # mxfp8 1x32 block scaling
       if mx_scales is not None:
         a_si, a_e8, b_si, b_e8 = mx_scales

extra/gemm/hk_gemm_frag.py

@@ -0,0 +1,395 @@
+"""
+HipKittens hk_bf16_gemm (extra/thunder/amd/gemm_bf16.cpp) reimplemented with tinygrad UOps.
+
+C[M, N] (bf16) = A[M, K] @ B[N, K]^T, fp32 accumulation, exactly the kittens kernel shape:
+  - 256x256 output tile per workgroup, K_STEP=64
+  - 8 warps in a 2x4 grid, each warp owns a 128x64 warp-tile
+  - v_mfma_f32_16x16x32_bf16 on CDNA4 (gfx950) / v_wmma_f32_16x16x16_bf16 (wave32, gfx12) on RDNA4,
+    fp32 accumulators
+  - shared tiles As/Bs with the kittens st_16x32_s swizzle (16x32 subtiles of 1024B)
+  - K stages (STAGES=1: synchronous single buffer)
+
+Validated on gfx1201 hardware (exact for identity-B, rounding-level noise otherwise), rendered
+and compiled to gfx950 with comgr for assembly comparison against gemm_bf16.cpp.
+
+What is NOT expressible vs the kittens C++:
+  - explicit s_waitcnt vmcnt()/lgkmcnt() pipelining and s_setprio: tinygrad models async copy
+    overlap with slot dependencies and emits full workgroup barriers; instruction scheduling
+    is left to clang/LLVM
+  - direct-to-LDS global loads (buffer_load_lds): tinygrad goes global->reg->LDS
+
+Pipelining status: STAGES=2 gives the kittens-shaped double-buffered pipeline (2 x 64KB LDS
+like gemm_bf16.cpp, copies overlap the previous pair's mma's), written with FA/gemm_fragment
+conventions: LDS buffers are (2, tile) placeholders indexed by symbolic parity (ko % 2),
+which sidesteps static slot choice, fill iterations, predication and duplicate static stores.
+Validated on the CDNA4 emulator for all tile counts (amt = K//64 in {1..32}, odd/even),
+single- AND multi-workgroup (bit-close to stages=1 / to hippkittens at rounding level).
+
+Bug hunt notes (all fixed on this branch; they were entangled for a long time):
+  1. The double-buffered pipeline REGISTER-SPILLS (255+ VGPRs vs 166 for stages=1), and the
+     mock emulator aliased the spill (scratch) segment of ALL waves of a workgroup onto one
+     64-lane region. On real HW each wavefront owns a per-lane segment of the scratch ring
+     (indexed by (wave_id, lane)); waves trampled each other's spilled accumulators, giving
+     the "only the last wave's output survives" signature. emu.py now allocates per-wave
+     scratch buffers.
+  2. The remaining "shape-dependent" corruption (NaNs, mispositioned values in contiguous
+     copies feeding the GEMM) came from tinygrad's devectorizer fusing adjacent bf16 stores
+     into 32-bit stores with UNALIGNED (2-byte) granularity: legal on AMD FLAT/GLOBAL (the
+     hardware splits them), but the emulator floored misaligned addresses to the word below.
+     _mem_store now handles unaligned 32-bit (and wider) accesses byte-exactly.
+  3. memory_coalescing (late/coalesce.py) assumed a single static store per (buffer, index)
+     ("attempting multiple stores"); aliased stores (a double-buffered LDS slot written in a
+     prologue AND a loop body) are now simply kept scalar instead of asserting/merging.
+  4. pm_split_ranges may only split ranges WITHOUT hardware meaning (WEAK/REDUCE/LOOP);
+     splitting LOCAL/WARP/THREAD/GLOBAL/GROUP_REDUCE/UPCAST ranges scrambles the
+     logical<->hardware mapping of hand-written kernels such as this one.
+
+  RDNA4 (gfx12) uses 8-element accumulator fragments, so the 8x4 tile grid needs
+  256 fp32 acc registers per thread -> guaranteed spills (0.85 TF vs 96 TF default on
+  gfx1201). The kernel is right-sized for CDNA4 (fragsz 4 -> 128 acc regs).
+
+Lane layouts (RDNA4 verified with probing on gfx1201 hardware; CDNA from the mfma docs):
+  CDNA (64 thr/warp, 16x16x32):  A/B frag: tile-row = l%16, k = (l//16)*8+i (i in 0..7)
+  RDNA4 (32 thr/warp, 16x16x16): A/B frag: tile-row = l%16, k = (l//16)*4+(i%4)+8*(i//4) (i in 0..7)
+  both:                        acc frag: CDNA m=(l//16)*4+i (i<4) / RDNA4 m=(l//16)*8+i (i<8), n=l%16
+
+The RDNA4 fragment k-set {k0..3, k0+8..11} is not contiguous, so on RDNA4 the LDS column layout
+is block-permuted (4-element blocks within each 16-col group are stored as [0,2,1,3]) making
+every fragment 8 contiguous halves (one 16B chunk) on both archs; the copy path applies the
+same permutation.
+
+NOTE: thread ids come from UOp.special (like mi350x_uop_matmul.py), not an AxisType.LOCAL
+RANGE. (pm_split_ranges now only splits WEAK/REDUCE/LOOP ranges, so LOCAL ranges would
+survive too, but UOp.special is the sanctioned way to tag hardware lane ids.)
+NOTE 2: WMMA operand/accumulator fragments must carry the fragment length in their UOp shape.
+NOTE 3: swizzled addresses are written in provably-contiguous "base + vector-offset" form,
+otherwise the devectorizer emits scalar ds_read_u16/ds_write_b16.
+"""
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.dtype import AddrSpace
+from tinygrad.renderer import Estimates
+from tinygrad.helpers import getenv, cdiv
+
+# ---- tile shape (identical to gemm_bf16.cpp; HK_TILE=128 overrides for small-LDS devices) ----
+BLOCK_M = BLOCK_N = getenv("HK_TILE", 256)
+K_STEP = 64
+WARPS_M, WARPS_N = 2, 4
+NUM_WARPS = WARPS_M * WARPS_N                                # 8
+WARP_TILE_M, WARP_TILE_N = BLOCK_M // WARPS_M, BLOCK_N // WARPS_N   # 128 x 64 (64 x 32 at HK_TILE=128)
+
+def arch_params(arch:str):
+  is_cdna = arch.startswith("gfx9")
+  if is_cdna:   # CDNA mfma 16x16x32 bf16: in-frag 8, acc 4 (16,16,16 on the acc side)
+    return dict(warp_threads=64, dims=(16,16,32), frag_in=8, frag_out=4,
+                acc_m=lambda l, i: (l//16)*4 + i, kperm=None, copy_vec=8)
+  # RDNA4 wmma 16x16x16 (wave32, gfx12 layout): in-frag 8 (permuted in LDS), acc 8
+  return dict(warp_threads=32, dims=(16,16,16), frag_in=8, frag_out=8,
+              acc_m=lambda l, i: (l//16)*8 + i, kperm=(0,2,1,3), copy_vec=4)
+
+# ---- kittens st_16x32_s swizzle (byte offset: off ^ (((off % 1024) >> 9) << 5)) ----
+# halves-index form: within a 1024B (16x32) subtile, halves-index bit4 ^= row bit3,
+# written in "base + vector-offset" form so the devectorizer can prove contiguity.
+def st_half_base(r, c, tile_cols:int):
+  """swizzled halves-index pre-vector-offset; c is the logical (permuted) column, 4-aligned."""
+  subtile_id = (r//16) * (tile_cols//32) + (c//32)
+  r16 = r % 16
+  flip = (r16 >> 3) & 1
+  return subtile_id*512 + r16*32 + (((c % 32) >> 2) ^ (flip << 2)) * 4
+
+def hk_bf16_gemm_kernel(C:UOp, A:UOp, B:UOp, *, arch:str, stages:int=1) -> UOp:
+  """C = A @ B^T ; A is (M,K), B is (N,K), C is (M,N). HipKittens tile shape."""
+  M, K = A.shape
+  N, K2 = B.shape
+  assert K == K2 and A.dtype == B.dtype == dtypes.bfloat16 and C.dtype == dtypes.bfloat16
+  assert not (M % BLOCK_M or N % BLOCK_N or K % K_STEP), f"dims must be multiples of {(BLOCK_M, BLOCK_N, K_STEP)}"
+
+  ap = arch_params(arch)
+  warp_threads, dims = ap["warp_threads"], ap["dims"]
+  FRAG_IN, FRAG_OUT, kperm, acc_m, CPV = ap["frag_in"], ap["frag_out"], ap["kperm"], ap["acc_m"], ap["copy_vec"]
+  NUM_THREADS = NUM_WARPS * warp_threads
+  TC_M, TC_N, TC_K = dims
+  MT, NT = WARP_TILE_M // TC_M, WARP_TILE_N // TC_N               # 8, 4 tiles per warp
+  # permute 4-half blocks within each 16-col group (RDNA4: [0,2,1,3] = swap middle blocks)
+  def perm_col(c):
+    if kperm is None: return c
+    return (c & ~15) | ((((c>>2) & 1) << 1 | ((c>>3) & 1)) << 2) | (c & 3)
+
+  bx, by = UOp.special(N//BLOCK_N, "gidx0"), UOp.special(M//BLOCK_M, "gidx1")
+  lane = UOp.special(warp_threads, "lidx0")
+  warp = UOp.special(NUM_WARPS, "lidx1")
+  warp_row, warp_col = warp // WARPS_N, warp % WARPS_N
+  tid = warp*warp_threads + lane
+
+  def smem(slot) -> UOp: return UOp.placeholder((BLOCK_M*K_STEP,), dtypes.bfloat16, slot, AddrSpace.LOCAL)
+  As = [smem(2*i) for i in range(stages)]
+  Bs = [smem(2*i+1) for i in range(stages)]
+
+  # per-warp accumulator: (MT x NT) 16x16 tiles of FRAG_OUT fp32 per thread
+  acc = UOp.placeholder((MT, NT, FRAG_OUT), dtypes.float32, 12, AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))   # FA-style init: self-store, keeps the value flow loop-carried
+
+  # global -> LDS copy: CPV halves per op (16B on CDNA, 8B on RDNA4), thread-major coalescing
+  OPS_PER_TILE = BLOCK_M*K_STEP//CPV
+  OPR = K_STEP//CPV
+  def copy_tile(dst:UOp, src:UOp, base_row:UOp, base_col:UOp, slot:int) -> UOp:
+    ir = UOp.range(cdiv(OPS_PER_TILE, NUM_THREADS), slot, AxisType.LOOP)
+    j = UOp.range(CPV, slot+1, AxisType.UPCAST)
+    chunk = ir*NUM_THREADS + tid
+    r, cb = chunk // OPR, chunk % OPR          # row, 4/8-col block
+    return dst[st_half_base(r, perm_col(cb*CPV), K_STEP) + j].store(src[base_row + r, base_col + cb*CPV + j]).end(ir, j)
+
+  def load_stage(sidx:int, ko, slot:int, barrier:bool) -> tuple[UOp, UOp]:
+    A_r = copy_tile(As[sidx], A, by*BLOCK_M, ko*K_STEP, slot)
+    B_r = copy_tile(Bs[sidx], B, bx*BLOCK_N, ko*K_STEP, slot+10)
+    bar = UOp.barrier(A_r, B_r) if barrier else UOp.group(A_r, B_r)
+    return As[sidx].after(bar), Bs[sidx].after(bar)
+
+  # ---- pipelined path (stages=2) ----
+  NIT = cdiv(OPS_PER_TILE, NUM_THREADS)      # copy ops per thread per tile
+  def setprio(n:int, slot:int) -> UOp:
+    """__builtin_amdgcn_s_setprio(n), like gemm_bf16.cpp: raise warp priority for the mma phase
+    so global/LDS traffic of the other waves doesn't starve issue slots."""
+    # distinct src slot per call site so identical-priority instructions at different k-tiles
+    # don't get UOp-hash-deduped into one placement (s_setprio is position-sensitive)
+    return UOp(Ops.CUSTOMI, dtypes.void, src=(UOp.const(dtypes.weakint, slot), UOp.const(dtypes.weakint, n)),
+               arg="__builtin_amdgcn_s_setprio({1}); // {0}")
+
+  def gload_write_tile(dst:UOp, src:UOp, base_row:UOp, kt, slot:int) -> UOp:
+    """store one global tile into an LDS slot (loads and stores share the vec range j)."""
+    j = UOp.range(CPV, slot, AxisType.UPCAST)
+    def one(ir:int) -> UOp:
+      chunk = ir*NUM_THREADS + tid
+      r, cb = chunk // OPR, chunk % OPR
+      return dst[st_half_base(r, perm_col(cb*CPV), K_STEP) + j].store(src[base_row + r, kt*K_STEP + cb*CPV + j])
+    return UOp.group(*[one(ir) for ir in range(NIT)]).end(j)
+
+  def compute(acc:UOp, A_l:UOp, B_l:UOp, afters:tuple[UOp, ...], pred:UOp|None=None, aoff:UOp=None, boff:UOp=None) -> UOp:
+    """One K_STEP=64 iteration: (K_STEP//TC_K) k-chunks unrolled, (MT x NT) mma each, like the kittens main loop.
+
+    pred (optional): a loop-range condition; accumulator stores are predicated on it so the
+    first (fill) iteration of a software pipeline can run the body with garbage LDS contents
+    without contaminating the accumulator."""
+    arow = warp_row*WARP_TILE_M + lane % 16    # fragment tile row in the LDS tile (m)
+    brow = warp_col*WARP_TILE_N + lane % 16    # (n)
+    ja = UOp.range(FRAG_IN, 701, AxisType.UPCAST)
+    jb = UOp.range(FRAG_IN, 702, AxisType.UPCAST)
+    acc_k = acc.after(*afters) if afters else acc
+    last_store = None
+    # in the permuted layout every fragment is 8 contiguous halves starting at an 8-aligned col
+    for kk in range(K_STEP//TC_K):
+      cc = kk*(TC_K//FRAG_IN) + (lane // 16)   # fragment chunk col (8 halves)
+      oa, ob = (aoff, boff) if aoff is not None else (None, None)
+      a_frags = [A_l[st_half_base(arow + mt*16, cc*8, K_STEP) + ja].contract(ja) if oa is None else
+                 A_l[oa + st_half_base(arow + mt*16, cc*8, K_STEP) + ja].contract(ja) for mt in range(MT)]
+      b_frags = [B_l[st_half_base(brow + nt*16, cc*8, K_STEP) + jb].contract(jb) if ob is None else
+                 B_l[ob + st_half_base(brow + nt*16, cc*8, K_STEP) + jb].contract(jb) for nt in range(NT)]
+      for mt in range(MT):
+        for nt in range(NT):
+          cur = acc_k[mt, nt]
+          out = UOp.wmma(a_frags[mt], b_frags[nt], cur, dims, 'AMD', warp_threads)
+          if pred is not None: out = pred.where(cur, out)
+          last_store = acc_k[mt, nt].store(out)
+          acc_k = acc_k.after(last_store)
+    return last_store
+
+  # ---- K loop ----
+  amt = cdiv(K, K_STEP)
+  _stages = stages
+  if _stages == 1:
+    ko = UOp.range(amt, 600, AxisType.LOOP)
+    A_l, B_l = load_stage(0, ko, 100, barrier=True)
+    last = compute(acc, A_l, B_l, afters=(ko,))
+    acc = acc.after(last.barrier().end(ko))
+  else:
+    # Double-buffered pipeline on FA/gemm_fragment conventions: each LDS buffer is a
+    # (2, tile) placeholder indexed by symbolic parity (ko % 2) -- no static slot choice,
+    # no duplicate static stores (memory_coalescing-safe), no fill iteration, no predication.
+    def smem2(slot) -> UOp: return UOp.placeholder((2*BLOCK_M*K_STEP,), dtypes.bfloat16, slot, AddrSpace.LOCAL)
+    A_l, B_l = smem2(0), smem2(1)
+
+    TILE_ELEMS = BLOCK_M * K_STEP
+    def copy_stage(dst:UOp, slot_off:UOp, src:UOp, base_row:UOp, kt, slot:int) -> UOp:
+      """store one global tile into dst + slot_off (flat element offset -- slot_off = parity*TILE_ELEMS).
+
+      Each thread's 8-element chunk is a single buffer_load_lds direct-to-LDS instruction
+      (the kittens '... offen lds' fill path), emitted via Ops.CUSTOMI so it bypasses the
+      devectorizer (a SHRINK store of a SHRINK load gets expanded to scalars before render)."""
+      ir = UOp.range(cdiv(OPS_PER_TILE, NUM_THREADS), slot+1, AxisType.LOOP)
+      chunk = ir*NUM_THREADS + tid
+      r, cc = chunk // OPR, chunk % OPR
+      if getenv("HK_G2L", 0) == 3:
+        # direct-to-LDS fill (kittens '... offen lds' path): the hardware writes each lane's
+        # chunk to the lane-linear LDS address (M0 + lane*size), so the swizzle is moved to
+        # the GLOBAL side: lane q's 16B chunk fetches the matrix element that st_half_base
+        # maps to the tile-linear position q. Verified bijective; the fragment-read layout
+        # (and therefore the read swizzle) is unchanged.
+        chunk = ir*NUM_THREADS + tid
+        p_ = chunk * CPV                             # tile-linear halves position of this lane's chunk
+        sub = p_ >> 9                                # 16x32 subtile id (512 halves)
+        r16 = (p_ & 511) >> 5
+        flip = (r16 >> 3) & 1
+        cb = ((p_ & 31) >> 2) ^ (flip << 2)
+        r_ = (sub >> 1) * 16 + r16
+        c_ = cb*4 + (sub & 1) * 32                   # global column (8-aligned)
+        off_g = (base_row + r_) * K + kt*K_STEP + c_
+        lds_el = slot_off + ir*NUM_THREADS*CPV       # elements; &buf[el*8] = chunk base byte addr
+        # feed the raw PARAM (unwrapping the scheduler's RESHAPE view, which would otherwise
+        # live unfused into the program and fail spec: 'movement ops not allowed in programs').
+        prm = src
+        while prm.op is not Ops.PARAM and len(prm.src): prm = prm.src[0]
+        nbytes = prm.max_numel() * prm.dtype.itemsize
+        gname = f"data{prm.arg.slot}_{prm.max_numel()}"
+        return UOp(Ops.CUSTOMI, dtypes.void, src=(prm, dst, lds_el, off_g),
+                   arg=(f"llvm_amdgcn_raw_buffer_load_lds(make_srsrc_((void*){gname}, {nbytes}), "
+                        f"(as3_uint32_ptr)(&({{1}}[({{2}})])), {CPV*2}, ((unsigned)({{3}}))*2U, 0, 0, 0);")).end(ir)
+      # default: elementwise global->LDS stores
+      off_l = slot_off + st_half_base(r, perm_col(cc*CPV), K_STEP)
+      off_g = (base_row + r) * K + kt*K_STEP + cc*CPV
+      j = UOp.range(CPV, slot, AxisType.UPCAST)
+      return dst[off_l + j].store(src[base_row + r, kt*K_STEP + cc*CPV + j]).end(ir, j)
+
+    ZERO = UOp.const(dtypes.weakint, 0)
+    # prologue: tile 0 into slot 0 of both buffers, barrier before first read
+    g0 = UOp.group(copy_stage(A_l, ZERO, A, by*BLOCK_M, ZERO, 100),
+                   copy_stage(B_l, ZERO, B, bx*BLOCK_N, ZERO, 110))
+    bar0 = UOp.barrier(g0)
+    # Double-buffered pipeline: slot ko%2 holds k-tile ko; the prefetch copy of tile ko+1
+    # (into the other slot) rides IN FRONT of the wmma's and overlaps them; one barrier per
+    # k-tile hand-off covers write(ko)->read(ko+1) [and read(ko)->write(ko+1) is closed by
+    # the ko-1 barrier already]. The parities/offsets are static python constants when
+    # HK_UNROLL (default on): straight-line like the kittens main loop; the rolled variant
+    # uses pm_split_ranges to split the ko LOOP range at the (ko % 2) boundary.
+    if getenv("HK_UNROLL", 1) and amt % (UN := getenv("HK_UNROLL_U", 8)) == 0:
+      # outer rolled loop of amt//U iterations, U python-unrolled k-tiles inside: nearly the
+      # kittens straight-line node shape (one barrier per k-tile) at a fraction of the
+      # full-unroll uop count (full unroll of amt=64 needs ~12 min of schedule time; U=8
+      # keeps every tile's prefetch + compute + hand-off barrier but stays seconds).
+      ko_o = UOp.range(amt // UN, 600, AxisType.LOOP)
+      pa, pb = A_l.after(bar0, ko_o), B_l.after(bar0, ko_o)
+      for i in range(UN):
+        kt = ko_o * UN + i
+        pr, pn = (i % 2) * TILE_ELEMS, ((i + 1) % 2) * TILE_ELEMS
+        kt_next = UOp.minimum(kt + 1, amt - 1)
+        ga0 = UOp.group(copy_stage(pa, UOp.const(dtypes.weakint, pn), A, by*BLOCK_M, kt_next, 300 + 4*i),
+                        copy_stage(pb, UOp.const(dtypes.weakint, pn), B, bx*BLOCK_N, kt_next, 302 + 4*i))
+        sp_hi = setprio(1, 300 + 4*i)                      # kittens: raised prio for the mma phase
+        last = compute(acc, pa, pb, afters=(ko_o, sp_hi), aoff=UOp.const(dtypes.weakint, pr), boff=UOp.const(dtypes.weakint, pr))
+        sp_lo = setprio(0, 301 + 4*i)
+        handoff = UOp.group(last, sp_lo, ga0).barrier()
+        acc = acc.after(handoff)
+        pa, pb = A_l.after(handoff, ga0), B_l.after(handoff, ga0)
+      acc = acc.after(UOp.group(handoff).end(ko_o))
+    else:
+      ko = UOp.range(amt, 600, AxisType.LOOP)
+      pr, pn = ko % 2, (ko+1) % 2                    # slot of the tile being computed / being prefetched
+      kt_next = UOp.minimum(ko+1, amt-1)             # clamped tail prefetch (its data is unused)
+      pa, pb = A_l.after(bar0, ko), B_l.after(bar0, ko)
+      ga = UOp.group(copy_stage(pa, pn*TILE_ELEMS, A, by*BLOCK_M, kt_next, 130),
+                     copy_stage(pb, pn*TILE_ELEMS, B, bx*BLOCK_N, kt_next, 140))
+      sp_hi = setprio(1, 150)
+      last = compute(acc, pa, pb, afters=(ko, sp_hi), aoff=pr*TILE_ELEMS, boff=pr*TILE_ELEMS)
+      acc = acc.after(UOp.group(last, setprio(0, 151), ga).barrier().end(ko))
+
+  # ---- epilogue: per-thread fragment stores, cast to bf16 (scalar per fragment element) ----
+  mt, nt = UOp.range(MT, 801, AxisType.LOOP), UOp.range(NT, 802, AxisType.LOOP)
+  def store_i(i:int) -> UOp:
+    crow = by*BLOCK_M + warp_row*WARP_TILE_M + mt*16 + acc_m(lane, i)
+    ccol = bx*BLOCK_N + warp_col*WARP_TILE_N + nt*16 + lane % 16
+    return C[crow, ccol].store(acc[mt, nt, i].cast(dtypes.bfloat16))
+  out_st = UOp.group(*[store_i(i) for i in range(FRAG_OUT)])
+  return out_st.end(mt, nt).sink(arg=KernelInfo(name="hk_bf16_gemm",
+    estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K+M*N)*2)))
+
+def hk_bf16_gemm_tiny(a:Tensor, b:Tensor, stages:int=1) -> Tensor:
+  """C = a @ b.T for bf16 a (M,K), b (N,K) with the HipKittens-shaped tinygrad kernel."""
+  arch = Device[a.device].renderer.target.arch
+  c = Tensor.empty(a.shape[0], b.shape[0], dtype=dtypes.bfloat16, device=a.device)
+  return c.custom_kernel(a, b, fxn=lambda C, A, B: hk_bf16_gemm_kernel(C, A, B, arch=arch, stages=stages))[0]
+
+if __name__ == "__main__":
+  import numpy as np
+  from tinygrad import Device
+  M = N = K = 512
+  # exact test: B = identity -> C must equal A bit-exactly
+  a = Tensor.randn(M, K, dtype=dtypes.bfloat16).contiguous()
+  bid = Tensor(np.eye(K, N, dtype=np.float32), dtype=dtypes.bfloat16).contiguous()
+  cid = hk_bf16_gemm_tiny(a, bid, stages=getenv("STAGES", 1)).realize()
+  assert np.array_equal(cid.float().numpy(), a.float().numpy()), "identity test failed"
+  # real test: bf16 gemm vs fp32 reference, rounding-level noise
+  b = Tensor.randn(N, K, dtype=dtypes.bfloat16).contiguous()
+  c = hk_bf16_gemm_tiny(a, b, stages=getenv("STAGES", 1)).realize()
+  ref = (a @ b.T).float().realize()
+  err = (c.float() - ref).abs().max().item()
+  print(f"identity exact, random max err: {err:.5f}")
+
+  # ---- benchmark mode: kittens hk_bf16_gemm vs tinygrad stages={1,2} vs the default scheduled gemm ----
+  # run on real hardware with:  DEV=AMD:HIP:gfx950 DEBUG=2 HK_BENCH=1 python extra/gemm/hk_gemm_frag.py
+  # sizes via HK_SIZES="2048x2048x2048,4096x4096x4096" (default 2048 cubed), iteration count via ITERS=20.
+  # timings come from GlobalCounters.time_sum_s (sum of kernel times; same source as the DEBUG=2 'tm' column).
+  if getenv("HK_BENCH"):
+    from tinygrad.helpers import GlobalCounters
+    from extra.gemm.cdna_asm_gemm import asm_gemm
+    from tinygrad import Device
+    dev, iters, warm = Device.DEFAULT, getenv("ITERS", 20), 3
+    arch = Device[dev].renderer.target.arch
+    assert arch.startswith("gfx9"), "CDNA only"
+    def bench(label:str, fn, M:int, N:int, K:int) -> float:
+      try:
+        for _ in range(warm): fn()
+        Device[dev].synchronize()
+        GlobalCounters.reset()
+        import time
+        t0 = time.perf_counter()
+        for _ in range(iters): fn()
+        Device[dev].synchronize()
+        wall = time.perf_counter() - t0
+      except Exception as e:
+        print(f"    {label:32s}  unsupported/failed: {type(e).__name__}: {e}")
+        return float('nan')
+      # prefer kernel-side time (GlobalCounters matches the DEBUG=2 'tm' column); fall back to wall clock
+      ms = (GlobalCounters.time_sum_s if GlobalCounters.time_sum_s > 0 else wall) * 1e3 / iters
+      tf = 2*M*N*K / (ms * 1e-3) / 1e12
+      print(f"    {label:32s} {ms:9.3f} ms  {tf:8.1f} TFLOPS")
+      return tf
+    for (M, N, K) in [tuple(map(int, s.split("x"))) for s in getenv("HK_SIZES", "2048x2048x2048").split(",")]:
+      print(f"  size ({M},{N},{K}), grid {M//BLOCK_M}x{N//BLOCK_N} WGs, amt={K//K_STEP} k-tiles/WG")
+      np.random.seed(0)
+      An, Bn = np.random.randn(M, K), np.random.randn(K, N)
+      A = Tensor(An, dtype=dtypes.bfloat16).contiguous().realize()          # (M,K)
+      Bk = Tensor(Bn, dtype=dtypes.bfloat16).contiguous().realize()          # (K,N) for kittens
+      Bt = Bk.T.contiguous().realize()                                       # (N,K) for ours
+      tf_kc = bench("kittens hk_bf16_gemm (asm_gemm)", lambda: asm_gemm(A, Bk).realize(), M, N, K)
+      tf_s1 = bench("tiny stages=1", lambda: hk_bf16_gemm_tiny(A, Bt, stages=1).realize(), M, N, K)
+      tf_s2 = bench("tiny stages=2", lambda: hk_bf16_gemm_tiny(A, Bt, stages=2).realize(), M, N, K)
+      tf_df = bench("tinygrad default (a @ Bt.T)", lambda: (A @ Bt.T).realize(), M, N, K)
+      err = (hk_bf16_gemm_tiny(A, Bt, stages=2).float() - asm_gemm(A, Bk).float()).abs().max().item()
+      print(f"    correctness tiny-s2 vs kittens max diff: {err:.5f}")
+      for nm, tf in [("s1", tf_s1), ("s2", tf_s2), ("default", tf_df)]:
+        if tf == tf and tf_kc == tf_kc: print(f"    tiny {nm:8s}/kittens: {tf/tf_kc:6.2%}")
+  # match the real HipKittens hk_bf16_gemm on the (mock) CDNA4 emulator at small sizes.
+  # run from the repo root with:  DEV=MOCK+AMD:HIP:gfx950 HK_COMPARE=1 python extra/gemm/hk_gemm_frag.py
+  if getenv("HK_COMPARE"):
+    from extra.gemm.cdna_asm_gemm import asm_gemm
+    assert Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"), "needs CDNA4 (mock emulator or hardware)"
+    def compare(M:int, N:int, K:int, seed:int=0, identity:bool=False):
+      np.random.seed(seed)
+      An = np.random.randn(M, K)
+      Bn = np.eye(K, N) if identity else np.random.randn(K, N)      # (K,N) as expected by asm_gemm
+      A = Tensor(An, dtype=dtypes.bfloat16).contiguous()
+      B = Tensor(Bn, dtype=dtypes.bfloat16).contiguous()             # (K,N) for asm_gemm
+      c_hkc = asm_gemm(A, B).realize().float().numpy()               # real HipKittens hk_bf16_gemm
+      c_hkt = hk_bf16_gemm_tiny(A, B.T.contiguous(), stages=getenv("STAGES", 2)).realize().float().numpy()
+      ref64 = A.float().numpy().astype(np.float64) @ B.float().numpy()
+      tag = "ident" if identity else "rand "
+      print(f"({M},{N},{K}) {tag}: tiny-vs-kittens {np.abs(c_hkt-c_hkc).max():9.6f}  "
+            f"tiny-vs-fp64 {np.abs(c_hkt-ref64).max():9.6f}  kittens-vs-fp64 {np.abs(c_hkc-ref64).max():9.6f}")
+      assert np.abs(c_hkt - ref64).max() < 0.26, "tiny kernel must match fp64 at rounding level"
+      assert np.abs(c_hkt - c_hkc).max() < 0.51, "tiny kernel must match hipkittens"
+    # NOTE: hk_bf16_gemm requires K % 128 == 0 (its prologue+epilogue unconditionally touch
+    # k-tiles num_tiles-1 and num_tiles-2); at other K it reads wrong-but-in-bounds global
+    # memory on the emulator and on real hardware, so only K%128==0 sizes are checked here.
+    compare(256, 256, 128, seed=1)   # single workgroup
+    compare(256, 256, 256, seed=2)
+    compare(512, 512, 128, seed=3)   # multi workgroup
+    compare(256, 256, 128, identity=True)   # bit-exact check

extra/gemm/mi350x_uop_matmul.py

@@ -79,7 +79,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   # this is the big accumulator
   acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG)
   assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const((0.0,)*4, dtypes.float), end=init_l)
+  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float, (0.0,)*4), end=init_l)
 
   # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
   def make_locals(slot) -> tuple[UOp, UOp]:

extra/gemm/moe_routing.py

@@ -34,7 +34,7 @@ def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp:
 
 def _ggather_zero_kernel(out:UOp) -> UOp:
   i = UOp.range(out.numel(), 0)
-  return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
+  return out.flatten().index(i).store(UOp.const(out.dtype, 0.0)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
 
 def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor:
   return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0]

extra/gemm_fragment.py

@@ -1,25 +1,55 @@
 """
 tilelang-style matmul_relu written with tinygrad UOp APIs.
 
+Demonstrates that tilelang's T.alloc_fragment is expressible with existing
+tinygrad primitives: a per-thread REG buffer, wrapped in one Ops.UNSHARD per
+sharded axis over the LOCAL thread-grid ranges to form the full logical tile.
+Here the 64 threads are an 8x8 grid and each thread owns an 8x8 sub-tile --
+the 2-D fragment layout tilelang infers. The kernel is written against the
+full-tile UNSHARD view, and multi_pm (the same pass that lowers multi-device
+UNSHARDs) resolves it into per-thread shard code.
+
 Reference tilelang kernel:
-  @tilelang.jit
-  def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
-                  dtype=T.float16, accum_dtype=T.float32):
-    M, N, K = T.const('M, N, K')
-    C = T.empty([M, N], dtype)
-    with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
-      A_shared = T.alloc_shared((block_M, block_K), dtype)
-      B_shared = T.alloc_shared((block_K, block_N), dtype)
-      C_local  = T.alloc_fragment((block_M, block_N), accum_dtype)
-      T.clear(C_local)
-      for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
-        T.copy(A[by * block_M, ko * block_K], A_shared)
-        T.copy(B[ko * block_K, bx * block_N], B_shared)
-        T.gemm(A_shared, B_shared, C_local)
-      for i, j in T.Parallel(block_M, block_N):
-        C_local[i, j] = T.max(C_local[i, j], 0)
-      T.copy(C_local, C[by * block_M, bx * block_N])
-    return C
+
+    @tilelang.jit
+    def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
+                    dtype=T.float16, accum_dtype=T.float32):
+        M, N, K = T.const('M, N, K')
+        C = T.empty([M, N], dtype)
+        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+            A_shared = T.alloc_shared((block_M, block_K), dtype)
+            B_shared = T.alloc_shared((block_K, block_N), dtype)
+            C_local  = T.alloc_fragment((block_M, block_N), accum_dtype)
+            T.clear(C_local)
+            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
+                T.copy(A[by * block_M, ko * block_K], A_shared)
+                T.copy(B[ko * block_K, bx * block_N], B_shared)
+                T.gemm(A_shared, B_shared, C_local)
+            for i, j in T.Parallel(block_M, block_N):
+                C_local[i, j] = T.max(C_local[i, j], 0)
+            T.copy(C_local, C[by * block_M, bx * block_N])
+        return C
+
+API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_kernel.py):
+
+    T.Kernel(gx, gy, threads=T)    -> AxisType.GLOBAL ranges (blocks) + AxisType.LOCAL ranges (thread grid)
+    T.alloc_shared(shape, dtype)   -> UOp.placeholder(shape, dtype, slot, AddrSpace.LOCAL)
+    T.alloc_fragment(shape, dt)    -> per-thread REG placeholder, wrapped in one Ops.UNSHARD per sharded axis over
+                                      the AxisType.LOCAL ranges: fragment.unshard((axis_y, axis_x), (ty, tx)).
+                                      The full logical tile is the shard with each sharded axis multiplied by its
+                                      range size, exactly like device sharding, but the sharding axes are thread
+                                      axes carried by the RANGE metadata instead of a device tuple. C_local[i, j]
+                                      with [i, j] in this thread's shard is INDEX on the UNSHARD, which multi_pm
+                                      resolves into INDEX on the per-thread REG shard, axis by axis.
+    T.copy(gmem_slice, smem)       -> smem[thread_idx].set(gmem_slice[thread_idx], end=copy_rng). set returns the
+                                      smem tile AFTER the copy; the implicit-barrier pass turns the store->load
+                                      dependency of the loop that consumes it into a workgroup barrier
+    T.gemm (no WMMA)               -> C_local[..].set(C_local.after(k)[..] + a_shared[..] * b_shared[..], end=k)
+                                      with k a loop-carried LOOP range (codegen builds the register accumulator
+                                      from this self-referential store automatically)
+    T.copy(fragment, gmem)         -> gmem.index(gidx).store(C_local[..]).end(all_ranges)
+    UNSHARD lowering               -> multi_pm in codegen (full_rewrite_to_sink): INDEX/AFTER/STORE ops on the
+                                      full-tile view become per-thread shard ops, no UNSHARD survives into the program.
 """
 
 from tinygrad.dtype import dtypes, AddrSpace, DType
@@ -31,17 +61,25 @@ from tinygrad.tensor import Tensor
 # tilelang builtins, expressed with tinygrad UOp APIs
 # ---------------------------------------------------------------------------
 
-def alloc_shared(shape:tuple[int, ...], dtype:DType, slot:int) -> UOp:
+def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp:
   """T.alloc_shared: one LOCAL buffer shared by all threads in the block."""
-  return UOp.placeholder(tuple(shape), dtype, slot, AddrSpace.LOCAL)
+  return UOp.placeholder(tuple(shape), dtype, next(UOp.unique_num), AddrSpace.LOCAL)
 
-def alloc_fragment(shape:tuple[int, ...], dtype:DType, slot:int, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
-  """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid."""
+def alloc_fragment(shape:tuple[int, ...], dtype:DType, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
+  """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid.
+
+  Each thread privately owns shape[axis]//threads elements along every sharded
+  axis in a REG buffer. The UNSHARDs over the LOCAL thread ranges present the
+  full logical tile: full_shape = shard_shape with each sharded axis multiplied
+  by its range size. This is exactly how UNSHARD carries a DEVICE axis today,
+  except the sharding axes are thread axes carried by the RANGE metadata.
+  """
   assert len(axes) == len(rngs)
   assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges"
+  assert all(shape[a] % (int(rng.vmax)+1) == 0 for a, rng in zip(axes, rngs))
   by_axis = dict(zip(axes, rngs))
   shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape))
-  fragment = UOp.placeholder(shard_shape, dtype, slot, AddrSpace.REG)
+  fragment = UOp.placeholder(shard_shape, dtype, next(UOp.unique_num), AddrSpace.REG)
   return fragment.unshard(axes, rngs)
 
 # ---------------------------------------------------------------------------
@@ -68,45 +106,64 @@ def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
   # with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) as (bx, by):
   bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL)
   by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL)
-
-  # 16*8 threads = 128 threads
+  # tx (N, 16) is the fast/inner LOCAL axis so a warp covers 16 cols x 2 rows --
+  # matching tilelang's (tidx>>4, tidx&15) warp composition. This keeps the 8 A_shared
+  # reads in a warp on only 2 row-groups (broadcast across 16 cols) instead of 8 rows
+  # (8-way bank conflict), since A_shared[row*512 + ...] all map to the same bank when 8
+  # distinct rows land in one warp.
   tx = UOp.range(TX, 2, AxisType.LOCAL)
   ty = UOp.range(TY, 3, AxisType.LOCAL)
 
-  # shared + fragment (regs)
-  A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype, 0)
-  B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype, 1)
-  C_local = alloc_fragment((TM, TY, TX, TN), dtypes.float32, 0, (1, 2), (ty, tx))
+  # A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), dtype)
+  # B_shared = T.alloc_shared((BLOCK_K, BLOCK_N), dtype)
+  A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype)
+  B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype)
 
-  # zero out the regs to start. this is expanded by the devectorizer
-  C_local = C_local.after(C_local.store(0.0))
+  # C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype) -- an 8x4 REG tile per thread of the 8x16 grid
+  C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, (0, 1), (ty, tx))
+
+  # T.clear(C_local) -- each thread zeroes its own fragment sub-tile
+  ic, jc = UOp.range(TM, 4, AxisType.LOOP), UOp.range(TN, 5, AxisType.UPCAST)
+  C_loc = C_local[ic*TM + ty, tx*TN + jc].set(0.0, end=(ic, jc))
 
   # for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
+  # (num_stages pipelining is async copy + multi-buffering; this is the synchronous single-buffer version)
   ko = UOp.range(cdiv(K, BLOCK_K), 6, AxisType.LOOP)
 
-  # index the outer matrices
-  a = a.rearrange("(m bm) (k bk) -> m k bm bk", bm=BLOCK_M, bk=BLOCK_K)[by, ko]
-  b = b.rearrange("(k bk) (n bn) -> k n bk bn", bk=BLOCK_K, bn=BLOCK_N)[ko, bx]
-  c = c.rearrange("(m bm) (n bn) -> m n bm bn", bm=BLOCK_M, bn=BLOCK_N)[by, bx]
+  # T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies its own 8x4 sub-tile.
+  # Row index is iar*TM + ty (strided by TM across ty), matching tilelang's layout: thread ty owns
+  # rows {ty, ty+8, ..., ty+56} not {ty*8, ..., ty*8+7}.
+  iar, ka = UOp.range(TM, 7, AxisType.LOOP), UOp.range(TN, 8, AxisType.UPCAST)
+  A_store = A_shared[iar*TM + ty, tx*TN + ka].store(a[by*BLOCK_M + iar*TM + ty, ko*BLOCK_K + tx*TN + ka]).end(iar, ka)
 
-  # T.copy: A_shared <- a, B_shared <- b
-  def with_threads(x:UOp): return x.rearrange("(tm ty) (tx tn) -> ty tx tm tn", tm=TM, tn=TN)[ty, tx]
-  A_shared = A_shared.after(with_threads(A_shared).store(with_threads(a)))
-  B_shared = B_shared.after(with_threads(B_shared).store(with_threads(b)))
+  # T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared)
+  kb, ibr = UOp.range(TM, 9, AxisType.LOOP), UOp.range(TN, 10, AxisType.UPCAST)
+  B_store = B_shared[kb*TM + ty, tx*TN + ibr].store(b[ko*BLOCK_K + kb*TM + ty, bx*BLOCK_N + tx*TN + ibr]).end(kb, ibr)
 
-  # T.gemm(A_shared, B_shared, C_local), no WMMA
-  kk = UOp.range(BLOCK_K, 11, AxisType.LOOP)
-  ir = UOp.range(TM, 12, AxisType.LOOP)
+  # get the shared after the stores (single barrier)
+  A_shared = A_shared.after(A_store, B_store)
+  B_shared = B_shared.after(A_store, B_store)
+
+  # T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment sub-tile.
+  # identical to custom_gemm: a self-referential store over the loop-carried kk range,
+  # which codegen turns into a register accumulator
+  # kk is the outer compute loop (axis 11) so that for each kk we read all 8 A rows and reuse
+  # the B[kk] read across them -- matching tilelang's ko > kk > row > col access order exactly.
+  kk, ir = UOp.range(BLOCK_K, 11, AxisType.LOOP), UOp.range(TM, 12, AxisType.LOOP)
   jj = UOp.range(TN, 13, AxisType.UPCAST)
-  acc = C_local.after(kk)[ir, ty, tx, jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32)
+  acc = C_loc.after(kk)[ir*TM + ty, tx*TN + jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32)
   # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
-  C_local = C_local[ir, ty, tx, jj].set(acc, end=(kk, ir, jj, ko))
+  C_loc = C_loc[ir*TM + ty, tx*TN + jj].set(acc, end=(kk, ir, jj, ko))
 
-  # c <- C_local (with relu and cast): every thread stores its shard's sub-view of the output tile
-  c_st = c.reshape(C_local.shape).store(C_local.relu().cast(c.dtype))
+  # for i, j in T.Parallel(BLOCK_M, BLOCK_N): C_local[i, j] = T.max(C_local[i, j], 0)
+  # T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N]) -- per-thread store of the fragment shard (relu fused into it)
+  # LOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them
+  ie, je = UOp.range(TM, 14, AxisType.LOOP), UOp.range(TN, 15, AxisType.UPCAST)
+  c_st = c[by*BLOCK_M + ie*TM + ty, bx*BLOCK_N + tx*TN + je].store(C_loc[ie*TM + ty, tx*TN + je].relu().cast(c.dtype))
 
-  # close the locals and globals
-  return c_st.end(tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
+  # all open ranges are closed at the final store (ko was closed above).
+  # the fragment UNSHARDs go to codegen as is: multi_pm there resolves the full-tile view into per-thread shard code
+  return c_st.end(je, ie, tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
 
 # ---------------------------------------------------------------------------
 # python wrapper: same signature as the tilelang function
@@ -131,8 +188,7 @@ if __name__ == "__main__":
   b = Tensor.randn(K, N, dtype=dtype_in).contiguous()
   ref = (a @ b).relu().realize()
 
-  for _ in range(10):
-    out = matmul_relu(a, b).realize()
+  out = matmul_relu(a, b).realize()
 
   import numpy as np
   np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2)

extra/hcq2/ops_amd2.py

@@ -37,7 +37,7 @@ class PM4Ops(FastEnum):
   RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702
 
 def pkt3(ctx, op:PM4Ops, *vals):
-  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
+  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(dtypes.uint32, x)
     for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))
 
 def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
@@ -157,19 +157,18 @@ def pm4_submit(ctx, lin):
   assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
 
   ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
-  done_idx, submit_idx = UOp.const(size_dw + 0, dtypes.int), UOp.const(size_dw + 1, dtypes.int)
-  init_counters = make_patches(ib, [((size_dw + i) * 4, UOp.const(0, dtypes.uint32)) for i in range(2)]).rtag("link")
-  submitted = (counter:=ib.after(init_counters).index(submit_idx)).load()
+  done_idx, submit_idx = UOp.const(dtypes.int, size_dw + 0), UOp.const(dtypes.int, size_dw + 1)
+  submitted = (counter:=ib.after(make_patches(ib, [((size_dw + i) * 4, UOp.const(dtypes.uint32, 0)) for i in range(2)])).index(submit_idx)).load()
   completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load()
   ib_free = completed.end(loop, completed != submitted)
 
-  bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(size_dw)), 2), (submitted + 1).cast(dtypes.uint64))
+  bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(dtypes.weakint, size_dw)), 2), (submitted + 1).cast(dtypes.uint64))
   cmdbuf = make_cmdbuf(lin.replace(src=lin.src + (bump_fence,)), devs, buf=ib, dep=ib_free)
 
   # the ring itself only carries a packet pointing at the ib, wrapping the ring
-  put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
+  put = put_ptr.index(zero:=UOp.const(dtypes.int, 0))
   pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
-  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(x, dtypes.uint32)) for off,x in enumerate(pkt)])
+  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(dtypes.uint32, x)) for off,x in enumerate(pkt)])
 
   # advance the put/write pointers past the packet
   bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
@@ -187,26 +186,26 @@ class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TR
 def sdma_copy(ctx, call):
   sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
   src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
-  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
+  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(dtypes.uint32, x) for off in range(0, sz, ctx.max_copy_size) for x in (
     ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
     ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, *data64_le(src_addr+off), *data64_le(dst_addr+off))))
 
 def sdma_wait(ctx, ins, dst, val):
   op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
      | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) for x in (
+  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(dtypes.uint32, x) for x in (
     op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
     ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))
 
 def sdma_store(ctx, ins, dst, val):
   op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
   return UOp(Ops.LINEAR, src=(
-    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
-    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
+    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
+    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(dtypes.uint32, x) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
 
 def sdma_timestamp(ctx, ins, dst):
   op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
+  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
 
 pm_sdma_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
@@ -219,7 +218,7 @@ pm_sdma_opsel = PatternMatcher([
 
 def sdma_submit(cmdbuf, devs):
   # the cmdbuf to submit + the patch writes that fill it
-  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int)
+  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0)
 
   # the sdma queue's ring and its host-side ring/write/put pointers
   for d in devs: q = Device[d].sdma_queue(0)
@@ -235,8 +234,8 @@ def sdma_submit(cmdbuf, devs):
 
   # zero the wrapped tail, then copy the cmdbuf into the ring
   zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
-  zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
-  i = UOp.range(UOp.const(size_dw, dtypes.int), 0, dtype=dtypes.int, src=(cmdbuf,))
+  zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(dtypes.uint32, 0)).end(zi)
+  i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,))
   copy_to_ring = ring.index(start_dw + i).store(cmdbuf.index(i).load()).end(i)
 
   # advance the put/write pointers past the zeroed tail and the cmdbuf

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -48,7 +48,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
   device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]
   if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
   else: raise NotImplementedError(f"no atomic max for device {device}")
-  amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
+  amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.weakint, 0))
   max_val = lds[0].load()
   atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
   return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))

test/amd/test_emu2_pcode.py

@@ -11,8 +11,8 @@ from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp
 
 def _srcs():
   """Create minimal source variables for pcode parsing."""
-  def u32(v=0): return UOp.const(v, dtypes.uint32)
-  return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(0, dtypes.uint64), 'laneId': u32()}
+  def u32(v=0): return UOp.const(dtypes.uint32, v)
+  return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(dtypes.uint64, 0), 'laneId': u32()}
 
 class TestBasicParsing(unittest.TestCase):
   """Test basic pcode parsing for common instruction patterns."""
@@ -44,8 +44,8 @@ class TestWithSources(unittest.TestCase):
 
   def test_v_add_f32_with_sources(self):
     """Test V_ADD_F32 with actual float constants."""
-    s0 = UOp.const(0x3f800000, dtypes.uint32)  # 1.0f
-    s1 = UOp.const(0x40000000, dtypes.uint32)  # 2.0f
+    s0 = UOp.const(dtypes.uint32, 0x3f800000)  # 1.0f
+    s1 = UOp.const(dtypes.uint32, 0x40000000)  # 2.0f
     _, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], {'S0': s0, 'S1': s1})
     self.assertEqual(len(assigns), 1)
     dest, val = assigns[0]
@@ -55,8 +55,8 @@ class TestWithSources(unittest.TestCase):
 
   def test_v_mul_f32_with_sources(self):
     """Test V_MUL_F32 with actual float constants."""
-    s0 = UOp.const(0x40000000, dtypes.uint32)  # 2.0f
-    s1 = UOp.const(0x40400000, dtypes.uint32)  # 3.0f
+    s0 = UOp.const(dtypes.uint32, 0x40000000)  # 2.0f
+    s1 = UOp.const(dtypes.uint32, 0x40400000)  # 3.0f
     _, assigns = parse_pcode(PCODE[VOP2Op.V_MUL_F32_E32], {'S0': s0, 'S1': s1})
     self.assertEqual(len(assigns), 1)
     dest, val = assigns[0]
@@ -67,36 +67,36 @@ class TestParseExpr(unittest.TestCase):
 
   def test_integer_literals(self):
     """Test parsing integer literals."""
-    self.assertEqual(parse_expr('0', {}).val, 0)
-    self.assertEqual(parse_expr('42', {}).val, 42)
-    self.assertEqual(parse_expr('42U', {}).val, 42)
+    self.assertEqual(parse_expr('0', {}).arg, 0)
+    self.assertEqual(parse_expr('42', {}).arg, 42)
+    self.assertEqual(parse_expr('42U', {}).arg, 42)
 
   def test_negative_integers(self):
     """Test parsing negative integer literals."""
     result = parse_expr('-1', {})
-    self.assertEqual(result.val, -1)
+    self.assertEqual(result.arg, -1)
     self.assertEqual(result.dtype, dtypes.int)
 
   def test_float_literals(self):
     """Test parsing float literals."""
     result = parse_expr('1.0F', {})
-    self.assertEqual(result.val, 1.0)
+    self.assertEqual(result.arg, 1.0)
     self.assertEqual(result.dtype, dtypes.float32)
 
   def test_hex_literals(self):
     """Test parsing hex literals."""
     result = parse_expr('0xFF', {})
-    self.assertEqual(result.val, 255)
+    self.assertEqual(result.arg, 255)
 
   def test_variable_lookup(self):
     """Test variable lookup in parse_expr."""
-    vrs = {'x': UOp.const(42, dtypes.uint32)}
+    vrs = {'x': UOp.const(dtypes.uint32, 42)}
     result = parse_expr('x', vrs)
-    self.assertEqual(result.val, 42)
+    self.assertEqual(result.arg, 42)
 
   def test_binary_ops(self):
     """Test parsing binary operations."""
-    vrs = {'a': UOp.const(10, dtypes.uint32), 'b': UOp.const(5, dtypes.uint32)}
+    vrs = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)}
 
     # Addition
     result = parse_expr('a + b', vrs)
@@ -105,11 +105,11 @@ class TestParseExpr(unittest.TestCase):
     # Subtraction with constant folding
     result = parse_expr('10 - 5', {})
     self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.val, 5)
+    self.assertEqual(result.arg, 5)
 
   def test_ternary(self):
     """Test parsing ternary expressions."""
-    vrs = {'cond': UOp.const(True), 'a': UOp.const(1, dtypes.uint32), 'b': UOp.const(0, dtypes.uint32)}
+    vrs = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)}
     result = parse_expr('cond ? a : b', vrs)
     self.assertEqual(result.op, Ops.WHERE)
 
@@ -127,7 +127,7 @@ class TestForLoopParsing(unittest.TestCase):
   def test_clz_parsing(self):
     """Test CLZ pcode parsing produces correct structure."""
     pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
-    S0 = UOp.const(0xFFFFFFFF, dtypes.uint32)  # All ones - CLZ should be 0
+    S0 = UOp.const(dtypes.uint32, 0xFFFFFFFF)  # All ones - CLZ should be 0
     _vrs, assigns = parse_pcode(pcode, {'S0': S0})
 
     self.assertEqual(len(assigns), 1)
@@ -139,7 +139,7 @@ class TestForLoopParsing(unittest.TestCase):
   def test_clz_with_zero(self):
     """Test CLZ with input 0 - should return -1."""
     pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
-    S0 = UOp.const(0, dtypes.uint32)
+    S0 = UOp.const(dtypes.uint32, 0)
     _vrs, assigns = parse_pcode(pcode, {'S0': S0})
 
     # Check that the innermost value (default) is -1 (may be wrapped in CAST)
@@ -150,7 +150,7 @@ class TestForLoopParsing(unittest.TestCase):
     # Unwrap CAST if present
     while val.op == Ops.CAST:
       val = val.src[0]
-    self.assertEqual(val.val, -1)
+    self.assertEqual(val.arg, -1)
 
   def test_ctz_parsing(self):
     """Test CTZ pcode parsing."""
@@ -158,7 +158,7 @@ class TestForLoopParsing(unittest.TestCase):
     if pcode is None:
       self.skipTest("V_CTZ_I32_B32_E32 pcode not available")
 
-    S0 = UOp.const(1, dtypes.uint32)  # LSB set - CTZ should be 0
+    S0 = UOp.const(dtypes.uint32, 1)  # LSB set - CTZ should be 0
     _vrs, assigns = parse_pcode(pcode, {'S0': S0})
     self.assertEqual(len(assigns), 1)
 
@@ -169,8 +169,8 @@ class TestDSPcodePatterns(unittest.TestCase):
     """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
     vmem = UOp.param(2, dtypes.uint32, (1024,))
     srcs = {
-      'ADDR': UOp.const(0, dtypes.uint64),
-      'DATA': UOp.const(0x3f800000, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint64, 0),
+      'DATA': UOp.const(dtypes.uint32, 0x3f800000),
       '_vmem': vmem,
     }
 
@@ -199,8 +199,8 @@ class TestDSPcodePatterns(unittest.TestCase):
     """Test MEM[addr].type read expression parsing."""
     # Create a mock LDS buffer
     lds = UOp.param(3, dtypes.uint32, (16384,))
-    addr = UOp.const(0, dtypes.uint32)
-    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)}
+    addr = UOp.const(dtypes.uint32, 0)
+    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}
 
     result = parse_expr('MEM[ADDR + OFFSET].b32', vrs)
     # Should be an INDEX operation into LDS
@@ -212,13 +212,13 @@ class TestDSPcodePatterns(unittest.TestCase):
     self.assertIsNotNone(pcode)
     assert pcode is not None
     srcs = {
-      'ADDR': UOp.const(0, dtypes.uint32),
-      'OFFSET0': UOp.const(0, dtypes.uint32),
-      'OFFSET1': UOp.const(1, dtypes.uint32),
-      'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32),
-      'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 0),
+      'OFFSET0': UOp.const(dtypes.uint32, 0),
+      'OFFSET1': UOp.const(dtypes.uint32, 1),
+      'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
+      'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
     }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
     _, assigns = parse_pcode(pcode, srcs)
     # Should have 2 MEM write assignments
     self.assertEqual(len(assigns), 2)
@@ -235,12 +235,12 @@ class TestDSPcodePatterns(unittest.TestCase):
     assert pcode is not None
     lds = UOp.param(3, dtypes.uint32, (16384,))
     srcs = {
-      'ADDR': UOp.const(0, dtypes.uint32),
-      'OFFSET0': UOp.const(0, dtypes.uint32),
-      'OFFSET1': UOp.const(1, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 0),
+      'OFFSET0': UOp.const(dtypes.uint32, 0),
+      'OFFSET1': UOp.const(dtypes.uint32, 1),
       '_lds': lds,
     }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
     _, assigns = parse_pcode(pcode, srcs)
     # Should have 2 RETURN_DATA assignments
     self.assertEqual(len(assigns), 2)
@@ -252,36 +252,36 @@ class TestDSPcodePatterns(unittest.TestCase):
     pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
     assert pcode is not None
     srcs = {
-      'ADDR': UOp.const(100, dtypes.uint32),
-      'OFFSET0': UOp.const(2, dtypes.uint32),
-      'OFFSET1': UOp.const(5, dtypes.uint32),
-      'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32),
-      'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 100),
+      'OFFSET0': UOp.const(dtypes.uint32, 2),
+      'OFFSET1': UOp.const(dtypes.uint32, 5),
+      'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
+      'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
     }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
     _, assigns = parse_pcode(pcode, srcs)
     # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
     # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
-    self.assertEqual(assigns[0][1][0].simplify().val, 108)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][0].simplify().val, 120)  # type: ignore[index]
+    self.assertEqual(assigns[0][1][0].simplify().arg, 108)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][0].simplify().arg, 120)  # type: ignore[index]
 
   def test_ds_store_data_values(self):
     """Test DS_STORE_2ADDR_B32 uses correct data values."""
     pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
     assert pcode is not None
     srcs = {
-      'ADDR': UOp.const(0, dtypes.uint32),
-      'OFFSET0': UOp.const(0, dtypes.uint32),
-      'OFFSET1': UOp.const(1, dtypes.uint32),
-      'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32),
-      'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 0),
+      'OFFSET0': UOp.const(dtypes.uint32, 0),
+      'OFFSET1': UOp.const(dtypes.uint32, 1),
+      'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
+      'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
     }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
     _, assigns = parse_pcode(pcode, srcs)
     # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
     # DATA[31:0] should preserve the value
-    self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][1].simplify().val, 0xBBBBBBBB)  # type: ignore[index]
+    self.assertEqual(assigns[0][1][1].simplify().arg, 0xAAAAAAAA)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][1].simplify().arg, 0xBBBBBBBB)  # type: ignore[index]
 
 class TestConditionalParsing(unittest.TestCase):
   """Test conditional (if/elsif/else) pcode parsing."""
@@ -290,9 +290,9 @@ class TestConditionalParsing(unittest.TestCase):
     """Test parsing ternary expression (which becomes WHERE)."""
     # S_CSELECT_B32: D0.u32 = SCC ? S0.u32 : S1.u32
     pcode = PCODE[SOP2Op.S_CSELECT_B32]
-    s0 = UOp.const(10, dtypes.uint32)
-    s1 = UOp.const(20, dtypes.uint32)
-    scc = UOp.const(1, dtypes.uint32)
+    s0 = UOp.const(dtypes.uint32, 10)
+    s1 = UOp.const(dtypes.uint32, 20)
+    scc = UOp.const(dtypes.uint32, 1)
     _vrs, assigns = parse_pcode(pcode, {'S0': s0, 'S1': s1, 'SCC': scc})
     self.assertEqual(len(assigns), 1)
     dest, val = assigns[0]
@@ -305,26 +305,26 @@ class TestConcatWidthParsing(unittest.TestCase):
 
   def test_permlanex16_altrow_concat(self):
     for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
-      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(dtypes.uint32, row)})
+      self.assertEqual(parsed.simplify().arg, expected)
 
   def test_permlane64_altlane_concat(self):
     for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
-      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(dtypes.uint32, lane)})
+      self.assertEqual(parsed.simplify().arg, expected)
 
   def test_permlane64_wave64_pcode_indices(self):
     vgpr = UOp.param(0, dtypes.uint32, (256,))
     srcs = {
-      'SRC0': UOp.const(0, dtypes.uint32),
-      'VDST': UOp.const(1, dtypes.uint32),
-      'EXEC_LO': UOp.const(0xFFFFFFFF, dtypes.uint32),
-      'EXEC': UOp.const(0xFFFFFFFFFFFFFFFF, dtypes.uint64),
+      'SRC0': UOp.const(dtypes.uint32, 0),
+      'VDST': UOp.const(dtypes.uint32, 1),
+      'EXEC_LO': UOp.const(dtypes.uint32, 0xFFFFFFFF),
+      'EXEC': UOp.const(dtypes.uint64, 0xFFFFFFFFFFFFFFFF),
       '_vgpr': vgpr,
       '_wave_size': 64,
-      'S0': UOp.const(0, dtypes.uint32),
-      'S1': UOp.const(0, dtypes.uint32),
-      'S2': UOp.const(0, dtypes.uint32),
+      'S0': UOp.const(dtypes.uint32, 0),
+      'S1': UOp.const(dtypes.uint32, 0),
+      'S2': UOp.const(dtypes.uint32, 0),
     }
 
     def load_idx(v: UOp) -> int:
@@ -333,12 +333,12 @@ class TestConcatWidthParsing(unittest.TestCase):
       self.assertEqual(simp.src[0].op, Ops.INDEX)
       idx = simp.src[0].src[1].simplify()
       self.assertEqual(idx.op, Ops.CONST)
-      return idx.val
+      return idx.arg
 
     _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
     self.assertEqual(len(assigns), 64)
     for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
-      self.assertEqual(assigns[lane][1][0].simplify().val, dst_idx)  # type: ignore[index]
+      self.assertEqual(assigns[lane][1][0].simplify().arg, dst_idx)  # type: ignore[index]
       self.assertEqual(load_idx(assigns[lane][1][1]), src_idx)  # type: ignore[index]
 
 class TestAllPcode(unittest.TestCase):
@@ -346,7 +346,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)
+    u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
     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(),
@@ -358,7 +358,7 @@ class TestAllPcode(unittest.TestCase):
             'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(),
             'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
             'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(),
-            '_lds': lds, '_vmem': lds, '_active': UOp.const(True)}
+            '_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)}
 
   def _parse_all_pcode(self, pcode_dict, arch: str, min_pct: float):
     """Parse all pcode. RuntimeError = parser limitation (ok), other exceptions = real bugs."""

test/backend/test_asm_gemm.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.helpers import getenv, system, DEV
-from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm, quantize_mxfp4
+from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
 from test.helpers import needs_second_gpu
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
 
@@ -150,45 +150,6 @@ class TestAsmGEMM(unittest.TestCase):
     with self.assertRaisesRegex(AssertionError, "not a multiple"):
       verify_asm_gemm(1, 256, 1000, 256)
 
-class TestMXFP4(unittest.TestCase):
-  def setUp(self):
-    if not is_cdna4() or DEV.interface.startswith("MOCK"):
-      self.skipTest("requires real amd machine")
-
-  def test_quantize(self):
-    import numpy as np
-    block = np.array([0, .26, .74, .75, 1.26, 1.75, 2.51, 3.5, 5.1, 6, -6] + [0] * 21, dtype=np.float32)
-    x = Tensor(np.tile(block, (32, 8)), dtype=dtypes.bfloat16)
-    packed, scale, _ = quantize_mxfp4(x)
-    p = packed.numpy()
-    codes = np.stack((p & 0xF, p >> 4), axis=-1).reshape(32, 256)
-    np.testing.assert_array_equal(codes[0, :11], [0, 1, 1, 2, 3, 4, 5, 6, 7, 7, 15])
-    np.testing.assert_array_equal(scale.numpy(), np.full((32, 8), 127, dtype=np.uint8))
-
-  def test_correctness(self):
-    import numpy as np
-    M = N = K = 256
-    rng = np.random.default_rng(1)
-    a = Tensor(rng.standard_normal((M, K), dtype=np.float32), dtype=dtypes.bfloat16)
-    b = Tensor(rng.standard_normal((N, K), dtype=np.float32), dtype=dtypes.bfloat16)
-    out = asm_gemm(a, b.T, mxfp4=True).realize()
-    # reference gemm
-    a_packed, scale_a, _ = quantize_mxfp4(a)
-    b_packed, scale_b, _ = quantize_mxfp4(b)
-    def unpack(x): return np.stack((x & 0xF, x >> 4), axis=-1).reshape(x.shape[0], -1)
-    code_a, code_b = unpack(a_packed.numpy()), unpack(b_packed.numpy())
-    lut = np.array([0, .5, 1, 1.5, 2, 3, 4, 6, -0., -.5, -1, -1.5, -2, -3, -4, -6], dtype=np.float32)
-    a_dequant = lut[code_a] * np.repeat(np.exp2(scale_a.numpy().astype(np.int16)-127), 32, axis=1)
-    b_dequant = lut[code_b] * np.repeat(np.exp2(scale_b.numpy().astype(np.int16)-127), 32, axis=1)
-    ref = Tensor(a_dequant @ b_dequant.T, dtype=dtypes.bfloat16).realize().numpy()
-    np.testing.assert_array_equal(out.numpy(), ref)
-
-  def test_empty(self):
-    M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336)
-    a = Tensor.empty(M, K, dtype=dtypes.bfloat16)
-    b = Tensor.empty(N, K, dtype=dtypes.bfloat16)
-    asm_gemm(a, b.T, mxfp4=True).realize()
-
 # test the Asm GEMM with Llama shapes, only run on the real machine for speed
 
 @unittest.skipUnless(has_hipcc(), "requires hipcc to compile")

test/backend/test_call.py

@@ -6,11 +6,11 @@ from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.uop.ops import KernelInfo
 
 def call_out_kernel(F:UOp, C:UOp) -> UOp:
-  call = F[0].load().call(UOp.const(3).cast(dtypes.int), C[0], ret_dtype=dtypes.void)
+  call = F[0].load().call(UOp.const(dtypes.int, 3), C[0], ret_dtype=dtypes.void)
   return C.after(call)[1].store(C.after(call)[0].load() + 1).sink(arg=KernelInfo(name="call_out"))
 
 def call_ret_kernel(F:UOp, C:UOp) -> UOp:
-  val = F[0].load().call(UOp.const(21).cast(dtypes.int), ret_dtype=dtypes.int)
+  val = F[0].load().call(UOp.const(dtypes.int, 21), ret_dtype=dtypes.int)
   return C[0].store(val * 2).sink(arg=KernelInfo(name="call_ret"))
 
 @unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "TODO: CALL is rendered in C style only")

test/backend/test_const_folding.py

@@ -16,7 +16,7 @@ def _check_ast_count(desired_count:int, t:Tensor):
 
 class TestMovedConstFolding(unittest.TestCase):
   def test_contiguous_deviceless_const(self):
-    t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
+    t = Tensor(UOp.const(dtypes.float, 2.0)).contiguous()
     self.assertIs(t.uop.op, Ops.CONST)
     self.assertIsNone(t.uop.device)
 
@@ -169,7 +169,7 @@ class TestMultiConstFolding(unittest.TestCase):
 class TestThreefryConstFolding(unittest.TestCase):
   def test_threefry(self):
     # THREEFRY(const,const) folds to a const once decomposed
-    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
+    x = threefry2x32(UOp.const(dtypes.uint64, 5), UOp.const(dtypes.uint64, 10))
     self.assertIs(x.simplify().op, Ops.CONST)
 
 class TestTautologicalCompare(unittest.TestCase):

test/backend/test_custom_kernel.py

@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
-import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
 from tinygrad.renderer.ptx import PTXRenderer
@@ -341,7 +340,7 @@ class TestCustomKernel(unittest.TestCase):
 
   def test_partial_invalid_store_keeps_uncovered_reads(self):
     x = Tensor([10., 20., 30., 40.])
-    after = x.uop.after(x.uop.shrink(((0, 2),)).store(Invalid))
+    after = x.uop.after(x.uop.shrink(((0, 2),)).store(UOp.const(dtypes.float, Invalid, shape=(2,))))
     self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.])
 
   def test_multi_after_invalid_store_dep_removed(self):
@@ -528,77 +527,6 @@ class TestUnshardIndex(unittest.TestCase):
     with self.assertRaisesRegex(RuntimeError, "cannot shard index"):
       self._run(kernel, (64, 8))
 
-def _run_fragment_kernel(testcase, kernel, out_shape, inputs=()):
-  c = Tensor.empty(*out_shape)
-  out = Tensor.custom_kernel(c, *inputs, fxn=kernel)[0]
-  try: return out.numpy()
-  except RuntimeError as e:
-    if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
-      testcase.skipTest("PTX does not support dynamic register indexing")
-    raise
-
-class TestUnshardAlu(unittest.TestCase):
-  """Tests for ALU on (fragment) UNSHARD values in schedule/multi.py's alu_multi.
-
-  An ALU with UNSHARD srcs lowers to per-shard ops when every src is one of:
-    same sharding:  peel the UNSHARD, keep the layout
-    scalar:         broadcast to every shard
-    whole unsharded same-shape value: takes its per-shard sub-view (shard_subview)
-  """
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_alu_scalar_broadcast(self):
-    # scalar srcs broadcast to every shard: frag*2.0 where frag is 1.5 per thread -> 3.0 everywhere
-    def kernel(C:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      # 8 values per thread, 8 threads -> 64-value full view
-      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
-      v = frag.after(frag.store(1.5)) * 2.0
-      return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_scalar", opts_to_apply=()))
-    out = _run_fragment_kernel(self, kernel, (64,))
-    np.testing.assert_allclose(out, 3.0)
-
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_alu_whole_value_subview(self):
-    # UNSHARD + whole unsharded same-shape value: each shard adds its own sub-view of A.
-    def kernel(C:UOp, A:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
-      v = frag.after(frag.store(0.0)) + A
-      return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_subview", opts_to_apply=()))
-    a = Tensor(np.arange(64, dtype=np.float32))
-    out = _run_fragment_kernel(self, kernel, (64,), inputs=(a,))
-    np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
-
-class TestUnshardStore(unittest.TestCase):
-  """Tests for STORE of a sharded value into an unsharded dest (store_value_multi in schedule/multi.py).
-
-  Every shard stores its value into its own contiguous sub-view of the dest, one SHRINK per sharded axis.
-  """
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_store_unshard_value(self):
-    # single-axis: 8 threads each own 8 values of the 64-value output tile
-    def kernel(C:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
-      v = frag.after(frag.store(0.0)) + 2.5
-      return C.store(v).end(ty).sink(arg=KernelInfo(name="store_unshard", opts_to_apply=()))
-    out = _run_fragment_kernel(self, kernel, (64,))
-    np.testing.assert_allclose(out, 2.5)
-
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_store_unshard_value_2axis(self):
-    # two sharded axes (the gemm fragment layout): thread (ty, tx) owns the (2, 1, 1, 2) sub-view of the
-    # (2, 4, 2, 2) output tile; the store must SHRINK dest on both sharded axes
-    def kernel(C:UOp, A:UOp) -> UOp:
-      ty = UOp.range(4, 0, AxisType.LOCAL)
-      tx = UOp.range(2, 1, AxisType.LOCAL)
-      frag = UOp.placeholder((2, 1, 1, 2), dtypes.float32, 0, AddrSpace.REG).unshard((1, 2), (ty, tx))
-      v = frag.after(frag.store(0.0)) + A
-      return C.store(v).end(tx, ty).sink(arg=KernelInfo(name="store_unshard_2axis", opts_to_apply=()))
-    a = Tensor(np.arange(32, dtype=np.float32).reshape(2, 4, 2, 2))
-    out = _run_fragment_kernel(self, kernel, (2, 4, 2, 2), inputs=(a,))
-    np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
-
 class TestUOpReduce(unittest.TestCase):
   def test_uop_sum(self):
     a = Tensor([1.0, 2, 3, 4, 5])

test/backend/test_isel.py

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

test/backend/test_linearizer.py

@@ -268,9 +268,9 @@ class TestLinearizer(unittest.TestCase):
     uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
     idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
     idxs = sorted(idxs, key=lambda uop: uop.arg)
-    assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0]
-    assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg
-    assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg
+    assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
+    assert (idxs[1].arg, idxs[1].src[0].arg) == ('gidx1', 5), idxs[1].arg
+    assert (idxs[2].arg, idxs[2].src[0].arg) == ('gidx2', 4), idxs[2].arg
 
   def test_sum_collapse(self):
     t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()

test/backend/test_linearizer_dumb.py

@@ -12,18 +12,18 @@ 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,))
-    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)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL)
     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)
+    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
+    c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE)
     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)
+    c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True)))
+    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
+    c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
     ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
     _ = to_program(ast, Device["METAL"].renderer)
 

test/backend/test_multitensor.py

@@ -76,13 +76,6 @@ class TestMultiTensor(unittest.TestCase):
     run_linear(linear)
     self.assertEqual(len(set(names)), 1, "function was relinearized")
 
-  def test_shard_beam(self):
-    cpu_2 = ("CPU:1", "CPU:2")
-    src = Tensor.ones(16).shard(cpu_2, 0).realize()
-    pad = src.to(cpu_2[::-1]).schedule_linear().src[0]
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): prg = compile_linear(UOp(Ops.LINEAR, src=(pad,))).src[0].src[0]
-    self.assertNotEqual(prg.src[0].arg.applied_opts, ())
-
   def test_shard_same_device(self):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_((d1, X.device), 0)

test/backend/test_ops.py

@@ -1728,8 +1728,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(10,10,10)], lambda x: x.log_softmax(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(10,10,10)], lambda x: x.log_softmax(1), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(10,10,10)], lambda x: x.log_softmax(2), atol=1e-7, grad_atol=1e-7)
-  def test_softmin(self):
-    helper_test_op([(45,65)], torch.nn.Softmin(dim=1), Tensor.softmin, atol=1e-7, grad_atol=1e-7)
 
   def test_normalize(self):
     helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x), lambda x: x.normalize(), atol=1e-7, grad_atol=1e-7)

test/backend/test_pickle.py

@@ -13,7 +13,7 @@ class TestPickle(unittest.TestCase):
 
   def test_pickle_pattern_matcher(self):
     pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
-    sink = UOp.const(2)
+    sink = UOp.const(dtypes.int, 2)
     tt = pm.rewrite(sink)
     pm_str = pickle.dumps(pm)
     pm2 = pickle.loads(pm_str)

test/backend/test_renderer_failures.py

@@ -24,7 +24,7 @@ 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,))
-  idx = UOp.const(0)
+  idx = UOp.const(dtypes.int, 0)
   ld = b.index(idx).load()
   alu = ld.alu(alu_op, *alu_src_uops)
   store = UOp.store(a.index(idx), alu)
@@ -35,7 +35,7 @@ class TestRendererFailures(unittest.TestCase):
   def test_gated_store_with_alu(self):
     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)))
+    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
     ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
@@ -45,7 +45,7 @@ class TestRendererFailures(unittest.TestCase):
     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)))
+    gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
     sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
     ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0]
     np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
@@ -54,7 +54,7 @@ class TestRendererFailures(unittest.TestCase):
 class TestCStyleFailures(unittest.TestCase):
   def test_inline_const_alu(self):
     # CPU doesn't use the max function
-    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int.min+1).cast(dtypes.int))
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
     self.assertEqual(ret[0], 1)
 
   def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
@@ -80,7 +80,7 @@ class TestWGSLFailures(unittest.TestCase):
   def test_multiply_infinity(self):
     # multiplying a positive constant by infinity should return infinity
     # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer
-    ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(float("inf")).cast(dtypes.float32))
+    ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(dtypes.float32, float("inf")))
     self.assertEqual(ret[0], float("inf"))
 
   # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
@@ -104,7 +104,7 @@ class TestPTXFailures(unittest.TestCase):
   def test_gated_store_with_if(self):
     a = UOp.param(0, dtypes.int, (4,))
     gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
-    val = UOp.const(1).cast(dtypes.int)
+    val = UOp.const(dtypes.int, 1)
     if_uop = UOp(Ops.IF, src=(gate_alu,))
     gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0, if_uop), val))
     sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())

test/backend/test_tensor.py

@@ -24,13 +24,13 @@ class TestTinygrad(unittest.TestCase):
     self.assertEqual(Tensor(3.14).shape, ())
 
   def test_deviceless_const_construct_device_repr(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float))
+    t = Tensor(UOp.const(dtypes.float, 2.0))
     self.assertIsNone(t.uop.device)
     self.assertIsNone(t.device)
     self.assertIn("<UOp None", repr(t))
 
   def test_deviceless_const_realize_noop(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float))
+    t = Tensor(UOp.const(dtypes.float, 2.0))
     uop = t.uop
     t.realize()
     self.assertIs(t.uop, uop)
@@ -728,12 +728,12 @@ class TestZeroShapeTensor(unittest.TestCase):
     self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
 
   def test_clone_deviceless_const(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone()
+    t = Tensor(UOp.const(dtypes.float, 2.0)).clone()
     np.testing.assert_equal(t.numpy(), 2.0)
     self.assertTrue(t.uop.has_buffer_identity())
 
   def test_numpy_deviceless_const(self):
-    np.testing.assert_equal(Tensor(UOp.const(2.0).cast(dtypes.float)).numpy(), 2.0)
+    np.testing.assert_equal(Tensor(UOp.const(dtypes.float, 2.0)).numpy(), 2.0)
 
   def test_clone_with_shrink(self):
     a = Tensor.rand(16, 16)
@@ -756,7 +756,7 @@ class TestZeroShapeTensor(unittest.TestCase):
     np.testing.assert_allclose(a.grad.numpy(), b.grad.numpy())
 
   def test_clone_deviceless_const_to_cpu(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone(device="CPU")
+    t = Tensor(UOp.const(dtypes.float, 2.0)).clone(device="CPU")
     self.assertEqual(t.device, "CPU")
     np.testing.assert_equal(t.numpy(), 2.0)
 

test/backend/test_uops.py

@@ -19,7 +19,7 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
   run_linear(UOp(Ops.LINEAR, src=(UOp.sink(*uops_list, arg=KernelInfo()).call(*buf_uops),)))
 
 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))
+  if op is Ops.CONST: uops.append(UOp.const(dtype, arg))
   elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
   else: uops.append(UOp(op, dtype, tuple(src), arg))
   return uops[-1]
@@ -43,7 +43,7 @@ def _test_single_value_const(vals, op, dts):
   buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0)
   loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
   alu = uop(uops, op, output_dtype, loads)
-  out = buf_store[UOp.const(0).cast(dtypes.int32)].store(alu)
+  out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
   run_uops([out], [buf])
   return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
@@ -221,12 +221,12 @@ 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,))
-    c1 = UOp.const(2)
-    c2 = UOp.const(3)
+    c1 = UOp.const(dtypes.int, 2)
+    c2 = UOp.const(dtypes.int, 3)
     l1 = g1.index(c1)
     a1 = UOp(Ops.MUL, src=(l1, c1))
     a2 = UOp(Ops.MUL, src=(l1, c2))
-    uops = to_uops_list([out.index(UOp.const(0)).store(a1), out.index(UOp.const(1)).store(a2)],
+    uops = to_uops_list([out.index(UOp.const(dtypes.int, 0)).store(a1), out.index(UOp.const(dtypes.int, 1)).store(a2)],
                         ren=Device[Device.DEFAULT].renderer)
     Device[Device.DEFAULT].renderer.render(uops)
     ops = [x.op for x in uops]
@@ -249,16 +249,16 @@ class TestAssembly(unittest.TestCase):
 
   def test_mulacc_shl(self):
     g1 = UOp.param(0, dtypes.int32, shape=(2,))
-    c1 = UOp.const(0)
-    c2 = UOp.const(1)
-    expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
+    c1 = UOp.const(dtypes.int, 0)
+    c2 = UOp.const(dtypes.int, 1)
+    expr = g1.index(c1) * UOp.const(dtypes.int, 4096) + g1.index(c2)
     uops = to_uops_list([expr], ren=Device[Device.DEFAULT].renderer)
     Device[Device.DEFAULT].renderer.render(uops)
     self.assertIn(Ops.MULACC, [x.op for x in uops])
 
   def test_use_cmpeq(self):
     g = UOp.param(0, dtypes.uint32, shape=(8,))
-    c = UOp.const(7)
+    c = UOp.const(dtypes.uint, 7)
     comp = g.index(c).ne(c).ne(True)
     uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
     Device[Device.DEFAULT].renderer.render(uops)

test/external/external_benchmark_op_conv.py

@@ -24,8 +24,8 @@ def vision_conv_143():
   c32 = ((c27<3)!=True)&(c27<67)
   c34 = UOp.param(1, dtypes.half, shape=(32, 1024, 4))
   c38 = c5//2
-  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(Invalid, dtypes.weakint))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
   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,))
@@ -50,8 +50,8 @@ def vision_conv_153():
   c32 = ((c27<3)!=True)&(c27<35)
   c34 = UOp.param(1, dtypes.half, shape=(16, 1024, 4))
   c38 = c5//2
-  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(Invalid, dtypes.weakint))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
   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,))

test/external/fuzz_symbolic.py

@@ -40,7 +40,7 @@ def random_int_expr(depth=10):
 def random_bool_expr(depth=10, expr1=None):
   if depth == 0: return True
   if expr1 is None: expr1 = random_int_expr(depth-1)
-  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(random.randint(-10, 10))])
+  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(dtypes.weakint, random.randint(-10, 10))])
   return random.choice(comp_ops)(expr1, expr2)
 
 
@@ -82,8 +82,8 @@ if __name__ == "__main__":
             f"v2=Variable(\"{u2.arg[0]}\", {u2.arg[1]}, {u2.arg[2]})\n" +\
             f"v3=Variable(\"{u3.arg[0]}\", {u3.arg[1]}, {u3.arg[2]})\n" +\
             f"expr = {expr}\n" +\
-            f"v1_val, v2_val, v3_val = UOp.const({n1.as_long()}), UOp.const({n2.as_long()})," +\
-                f"UOp.const({n3.as_long()})\n" +\
+            f"v1_val, v2_val, v3_val = UOp.const(dtypes.weakint, {n1.as_long()}), UOp.const(dtypes.weakint, {n2.as_long()})," +\
+                f"UOp.const(dtypes.weakint, {n3.as_long()})\n" +\
             "num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
             "rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
             "assert num==rn, f\"{num} != {rn}\"\n"

test/helpers.py

@@ -90,7 +90,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
     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,))
-  prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
+  prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 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)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]

test/mockgpu/amd/emu.py

@@ -6,7 +6,7 @@
 #   arg=3: lds - local data share
 #   arg=4: scratch - per-lane scratch memory
 from __future__ import annotations
-import ctypes, functools, re, platform, subprocess, tempfile
+import ctypes, functools, re, platform, subprocess, tempfile, os
 from typing import Callable
 
 # Set/restore DAZ+FTZ (denormals-are-zero + flush-to-zero) to match RDNA3 default float mode
@@ -192,16 +192,16 @@ def _init_sqtt_encoder():
 
   return emit, finish, finalize
 
-def _c(val, dtype=dtypes.uint32): return UOp.const(val, dtype)
+def _c(val, dtype=dtypes.uint32): return UOp.const(dtype, val)
 
 def _u64(lo: UOp, hi: UOp) -> UOp:
   """Combine two 32-bit UOps into a 64-bit UOp."""
-  return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
+  return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
 
 def _split64(val: UOp) -> tuple[UOp, UOp]:
   """Split a 64-bit value into (lo, hi) 32-bit values."""
   v64 = val.bitcast(dtypes.uint64) if val.dtype == dtypes.float64 else val.cast(dtypes.uint64) if val.dtype != dtypes.uint64 else val
-  return v64.cast(dtypes.uint32), (v64 >> UOp.const(32, dtypes.uint64)).cast(dtypes.uint32)
+  return v64.cast(dtypes.uint32), (v64 >> UOp.const(dtypes.uint64, 32)).cast(dtypes.uint32)
 
 _SRC_MOD_TYPES = {16: (dtypes.uint16, dtypes.half, 0x7FFF), 32: (dtypes.uint32, dtypes.float32, 0x7FFFFFFF),
                   64: (dtypes.uint64, dtypes.float64, 0x7FFFFFFFFFFFFFFF)}
@@ -210,7 +210,7 @@ def _apply_src_mods(val: UOp, mod_bit: int, abs_bits: int, neg_bits: int, bits:
   if not (abs_bits & (1 << mod_bit)) and not (neg_bits & (1 << mod_bit)): return val
   ut, ft, mask = _SRC_MOD_TYPES[bits]
   fv = val.cast(ut).bitcast(ft) if bits == 16 else val.bitcast(ft) if val.dtype == ut else val
-  if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(mask, ut)).bitcast(ft)
+  if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(ut, mask)).bitcast(ft)
   if neg_bits & (1 << mod_bit): fv = fv.neg()
   return fv.bitcast(ut).cast(dtypes.uint32) if bits == 16 else fv.bitcast(ut)
 
@@ -251,7 +251,7 @@ def _to_u32(val: UOp) -> UOp:
   if val.dtype.itemsize == 4: return val.bitcast(dtypes.uint32)  # same size: bitcast (float32->uint32)
   return val.cast(dtypes.uint32)  # different size: cast (bool, int16, etc)
 def _lane_active(exec_mask: UOp, lane: UOp) -> UOp:
-  if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(1, dtypes.uint64)).ne(UOp.const(0, dtypes.uint64))
+  if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(dtypes.uint64, 1)).ne(UOp.const(dtypes.uint64, 0))
   return ((exec_mask >> lane.cast(dtypes.uint32)) & _c(1)).ne(_c(0))
 def _hi16(v: UOp) -> UOp: return (v >> _c(16)) & _c(0xFFFF)
 def _cond(cond, if_true, if_false):
@@ -264,9 +264,9 @@ def _set_lane_bit(old: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> UOp:
   """Set/clear a single bit in a mask based on lane index, respecting exec mask."""
   if old.dtype in (dtypes.uint64, dtypes.int64):
     dt = dtypes.uint64
-    mask = UOp.const(1, dt) << lane.cast(dt)
+    mask = UOp.const(dt, 1) << lane.cast(dt)
     new_bit = _to_u32(val).cast(dt) << lane.cast(dt)
-    cleared = old.cast(dt) & (mask ^ UOp.const(0xFFFFFFFFFFFFFFFF, dt))
+    cleared = old.cast(dt) & (mask ^ UOp.const(dt, 0xFFFFFFFFFFFFFFFF))
     return _lane_active(exec_mask, lane).where(cleared | new_bit, old.cast(dt))
   mask = _c(1) << lane.cast(dtypes.uint32)
   new_bit = _to_u32(val) << lane.cast(dtypes.uint32)
@@ -365,18 +365,40 @@ def _write_64bit(val: UOp, wfn, reg_or_addr, is_mem: bool, *args) -> list[UOp]:
   """Write a 64-bit value as two 32-bit writes. args passed to wfn after reg/addr and lo/hi value."""
   lo, hi = _split64(val)
   incr = 4 if is_mem else 1  # 4 bytes for memory addresses, 1 for register indices
-  return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(incr, reg_or_addr.dtype) if isinstance(reg_or_addr, UOp) else incr), hi, *args)]
+  return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(reg_or_addr.dtype, incr) if isinstance(reg_or_addr, UOp) else incr), hi, *args)]
 
 def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = False) -> list[UOp]:
   """Write value, splitting 64-bit if needed. bits=64 for 64-bit writes, otherwise 32-bit."""
   return _write_64bit(val, wfn, reg_or_addr, is_mem, *args) if bits == 64 else [wfn(reg_or_addr, _to_u32(val), *args)]
 
 def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32, data_bits: int = 32) -> list[UOp]:
-  """Conditional memory store with sub-word support. Returns list of store UOps."""
+  """Conditional memory store with sub-word and unaligned support. Returns list of store UOps.
+
+  FLAT/GLOBAL accesses on AMD hardware are allowed to be unaligned, so a 32-bit store at a
+  byte offset of 1-3 spans two words (handled in 64-bit domain to keep shifts in range)."""
   adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32
-  word_addr = addr >> UOp.const(2, adt)
+  if data_bits > 32:   # wider stores decompose into dwords; each dword handles its own alignment
+    ws = val.cast(dtypes.uint64) if data_bits > 64 else val
+    return [s for i in range(data_bits // 32)
+            for s in _mem_store(mem, addr + UOp.const(adt, i * 4), ws >> UOp.const(ws.dtype, 32 * i) if i else ws, active, addr_bits, 32)]
+  word_addr = addr >> UOp.const(adt, 2)
   idx = mem.index(word_addr.valid(active))
-  if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))]
+  if data_bits == 32:
+    byte_off = (addr & UOp.const(adt, 3)).cast(dtypes.uint32)
+    is_unaligned = byte_off.ne(UOp.const(dtypes.uint32, 0))
+    if addr.divides(4) is not None: return [idx.store(active.where(_to_u32(val), idx))]
+    shift = byte_off * UOp.const(dtypes.uint32, 8)
+    val64 = _to_u32(val).cast(dtypes.uint64)
+    # word0 keeps its low byte_off*8 bits, gets val's low bits shifted in; word1 gets the rest
+    low_keep = (UOp.const(dtypes.uint32, 1) << shift) - UOp.const(dtypes.uint32, 1)
+    lo_bits = ((val64 << shift.cast(dtypes.uint64)) & UOp.const(dtypes.uint64, 0xFFFFFFFF)).cast(dtypes.uint32)
+    new_word0 = (idx & low_keep) | lo_bits
+    store0 = idx.store(active.where(is_unaligned.where(new_word0, _to_u32(val)), idx))
+    idx1 = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int64).valid(active & is_unaligned))
+    spill = (val64 >> (UOp.const(dtypes.uint64, 32) - shift.cast(dtypes.uint64))).cast(dtypes.uint32)
+    keep = UOp.const(dtypes.uint32, 0xFFFFFFFF) << shift
+    new_word1 = (idx1 & keep) | spill
+    return [store0, idx1.store((active & is_unaligned).where(new_word1, idx1))]
   # Sub-word store: read-modify-write with mask
   byte_pos = addr.cast(dtypes.uint32) & _c(3)
   byte_shift = byte_pos * _c(8)
@@ -388,7 +410,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
   is_cross = byte_pos.eq(_c(3))
   cross_word0 = (idx & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24))
   store0 = idx.store(active.where(is_cross.where(cross_word0, new_word), idx))
-  next_idx = mem.index((word_addr + UOp.const(1, adt)).valid(active & is_cross))
+  next_idx = mem.index((word_addr + UOp.const(adt, 1)).valid(active & is_cross))
   cross_word1 = (next_idx & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF))
   return [store0, next_idx.store((active & is_cross).where(cross_word1, next_idx))]
 
@@ -397,8 +419,8 @@ def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int
   stores = []
   val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
   for i in range(data_bits // 8):
-    byte_val = (val_u32 >> UOp.const(i * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-    stores.append(mem.index((addr + UOp.const(i, dtypes.uint64)).valid(active)).store(byte_val.cast(dtypes.uint8)))
+    byte_val = (val_u32 >> UOp.const(dtypes.uint32, i * 8)) & UOp.const(dtypes.uint32, 0xFF)
+    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).valid(active)).store(byte_val.cast(dtypes.uint8)))
   return stores
 
 def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]:
@@ -462,8 +484,8 @@ class _Ctx:
   def inst_word(self, dword_idx: int) -> UOp:
     """Read instruction dword from vmem at PC + dword_idx*4."""
     pc = self.rpc()
-    addr = pc if dword_idx == 0 else pc + UOp.const(dword_idx * 4, dtypes.uint64)
-    return self.vmem.index(addr >> UOp.const(2, dtypes.uint64)).load()
+    addr = pc if dword_idx == 0 else pc + UOp.const(dtypes.uint64, dword_idx * 4)
+    return self.vmem.index(addr >> UOp.const(dtypes.uint64, 2)).load()
 
   def inst_field(self, field) -> UOp:
     """Extract field bits from instruction encoding. Tracks field for canonical key computation."""
@@ -475,15 +497,15 @@ class _Ctx:
     word = self.inst_word(dword_idx)
     if lo // 32 == hi // 32:  # Same dword
       mask = (1 << (hi - lo + 1)) - 1
-      shifted = word if lo_in_dword == 0 else word >> UOp.const(lo_in_dword, dtypes.uint32)
-      return shifted & UOp.const(mask, dtypes.uint32)
+      shifted = word if lo_in_dword == 0 else word >> UOp.const(dtypes.uint32, lo_in_dword)
+      return shifted & UOp.const(dtypes.uint32, mask)
     else:  # Spans two dwords
       lo_bits = 32 - lo_in_dword
       lo_mask = (1 << lo_bits) - 1
       hi_mask = (1 << (hi_in_dword + 1)) - 1
-      lo_part = (word >> UOp.const(lo_in_dword, dtypes.uint32)) & UOp.const(lo_mask, dtypes.uint32)
-      hi_part = self.inst_word(dword_idx + 1) & UOp.const(hi_mask, dtypes.uint32)
-      return lo_part | (hi_part << UOp.const(lo_bits, dtypes.uint32))
+      lo_part = (word >> UOp.const(dtypes.uint32, lo_in_dword)) & UOp.const(dtypes.uint32, lo_mask)
+      hi_part = self.inst_word(dword_idx + 1) & UOp.const(dtypes.uint32, hi_mask)
+      return lo_part | (hi_part << UOp.const(dtypes.uint32, lo_bits))
 
   def inst_field_signed(self, field) -> UOp:
     """Extract field and sign-extend based on field width."""
@@ -584,7 +606,7 @@ class _Ctx:
       inline = is_float_const.where(float_inline.bitcast(dtypes.uint64), int_inline.bitcast(dtypes.uint64))
       # Literal handling: F64 VOP puts literal in high 32 bits; B64/I64/U64 VOP and SOP zero-extend
       if literal is not None:
-        lit_val = literal.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64) if is_f64 else literal.cast(dtypes.uint64)
+        lit_val = literal.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32) if is_f64 else literal.cast(dtypes.uint64)
         inline = off.eq(_c(255)).where(lit_val, inline)
       scalar_val = (off < _c(128)).where(sgpr_val, inline)
     else:
@@ -602,7 +624,7 @@ class _Ctx:
 
   def inc_pc(self) -> list[UOp]:
     """Increment PC by instruction size in bytes. Returns [store]."""
-    new_pc = self.rpc() + UOp.const(self.inst_size, dtypes.uint64)
+    new_pc = self.rpc() + UOp.const(dtypes.uint64, self.inst_size)
     lo, hi = _split64(new_pc)
     return [self.wsgpr_dyn(_c(PC_LO_IDX), lo), self.wsgpr_dyn(_c(PC_HI_IDX), hi)]
 
@@ -668,7 +690,7 @@ class _Ctx:
     if 'VCC' not in srcs: srcs['VCC'] = self.rmask(_c(vcc_reg))
     srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane, 'VDST': vdst_reg,
                   'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0), 'ROUND_NEAREST_EVEN': _c(0), '_vgpr': self.vgpr, '_wave_size': self.wave_size,
-                  'MAX_FLOAT_F32': UOp.const(3.4028234663852886e38, dtypes.float32),
+                  'MAX_FLOAT_F32': UOp.const(dtypes.float32, 3.4028234663852886e38),
                   # CDNA SDWA byte/word select constants (E32 always uses BYTE0/WORD0 defaults)
                   'SDWA_SRC0_SEL': _c(0), 'BYTE0': _c(0), 'BYTE1': _c(1), 'BYTE2': _c(2), 'BYTE3': _c(3),
                   'WORD0': _c(0), 'WORD1': _c(1)})  # rounding mode and SDWA constants
@@ -709,7 +731,7 @@ class _Ctx:
       # VGPR bit-slice assignment: VGPR[lane][reg][hi:lo] = (vgpr_idx, rhs_val, hi, lo[, cond]) -> read-modify-write
       if dest.startswith('VGPR[') and re.search(r'\[\d+:\d+\]', dest):
         # VGPR bit-slice: (vgpr_idx, rhs_val, hi_bit, lo_bit) - hi/lo are UOp constants
-        hi_bit, lo_bit = int(val[2].val), int(val[3].val)
+        hi_bit, lo_bit = int(val[2].arg), int(val[3].arg)
         width = hi_bit - lo_bit + 1
         old = self.vgpr.index(val[0]).load()
         new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32)
@@ -732,24 +754,24 @@ class _Ctx:
             d0_width, slice_mask = d0_hi_bit - d0_lo_bit + 1, (1 << (d0_hi_bit - d0_lo_bit + 1)) - 1
             val_bits = val.bitcast(dtypes.uint16).cast(dtypes.uint32) if val.dtype == dtypes.half else \
                        val.cast(dtypes.uint32) if val.dtype in (dtypes.uint16, dtypes.int16) else \
-                       val.cast(dtypes.uint32) & UOp.const(slice_mask, dtypes.uint32)
+                       val.cast(dtypes.uint32) & UOp.const(dtypes.uint32, slice_mask)
             raw_stores.append(('vgpr_slice', (d0_lo_bit, d0_width, val_bits)))
             continue
         # For integer ops with clamp, use pre-computed saturated value; for floats, clamp to [0,1]
         if int_saturate is not None: val = int_saturate
         elif clmp and val.dtype in (dtypes.float32, dtypes.half, dtypes.float64):
-          clamped = val.maximum(UOp.const(0.0, val.dtype)).minimum(UOp.const(1.0, val.dtype))
-          val = _FUNCS['isNAN'](val).where(UOp.const(0.0, val.dtype), clamped)
+          clamped = val.maximum(UOp.const(val.dtype, 0.0)).minimum(UOp.const(val.dtype, 1.0))
+          val = _FUNCS['isNAN'](val).where(UOp.const(val.dtype, 0.0), clamped)
         if val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64):
           lo, hi = _split64(val)
           raw_stores.extend([('vgpr', self.wvgpr_dyn(vdst_reg, lane, lo, exec_mask)),
                              ('vgpr', self.wvgpr_dyn(vdst_reg + _c(1), lane, hi, exec_mask))])
         elif val.dtype in (dtypes.half, dtypes.uint16, dtypes.int16):
           result, old_val = _val_to_u32(val), self.rvgpr_dyn(vdst_reg, lane)
-          hi_result = (old_val & UOp.const(0xFFFF, dtypes.uint32)) | (result << UOp.const(16, dtypes.uint32))
+          hi_result = (old_val & UOp.const(dtypes.uint32, 0xFFFF)) | (result << UOp.const(dtypes.uint32, 16))
           # GFX9/CDNA zeroes upper 16 bits on lo-half write; RDNA preserves them
-          lo_result = (result & UOp.const(0xFFFF, dtypes.uint32)) if self.wave_size == 64 else \
-                      (old_val & UOp.const(0xFFFF0000, dtypes.uint32)) | (result & UOp.const(0xFFFF, dtypes.uint32))
+          lo_result = (result & UOp.const(dtypes.uint32, 0xFFFF)) if self.wave_size == 64 else \
+                      (old_val & UOp.const(dtypes.uint32, 0xFFFF0000)) | (result & UOp.const(dtypes.uint32, 0xFFFF))
           result = opsel_dst_hi.where(hi_result, lo_result) if isinstance(opsel_dst_hi, UOp) else hi_result if opsel_dst_hi else lo_result
           raw_stores.append(('vgpr', self.wvgpr_dyn(vdst_reg, lane, result, exec_mask)))
         else: raw_stores.append(('vgpr', self.wvgpr_dyn(vdst_reg, lane, _val_to_u32(val), exec_mask)))
@@ -767,8 +789,8 @@ class _Ctx:
     if slice_stores:
       result = self.rvgpr_dyn(vdst_reg, lane)
       for lo_bit, width, val_bits in slice_stores:
-        mask = UOp.const(((1 << width) - 1) << lo_bit, dtypes.uint32)
-        result = (result & (mask ^ UOp.const(0xFFFFFFFF, dtypes.uint32))) | (val_bits << UOp.const(lo_bit, dtypes.uint32))
+        mask = UOp.const(dtypes.uint32, ((1 << width) - 1) << lo_bit)
+        result = (result & (mask ^ UOp.const(dtypes.uint32, 0xFFFFFFFF))) | (val_bits << UOp.const(dtypes.uint32, lo_bit))
       lane_stores.append(self.wvgpr_dyn(vdst_reg, lane, result, exec_mask))
     # VCC/EXEC mask writes must be computed BEFORE VGPR stores to avoid reading modified VGPRs.
     # When vdst overlaps with src operands (e.g. v_add_co_u32 v[0], vcc, s[8], v[0]), the carry
@@ -790,8 +812,8 @@ class _Ctx:
 def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
   simm16 = ctx.inst_field_signed(type(inst).simm16).cast(dtypes.int16)
   if inst.op in (ir3.SOPPOp.S_ENDPGM, ir4.SOPPOp.S_ENDPGM, irc.SOPPOp.S_ENDPGM):
-    return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(0xFFFFFFFF, dtypes.uint32)),
-                          ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(0xFFFFFFFF, dtypes.uint32)))
+    return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)),
+                          ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)))
   # S_BARRIER: advance PC past the barrier instruction. The execution loop detects barriers before executing and handles synchronization.
   barrier_ops = {ir3.SOPPOp.S_BARRIER, irc.SOPPOp.S_BARRIER}
   if hasattr(ir4.SOPPOp, 'S_BARRIER_WAIT'): barrier_ops.add(ir4.SOPPOp.S_BARRIER_WAIT)
@@ -804,8 +826,8 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
     pc_bytes = ctx.rpc()  # PC is already 64-bit byte address
     vcc, exec_val = ctx.rmask(_c(VCC_LO.offset)), ctx.rexec()
     srcs: dict[str, UOp|int] = {'PC': pc_bytes.cast(dtypes.int64), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc,
-            'VCCZ': vcc.eq(UOp.const(0, vcc.dtype)).cast(dtypes.uint32),
-            'EXECZ': exec_val.eq(UOp.const(0, exec_val.dtype)).cast(dtypes.uint32)}
+            'VCCZ': vcc.eq(UOp.const(vcc.dtype, 0)).cast(dtypes.uint32),
+            'EXECZ': exec_val.eq(UOp.const(exec_val.dtype, 0)).cast(dtypes.uint32)}
     for dest, val in parse_pcode(pcode, srcs)[1]:
       if dest == 'PC' or dest.startswith('PC.'):
         lo, hi = _split64(val.cast(dtypes.uint64))
@@ -833,12 +855,12 @@ def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp:
   part = op_name.rsplit('_', 1)[1]  # B32, DWORD, DWORDX2, U8, I8, etc.
   nval = int(part.removeprefix('DWORD').removeprefix('X') or '1') if 'DWORD' in part else int(part[1:]) / 32 * (-1 if part[0] == 'I' else 1)
   ndwords = max(1, int(abs(nval)))
-  dword_base = addr >> UOp.const(2, dtypes.uint64)
-  vals = [ctx.vmem.index(dword_base + UOp.const(i, dtypes.uint64)) for i in range(ndwords)]
+  dword_base = addr >> UOp.const(dtypes.uint64, 2)
+  vals = [ctx.vmem.index(dword_base + UOp.const(dtypes.uint64, i)) for i in range(ndwords)]
   if abs(nval) < 1:
     nbits = int(abs(nval) * 32)
-    byte_off = (addr & UOp.const(3, dtypes.uint64)).cast(dtypes.uint32) * UOp.const(8, dtypes.uint32)
-    extracted = (vals[0] >> byte_off) & UOp.const((1 << nbits) - 1, dtypes.uint32)
+    byte_off = (addr & UOp.const(dtypes.uint64, 3)).cast(dtypes.uint32) * UOp.const(dtypes.uint32, 8)
+    extracted = (vals[0] >> byte_off) & UOp.const(dtypes.uint32, (1 << nbits) - 1)
     vals[0] = extracted.cast({8: dtypes.int8, 16: dtypes.int16}[nbits]).cast(dtypes.int32).bitcast(dtypes.uint32) if nval < 0 else extracted
   stores = [ctx.wsgpr_dyn(sdata_reg + _c(i), vals[i]) for i in range(ndwords)]
   return UOp.sink(*stores, *ctx.inc_pc())
@@ -941,7 +963,7 @@ def _dpp16_ctrl(lane: UOp, dpp: int, row_mask: int, bank_mask: int, wave_size: i
   enabled = (((_c(row_mask) >> row.cast(dtypes.uint32)) & _c(1)).ne(_c(0)) &
              (((_c(bank_mask) >> bank.cast(dtypes.uint32)) & _c(1)).ne(_c(0))))
   op, arg = decode_dpp16(dpp)
-  src_lane, valid = lane_i, UOp.const(True)
+  src_lane, valid = lane_i, UOp.const(dtypes.bool, True)
 
   if op == 'quad_perm':
     assert isinstance(arg, tuple)
@@ -967,7 +989,7 @@ def _load_dpp16_src0(ctx: _Ctx, inst, lane: UOp, fallback: UOp) -> UOp:
                                          getattr(inst, 'bank_mask', 0xf) or 0xf, ctx.wave_size)
   safe_src_lane = (enabled & valid).where(src_lane, _c(0, dtypes.int))
   swizzled = ctx.rvgpr_dyn(ctx.inst_field(type(inst).vsrc0), safe_src_lane)
-  invalid = UOp.const(0, fallback.dtype) if getattr(inst, 'bc', 0) else fallback
+  invalid = UOp.const(fallback.dtype, 0) if getattr(inst, 'bc', 0) else fallback
   return enabled.where(valid.where(swizzled, invalid), fallback)
 
 def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc.VOPC_SDWA_SDST, ctx: _Ctx) -> UOp:
@@ -1159,7 +1181,7 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16
         s1 = _apply_src_mods(s1, 0, 1 if getattr(inst, 'src1_abs', 0) else 0, 1 if getattr(inst, 'src1_neg', 0) else 0, bits['s1'])
       s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0'])
       s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1'])
-    for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(0, dtypes.uint64)})[1]:
+    for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(dtypes.uint64, 0)})[1]:
       if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32)
     return _c(0)
 
@@ -1189,8 +1211,8 @@ def _compile_bitop3(inst, ctx: _Ctx, exec_mask: UOp, bits: dict, op_name: str) -
   is_16 = 'B16' in op_name
   dt, mask = (dtypes.uint16, 0xFFFF) if is_16 else (dtypes.uint32, 0xFFFFFFFF)
   s0, s1, s2 = src0.cast(dt), src1.cast(dt), src2.cast(dt)
-  def bnot(v): return v ^ UOp.const(mask, dt)
-  result = UOp.const(0, dt)
+  def bnot(v): return v ^ UOp.const(dt, mask)
+  result = UOp.const(dt, 0)
   for i in range(8):
     if not (ttbl & (1 << i)): continue
     result = result | ((s0 if i & 4 else bnot(s0)) & (s1 if i & 2 else bnot(s1)) & (s2 if i & 1 else bnot(s2)))
@@ -1244,7 +1266,7 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp:
   src0 = _apply_src_mods(src0, 0, abs_bits, neg_bits, bits['s0'])
   src1 = _apply_src_mods(src1, 1, abs_bits, neg_bits, bits['s1'])
   src2 = _apply_src_mods(src2, 2, abs_bits, neg_bits, bits['s2'])
-  srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(opsel, dtypes.uint32)}
+  srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(dtypes.uint32, opsel)}
   if 'CNDMASK' in op_name and src2 is not None: srcs['VCC'] = src2
   # FMAC instructions need D0 (accumulator) from destination register
   if 'FMAC' in op_name: srcs['D0'] = ctx.rvgpr_dyn(vdst_reg, lane)
@@ -1398,7 +1420,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
   # Layout: tmp[0..n_a_elems-1] = A[m][k], tmp[n_a_elems..n_a_elems+n_b_elems-1] = B[n][k]
   # Within each group of lanes, lane%grp_sub gives M/N index, lane//grp_sub gives sub-block
   grp_sub = min(M, 16)  # lanes within group mapped to M/N dimension
-  b_off = UOp.const(n_a_elems, dtypes.int)
+  b_off = UOp.const(dtypes.int, n_a_elems)
   acc_dt = dtypes.int32 if is_int_out else dtypes.float32
   # Use uint32 temp array to prevent optimizer from eliminating f16→f32 bitcast chains.
   # The optimizer folds bitcast(uint32→float32) stores to float32 arrays, losing the conversion.
@@ -1407,44 +1429,44 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
   def cvt_elem(raw: UOp, sub_idx: int) -> UOp:
     if is_i8:
       # Extract i8, sign-extend to i32
-      byte_val = (raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-      return (byte_val.cast(dtypes.int32) ^ UOp.const(0x80, dtypes.int32)) - UOp.const(0x80, dtypes.int32)
+      byte_val = (raw >> UOp.const(dtypes.uint32, sub_idx * 8)) & UOp.const(dtypes.uint32, 0xFF)
+      return (byte_val.cast(dtypes.int32) ^ UOp.const(dtypes.int32, 0x80)) - UOp.const(dtypes.int32, 0x80)
     elif is_f32_src:
       return raw  # already uint32 (f32 bit pattern)
     elif is_fp8:
-      return ((raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)).cast(dtypes.uint32)
+      return ((raw >> UOp.const(dtypes.uint32, sub_idx * 8)) & UOp.const(dtypes.uint32, 0xFF)).cast(dtypes.uint32)
     elif is_bf16:
       # bf16→f32 bits: just shift left by 16 (bf16 is upper 16 bits of f32)
-      return ((raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)) << UOp.const(16, dtypes.uint32)
+      return ((raw >> UOp.const(dtypes.uint32, sub_idx * 16)) & UOp.const(dtypes.uint32, 0xFFFF)) << UOp.const(dtypes.uint32, 16)
     else:
       # f16→f32 conversion using float arithmetic to avoid UOp optimizer eliminating the conversion.
       # The optimizer folds bitcast(uint32→float32) chains, so we compute the float value directly.
-      h = (raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)
-      sign = (h >> UOp.const(15, dtypes.uint32)) & UOp.const(1, dtypes.uint32)
-      exp = (h >> UOp.const(10, dtypes.uint32)) & UOp.const(0x1F, dtypes.uint32)
-      mant = h & UOp.const(0x3FF, dtypes.uint32)
+      h = (raw >> UOp.const(dtypes.uint32, sub_idx * 16)) & UOp.const(dtypes.uint32, 0xFFFF)
+      sign = (h >> UOp.const(dtypes.uint32, 15)) & UOp.const(dtypes.uint32, 1)
+      exp = (h >> UOp.const(dtypes.uint32, 10)) & UOp.const(dtypes.uint32, 0x1F)
+      mant = h & UOp.const(dtypes.uint32, 0x3FF)
       # Use bf16 path: shift left by 16 to create bf16 bits, then shift mantissa and adjust exponent in float domain
       # bf16 bits = (sign << 15) | (exp_bf16 << 7) | mant_bf16 -- but f16 and bf16 have different formats
       # Instead: construct f32 bits properly, use a local uint32 array to force materialization
-      f32_bits = (sign << UOp.const(31, dtypes.uint32)) | \
-                 ((exp + UOp.const(112, dtypes.uint32)) << UOp.const(23, dtypes.uint32)) | \
-                 (mant << UOp.const(13, dtypes.uint32))
-      is_zero = exp.eq(UOp.const(0, dtypes.uint32))
+      f32_bits = (sign << UOp.const(dtypes.uint32, 31)) | \
+                 ((exp + UOp.const(dtypes.uint32, 112)) << UOp.const(dtypes.uint32, 23)) | \
+                 (mant << UOp.const(dtypes.uint32, 13))
+      is_zero = exp.eq(UOp.const(dtypes.uint32, 0))
       # Return uint32 (f32 bit pattern) — stored directly to uint32 temp array, bitcast to float on read
-      return is_zero.where(UOp.const(0, dtypes.uint32), f32_bits)
+      return is_zero.where(UOp.const(dtypes.uint32, 0), f32_bits)
 
   read_lane = ctx.range()
   # For 32x32: lane%16 = M/N index within 16-wide block, lane//16 = which of 4 quarter-waves
   # Groups: lanes 0-31 = group 0, lanes 32-63 = group 1
   # Within group: (lane%32)%16 = M/N[0-15], (lane%32)//16 selects M/N[0-15] or [16-31]
-  lane_in_grp = read_lane % UOp.const(grp_size, dtypes.int)
-  grp_idx = read_lane // UOp.const(grp_size, dtypes.int)
+  lane_in_grp = read_lane % UOp.const(dtypes.int, grp_size)
+  grp_idx = read_lane // UOp.const(dtypes.int, grp_size)
 
   if M == 32:
     # 32x32: lane_in_grp%16 = sub-row/col (0-15), lane_in_grp//16 = block (0=rows 0-15, 1=rows 16-31)
-    sub_mn = lane_in_grp % UOp.const(16, dtypes.int)
-    block_mn = lane_in_grp // UOp.const(16, dtypes.int)
-    mn_idx = block_mn * UOp.const(16, dtypes.int) + sub_mn  # actual M/N index (0-31)
+    sub_mn = lane_in_grp % UOp.const(dtypes.int, 16)
+    block_mn = lane_in_grp // UOp.const(dtypes.int, 16)
+    mn_idx = block_mn * UOp.const(dtypes.int, 16) + sub_mn  # actual M/N index (0-31)
   else:
     mn_idx = lane_in_grp  # for 16x16 and 4x4
 
@@ -1456,18 +1478,18 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
                                 ctx.rsrc_dyn(src0_off, _c(0, dtypes.int), 32))
     a_val = cvt_elem(a_raw, sub_idx)
     if M == 4:
-      a_idx = grp_idx * UOp.const(M * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
+      a_idx = grp_idx * UOp.const(dtypes.int, M * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl)
     else:
-      a_idx = mn_idx * UOp.const(K, dtypes.int) + grp_idx * UOp.const(k_per_grp, dtypes.int) + UOp.const(kl, dtypes.int)
+      a_idx = mn_idx * UOp.const(dtypes.int, K) + grp_idx * UOp.const(dtypes.int, k_per_grp) + UOp.const(dtypes.int, kl)
     read_stores.append(tmp.index(a_idx).store(a_val))
 
     b_raw = src1_is_vgpr.where(ctx.rvgpr_dyn(src1_r + _c(reg_idx), read_lane),
                                 ctx.rsrc_dyn(src1_off, _c(0, dtypes.int), 32))
     b_val = cvt_elem(b_raw, sub_idx)
     if M == 4:
-      b_idx = b_off + grp_idx * UOp.const(N * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
+      b_idx = b_off + grp_idx * UOp.const(dtypes.int, N * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl)
     else:
-      b_idx = b_off + mn_idx * UOp.const(K, dtypes.int) + grp_idx * UOp.const(k_per_grp, dtypes.int) + UOp.const(kl, dtypes.int)
+      b_idx = b_off + mn_idx * UOp.const(dtypes.int, K) + grp_idx * UOp.const(dtypes.int, k_per_grp) + UOp.const(dtypes.int, kl)
     read_stores.append(tmp.index(b_idx).store(b_val))
 
   read_phase = UOp.group(*read_stores).end(read_lane)
@@ -1488,11 +1510,11 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
     # Lane mapping: n = (lane%32)%16 + ((lane%32)//16)*16, gives column 0-31
     # Row groups: 4 groups of 4, covering rows 0-31. Group g (0-3): rows g*4 .. g*4+3
     # group assignment: lane//16 gives quarter (0-3), each quarter maps to 4 rows
-    c_lane_in_32 = compute_lane % UOp.const(32, dtypes.int)
-    c_sub = c_lane_in_32 % UOp.const(16, dtypes.int)
-    c_block = c_lane_in_32 // UOp.const(16, dtypes.int)
-    n_idx = c_block * UOp.const(16, dtypes.int) + c_sub
-    c_half = compute_lane // UOp.const(32, dtypes.int)  # 0 or 1
+    c_lane_in_32 = compute_lane % UOp.const(dtypes.int, 32)
+    c_sub = c_lane_in_32 % UOp.const(dtypes.int, 16)
+    c_block = c_lane_in_32 // UOp.const(dtypes.int, 16)
+    n_idx = c_block * UOp.const(dtypes.int, 16) + c_sub
+    c_half = compute_lane // UOp.const(dtypes.int, 32)  # 0 or 1
 
     for out_reg in range(16):
       # Each half covers 8 rows. out_reg 0-3: rows 0-3 (half0) or 16-19 (half1)
@@ -1503,7 +1525,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
       # acc[r] at lane l -> D[m][n] where n = (l%32)%16 + ((l%32)//16)*16
       # m = (l//32)*16 + (r//4)*4 + (r%4)  ... giving rows in blocks of 4
       # So: m_base = half * 16 + (out_reg // 4) * 4 + (out_reg % 4)
-      m_base = c_half * UOp.const(16, dtypes.int) + UOp.const((out_reg // 4) * 4 + (out_reg % 4), dtypes.int)
+      m_base = c_half * UOp.const(dtypes.int, 16) + UOp.const(dtypes.int, (out_reg // 4) * 4 + (out_reg % 4))
 
       acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr)
       if is_int_out: acc_v = acc_v.cast(dtypes.int32)
@@ -1511,8 +1533,8 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
       acc = src2_is_vgpr.where(acc_v, acc_scalar)
 
       for k in range(K):
-        a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-        b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
+        a_val = tmp2.index(m_base * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
+        b_val = tmp2.index(b_off + n_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
         acc = acc + a_val * b_val
 
       if is_int_out:
@@ -1523,8 +1545,8 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
           vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask))
   else:
     # 16x16 and 4x4: each lane computes out_per_lane outputs
-    n_idx = compute_lane % UOp.const(grp_sub, dtypes.int)
-    c_grp = compute_lane // UOp.const(grp_sub, dtypes.int)
+    n_idx = compute_lane % UOp.const(dtypes.int, grp_sub)
+    c_grp = compute_lane // UOp.const(dtypes.int, grp_sub)
 
     for out_reg in range(out_per_lane):
       acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr)
@@ -1534,17 +1556,17 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
 
       if M == 4:
         # 4x4: each group is independent. A/B indexed per-group.
-        m_base = c_grp * UOp.const(M * K, dtypes.int) + UOp.const(out_reg * K, dtypes.int)
+        m_base = c_grp * UOp.const(dtypes.int, M * K) + UOp.const(dtypes.int, out_reg * K)
         for k in range(K):
-          a_val = tmp2.index(m_base + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-          b_val = tmp2.index(b_off + c_grp * UOp.const(N*K, dtypes.int) + n_idx * UOp.const(K, dtypes.int)+UOp.const(k, dtypes.int)).bitcast(acc_dt)
+          a_val = tmp2.index(m_base + UOp.const(dtypes.int, k)).bitcast(acc_dt)
+          b_val = tmp2.index(b_off + c_grp * UOp.const(dtypes.int, N*K) + n_idx * UOp.const(dtypes.int, K)+UOp.const(dtypes.int, k)).bitcast(acc_dt)
           acc = acc + a_val * b_val
       else:
         # 16x16: K is split across groups. Shared MxK/NxK arrays.
-        m_base = c_grp * UOp.const(out_per_lane, dtypes.int) + UOp.const(out_reg, dtypes.int)
+        m_base = c_grp * UOp.const(dtypes.int, out_per_lane) + UOp.const(dtypes.int, out_reg)
         for k in range(K):
-          a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-          b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
+          a_val = tmp2.index(m_base * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
+          b_val = tmp2.index(b_off + n_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
           acc = acc + a_val * b_val
 
       if is_int_out:
@@ -1570,8 +1592,8 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
   is_rdna4 = isinstance(inst, ir4.VOP3P)
   # read 16x16 F16/BF16 matrix from VGPRs → flat f32 array[row*16+k]
   def read_f16_val(src, lane, vgpr, half):
-    v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(lane, dtypes.int))
-    return cvt((v >> UOp.const(16, dtypes.uint32)) if half else (v & UOp.const(0xFFFF, dtypes.uint32)))
+    v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(dtypes.int, lane))
+    return cvt((v >> UOp.const(dtypes.uint32, 16)) if half else (v & UOp.const(dtypes.uint32, 0xFFFF)))
 
   # RDNA3: 16 lanes × 8 VGPRs × 2 halves, k maps linearly
   # RDNA4: 32 lanes × 4 VGPRs × 2 halves, k bits are scrambled (k[2] goes to lane bit 4)
@@ -1593,20 +1615,20 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
              for m in range(16) for n in range(16) for lane, vgpr in [d_map(m, n)]]
     mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)]
     def f32_to_f16_bits(v: UOp) -> UOp: return v.cast(dtypes.half).bitcast(dtypes.uint16).cast(dtypes.uint32)
-    def f32_to_bf16_bits(v: UOp) -> UOp: return (v.bitcast(dtypes.uint32) >> UOp.const(16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)
+    def f32_to_bf16_bits(v: UOp) -> UOp: return (v.bitcast(dtypes.uint32) >> UOp.const(dtypes.uint32, 16)) & UOp.const(dtypes.uint32, 0xFFFF)
     out_cvt = f32_to_bf16_bits if is_bf16 else f32_to_f16_bits
     if is_rdna4:  # pack 2 f16 per VGPR: adjacent m values share (lane, vgpr) since vgpr=m&7, half=m&1
-      stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1] // 2), UOp.const(d_map(m, n)[0], dtypes.int),
-                out_cvt(mat_d[m*16+n]) | (out_cvt(mat_d[(m+1)*16+n]) << UOp.const(16, dtypes.uint32)), exec_mask)
+      stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1] // 2), UOp.const(dtypes.int, d_map(m, n)[0]),
+                out_cvt(mat_d[m*16+n]) | (out_cvt(mat_d[(m+1)*16+n]) << UOp.const(dtypes.uint32, 16)), exec_mask)
                 for n in range(16) for m in range(0, 16, 2)]
     else:  # (rdna3) 1 f16 per VGPR (lo half only)
-      stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), out_cvt(mat_d[m*16+n]), exec_mask)
+      stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(dtypes.int, d_map(m, n)[0]), out_cvt(mat_d[m*16+n]), exec_mask)
                 for m in range(16) for n in range(16)]
   else: # f32
-    mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int)).bitcast(dtypes.float32)
+    mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(dtypes.int, d_map(m, n)[0])).bitcast(dtypes.float32)
              for m in range(16) for n in range(16)]
     mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)]
-    stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), mat_d[m*16+n].bitcast(dtypes.uint32), exec_mask)
+    stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(dtypes.int, d_map(m, n)[0]), mat_d[m*16+n].bitcast(dtypes.uint32), exec_mask)
               for m in range(16) for n in range(16)]
   return UOp.sink(*stores, *ctx.inc_pc())
 
@@ -1688,8 +1710,8 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
       scalar_hi_sel = src_lo if not opsel_hi_bit else is_sgpr_pair.where(sgpr_hi, src_lo)
       lo = is_vgpr.where(vgpr_hi if opsel_lo else vgpr_lo, scalar_lo_sel)
       hi = is_vgpr.where(vgpr_hi if opsel_hi_bit else vgpr_lo, scalar_hi_sel)
-      if neg_lo: lo = lo ^ UOp.const(0x80000000, dtypes.uint32)
-      if neg_hi_bit: hi = hi ^ UOp.const(0x80000000, dtypes.uint32)
+      if neg_lo: lo = lo ^ UOp.const(dtypes.uint32, 0x80000000)
+      if neg_hi_bit: hi = hi ^ UOp.const(dtypes.uint32, 0x80000000)
       return _u64(lo, hi)
     srcs = {'S0': build_pk_f32(src0, src_offs[0], opsel & 1, opsel_hi & 1, neg & 1, neg_hi & 1),
             'S1': build_pk_f32(src1, src_offs[1], opsel & 2, opsel_hi & 2, neg & 2, neg_hi & 2),
@@ -1700,35 +1722,35 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
     def apply_abs(v, bit, opsel_hi_bit, opsel_bit):
       if not (neg_hi & bit): return v
       # Apply abs based on whether source is f32 or f16
-      if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(0x7FFFFFFF, dtypes.uint32)  # f32 abs
-      if opsel & opsel_bit: return v & UOp.const(0x7FFF0000, dtypes.uint32)  # f16 hi abs (preserve lo)
-      return v & UOp.const(0xFFFF7FFF, dtypes.uint32)  # f16 lo abs (preserve hi)
+      if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(dtypes.uint32, 0x7FFFFFFF)  # f32 abs
+      if opsel & opsel_bit: return v & UOp.const(dtypes.uint32, 0x7FFF0000)  # f16 hi abs (preserve lo)
+      return v & UOp.const(dtypes.uint32, 0xFFFF7FFF)  # f16 lo abs (preserve hi)
     def apply_neg_mix(v, bit, opsel_hi_bit, opsel_bit):
       if not (neg & bit): return v
-      if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(0x80000000, dtypes.uint32)  # f32 neg
-      if opsel & opsel_bit: return v ^ UOp.const(0x80000000, dtypes.uint32)  # f16 hi neg
-      return v ^ UOp.const(0x00008000, dtypes.uint32)  # f16 lo neg
+      if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(dtypes.uint32, 0x80000000)  # f32 neg
+      if opsel & opsel_bit: return v ^ UOp.const(dtypes.uint32, 0x80000000)  # f16 hi neg
+      return v ^ UOp.const(dtypes.uint32, 0x00008000)  # f16 lo neg
     s0_mod = apply_neg_mix(apply_abs(src0, 1, 1, 1), 1, 1, 1)
     s1_mod = apply_neg_mix(apply_abs(src1, 2, 2, 2), 2, 2, 2)
     s2_mod = apply_neg_mix(apply_abs(src2, 4, 4, 4), 4, 4, 4)
     srcs = {'S@0': s0_mod, 'S@1': s1_mod, 'S@2': s2_mod,
-            'OPSEL_HI': UOp.const(combined_opsel_hi, dtypes.uint32), 'OPSEL': UOp.const(opsel, dtypes.uint32)}
+            'OPSEL_HI': UOp.const(dtypes.uint32, combined_opsel_hi), 'OPSEL': UOp.const(dtypes.uint32, opsel)}
   else:
     def get_half_bits(val: UOp, use_hi: bool, apply_neg: bool = False) -> UOp:
-      bits = ((val >> UOp.const(16, dtypes.uint32)) if use_hi else val) & UOp.const(0xFFFF, dtypes.uint32)
+      bits = ((val >> UOp.const(dtypes.uint32, 16)) if use_hi else val) & UOp.const(dtypes.uint32, 0xFFFF)
       if apply_neg: bits = bits.cast(dtypes.uint16).bitcast(dtypes.half).neg().bitcast(dtypes.uint16).cast(dtypes.uint32)
       return bits
     def build_remapped_src(src: UOp, opsel_lo_bit: int, opsel_hi_bit: int, neg_lo_bit: int, neg_hi_bit: int) -> UOp:
       lo = get_half_bits(src, bool(opsel_lo_bit), bool(neg_lo_bit))
       hi = get_half_bits(src, bool(opsel_hi_bit), bool(neg_hi_bit))
-      return lo | (hi << UOp.const(16, dtypes.uint32))
+      return lo | (hi << UOp.const(dtypes.uint32, 16))
     # DOT IU instructions use NEG bits for signed/unsigned selection, not fp16 negation
     is_dot_iu = 'DOT' in op_name and 'IU' in op_name
     n0, n1, n2, nh0, nh1, nh2 = (0, 0, 0, 0, 0, 0) if is_dot_iu else (neg & 1, neg & 2, neg & 4, neg_hi & 1, neg_hi & 2, neg_hi & 4)
     srcs = {'S0': build_remapped_src(src0, opsel & 1, opsel_hi & 1, n0, nh0),
             'S1': build_remapped_src(src1, opsel & 2, opsel_hi & 2, n1, nh1),
             'S2': build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, n2, nh2)}
-    if is_dot_iu: srcs['NEG'] = UOp.const(neg, dtypes.uint32)
+    if is_dot_iu: srcs['NEG'] = UOp.const(dtypes.uint32, neg)
   return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask)
 
 def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp:
@@ -1781,7 +1803,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
   # CDNA acc bit: when set, VGPR operands (vdst/vdata) target ACCVGPR file instead of VGPR
   use_acc = bool(getattr(inst, 'acc', 0))
   mem = ctx.lds if is_lds else ctx.scratch if is_scratch else ctx.vmem
-  addr_shift = UOp.const(2, dtypes.uint32 if is_lds else dtypes.uint64)
+  addr_shift = UOp.const(dtypes.uint32 if is_lds else dtypes.uint64, 2)
 
   # Extract register info - all dynamic for deduplication
   if is_lds:
@@ -1831,20 +1853,20 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
       return addr
     offset64 = offset.cast(dtypes.uint64)
     # Dynamic saddr check: saddr < 124 means valid SGPR, otherwise use VGPR pair for address
-    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False)
+    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False)
     if is_scratch:
       scratch_stride = ctx.rsgpr_dyn(_c(SCRATCH_STRIDE_IDX)).cast(dtypes.uint64)
       base = lane.cast(dtypes.uint64) * scratch_stride
       # SVE (Scratch VGPR Enable): when SVE=1, VADDR is used as offset; when SVE=0, VADDR is ignored
       sve = getattr(inst, 'sve', 0)
       vaddr = ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64)
-      addr_offset = vaddr if sve == 1 else UOp.const(0, dtypes.uint64)
+      addr_offset = vaddr if sve == 1 else UOp.const(dtypes.uint64, 0)
       # Add saddr value only if use_saddr is true (saddr < 124)
-      saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(0, dtypes.uint64)) \
-        if saddr_reg is not None else UOp.const(0, dtypes.uint64)
+      saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(dtypes.uint64, 0)) \
+        if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
       return base + addr_offset + saddr_contrib + offset64
     # FLAT/GLOBAL: choose between SGPR base (saddr) or VGPR pair (addr) based on saddr validity
-    saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64)
+    saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
     vaddr_base = _u64(ctx.rvgpr_dyn(addr_reg, lane), ctx.rvgpr_dyn(addr_reg + _c(1), lane))
     # When saddr is valid: base = saddr pair, vaddr is 32-bit offset; otherwise: base = 0, vaddr is 64-bit address
     base_addr = use_saddr.where(saddr_base + ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64), vaddr_base)
@@ -1874,18 +1896,18 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
         data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA1': ctx.rvgpr_dyn(vdata_reg + _c(1), lane),
                 'DATA2': ctx.rvgpr_dyn(vdata_reg + _c(2), lane)}
       elif data_bits_mem <= 32:
-        data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA2': ctx.rvgpr_dyn(data1_reg, lane) if has_data1 else UOp.const(0, dtypes.uint32)}
+        data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA2': ctx.rvgpr_dyn(data1_reg, lane) if has_data1 else UOp.const(dtypes.uint32, 0)}
       else:
         data = {'DATA': _u64(ctx.rvgpr_dyn(vdata_reg, lane), ctx.rvgpr_dyn(vdata_reg + _c(1), lane)),
-                'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(0, dtypes.uint64)}
+                'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(dtypes.uint64, 0)}
       # RDNA3 uses ADDR/OFFSET, RDNA4 uses vgpr_a/offset (lowercase) + CalcDsAddr function
       return {'ADDR': addr, 'ADDR_BASE': addr, 'OFFSET': offset, 'OFFSET0': offset0, 'OFFSET1': offset1, '_lds': mem, 'laneId': lane,
               'vgpr_a': ctx.rvgpr_dyn(addr_reg, lane), 'offset': offset, 'offset0': offset0, 'offset1': offset1, **data}
     active = _lane_active(exec_mask, lane)
     # saddr < 124 means valid SGPR pair, otherwise use 0 (NULL means no saddr contribution)
-    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False)
-    saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64)
-    saddr_base = use_saddr.where(saddr_raw, UOp.const(0, dtypes.uint64))
+    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False)
+    saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
+    saddr_base = use_saddr.where(saddr_raw, UOp.const(dtypes.uint64, 0))
     # Sign-extend offset to 64-bit for the final address calculation
     ioffset64 = offset.cast(dtypes.int64).cast(dtypes.uint64)
     # v_addr for CalcGlobalAddr: when saddr valid, use low 32 bits as offset; otherwise full 64-bit address. Include ioffset.
@@ -1900,13 +1922,13 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
     # acc bit: read/write ACCVGPR instead of VGPR for data operands
     _rvdata = (lambda r, l, *a: ctx.raccvgpr_dyn(r, l)) if use_acc else ctx.rvgpr_dyn
     vdata = _rvdata(vdata_reg, lane).cast(dtypes.uint64) if 'STORE' in op_name \
-      else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(0, dtypes.uint32)
+      else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(dtypes.uint32, 0)
     if 'STORE' in op_name and data_bits_mem >= 64:
-      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
+      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
     srcs = {'ADDR': addr, 'VDATA': vdata, '_vmem': mem, '_active': active,
             'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base, 'SADDR': saddr_base, 'OFFSET': offset}
     for i in range(data_bits_mem // 32):
-      srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(0, dtypes.uint32)
+      srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(dtypes.uint32, 0)
     return srcs
 
   def make_stores(dest: str, val: UOp, lane: UOp, active: UOp, writes_return_data: bool) -> list[UOp]:
@@ -1983,7 +2005,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
   offset, offen, idxen = ctx.inst_field(type(inst).offset), ctx.inst_field(type(inst).offen), ctx.inst_field(type(inst).idxen)
 
   # V# descriptor: base[0:1], num_records[2], stride=word3[13:0]
-  base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(0xFFFFFFFFFFFF, dtypes.uint64)
+  base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(dtypes.uint64, 0xFFFFFFFFFFFF)
   num_records = ctx.rsgpr_dyn(srsrc + _c(2))
   stride = (ctx.rsgpr_dyn(srsrc + _c(3)) & _c(0x3FFF)).cast(dtypes.uint64)
 
@@ -2000,29 +2022,31 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
   buffer_offset = (stride * index + voff + offset.cast(dtypes.uint64)).cast(dtypes.uint32)
   in_bounds = active & buffer_offset.__lt__(num_records)
   addr = base + soff + buffer_offset.cast(dtypes.uint64)
-  addr = in_bounds.where(addr, UOp.const(0, dtypes.uint64))  # safe address when OOB
+  addr = in_bounds.where(addr, UOp.const(dtypes.uint64, 0))  # safe address when OOB
   mem = ctx.vmem
 
   stores: list[UOp] = []
   if is_lds and not is_store:
-    # LDS load: buffer -> LDS (bypass VGPRs), LDS addr = M0[17:0] + lane * elem_size
+    # LDS load: buffer -> LDS (bypass VGPRs), LDS addr = M0[17:0] + lane * elem_size.
+    # HW never takes a per-lane LDS address: kittens' direct fill sets M0 (s_mov_b32 m0, sN)
+    # before every lds instruction, giving lane-linear chunks with the swizzle on the GLOBAL side.
     lds_base = ctx.rsgpr_dyn(_c(124)) & _c(0x3FFFF)
     lds_addr = lds_base + lane.cast(dtypes.uint32) * _c(n_dwords * 4)
     for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
       val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64)).load(), _c(0))
       lds_idx = (lds_addr + _c(i * 4)) >> _c(2)
       lds_slot = ctx.lds.index(lds_idx.valid(active))
       stores.append(lds_slot.store(active.where(val, lds_slot)))
   elif is_store:
     for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
       idx = mem.index(word_addr.cast(dtypes.int64).valid(in_bounds))
       val = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(vdata + _c(i), lane)
       stores.append(idx.store(in_bounds.where(_to_u32(val), idx)))
   else:
     for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
       val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64).valid(in_bounds)).load(), _c(0))
       stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(vdata + _c(i), lane, val, exec_mask))
   return UOp.sink(UOp.group(*stores).end(lane), *ctx.inc_pc())
@@ -2114,6 +2138,8 @@ def _decode_at(pc: int, arch: str):
 F32_INLINE = {240: 0x3f000000, 241: 0xbf000000, 242: 0x3f800000, 243: 0xbf800000,  # 0.5, -0.5, 1.0, -1.0
               244: 0x40000000, 245: 0xc0000000, 246: 0x40800000, 247: 0xc0800000, 248: 0x3e22f983}  # 2.0, -2.0, 4.0, -4.0, 1/(2*pi)
 
+_inst_hist: dict = {}
+
 class WaveState:
   __slots__ = ('vgpr_buf', 'sgpr_buf', 'accvgpr_buf', '_vgpr_mv', '_sgpr_mv', 'n_lanes', 'wave_size')
 
@@ -2209,7 +2235,8 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
   # Use Buffer objects with external_ptr=0 for vmem
   vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
   lds_buf = Buffer('CPU', max(lds_size // 4, 1), dtypes.uint32).ensure_allocated()
-  scratch_buf = Buffer('CPU', scratch_size * wave_size, dtypes.uint8).ensure_allocated() if scratch_size else None
+  ctypes.memset(lds_buf._buf.va_addr, 0, max(lds_size, 4))
+  # NOTE: scratch (private/spill) memory is per-wavefront; buffers are allocated in the wave loop below
 
   # Initialize SQTT encoder — emits packets inline as instructions execute (only when profiling)
   if PROFILE:
@@ -2227,6 +2254,8 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
         print(colored(msg, 'green') if len(_canonical_runner_cache) > prev_len else msg)
     return program[pc]
 
+  if os.getenv("EMU_TRACE_INST") or os.getenv("EMU_WATCH_PC") is not None:
+    print(f"[emu-dispatch] gx={gx} gy={gy} gz={gz} lx={lx} ly={ly} lz={lz} scratch={scratch_size}", flush=True)
   # Set DAZ+FTZ during emulator execution, restore afterward to avoid breaking hypothesis tests
   # Only trace the first workgroup (like real HW traces one CU/SIMD), subsequent workgroups run but don't add to trace
   tracing = bool(PROFILE)
@@ -2237,18 +2266,27 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
         for gidx in range(gx):
           # Initialize all wavefronts for this workgroup
           waves: list[tuple[WaveState, list]] = []
+          wave_scratch_bufs: list[Buffer] = []   # keep alive for the dispatch
           for wave_start in range(0, total_threads, wave_size):
             st = _init_wave(lib, wave_start, total_threads, lx, ly, lz, args_ptr, rsrc2, scratch_size, arch, gidx, gidy, gidz, user_data,
                             wave_size)
+            # each wavefront owns its private (spill) scratch segment: on real HW the ring is indexed by
+            # (wave_id, lane), and scratch addresses here are lane*stride within the wave's segment.
+            if scratch_size:
+              wave_scratch_bufs.append(Buffer('CPU', scratch_size * wave_size, dtypes.uint8).ensure_allocated())
+              ctypes.memset(wave_scratch_bufs[-1]._buf.va_addr, 0, scratch_size * wave_size)
             c_bufs = [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
                       ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(lds_buf._buf.va_addr),
-                      ctypes.c_uint64(scratch_buf._buf.va_addr if scratch_buf else 0),
+                      ctypes.c_uint64(wave_scratch_bufs[-1]._buf.va_addr if scratch_size else 0),
                       ctypes.c_uint64(st.accvgpr_buf._buf.va_addr)]
             waves.append((st, c_bufs))
 
           # Execute wavefronts with barrier synchronization
           # Each wave runs until it hits s_barrier or s_endpgm. When all waves have stopped, release barrier waves.
           done = [False] * len(waves)
+          _exec_dumped: set = set()
+          _trace_on = bool(os.getenv("EMU_TRACE_INST")) and total_threads >= int(os.getenv("EMU_TRACE_MIN_THREADS", "1"))
+          if _trace_on: _inst_hist.clear()
           for total_inst in range(10_000_000):
             if all(done): break
             for wi, (st, c_bufs) in enumerate(waves):
@@ -2259,9 +2297,23 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
                 if pc == ENDPGM_PC:
                   done[wi] = True
                   if tracing: sqtt_finish(wi)
+                  if os.getenv("EMU_EXEC_DUMP") and wi not in _exec_dumped:
+                    _exec_dumped.add(wi)
+                    print(f"[exec-dump] gid=({gidx},{gidy},{gidz}) wave{wi} "
+                          f"exec_lo={st._read_sgpr(EXEC_LO.offset):08x} exec_hi={st._read_sgpr(EXEC_LO.offset+1):08x}")
                   break
                 fxn, globals_list, is_barrier, inst = _ensure_compiled(pc)
                 if DEBUG >= 5: print(f"  exec gid=({gidx},{gidy},{gidz}) w={wi} PC={pc - lib}: {inst!r}", flush=True)
+                if _trace_on:
+                  key = (gidx, gidy, gidz, wi)
+                  _inst_hist.setdefault(key, []).append((pc - lib, type(inst).__name__, getattr(inst, 'op', None) and inst.op.name))
+                wpc, wwave = os.getenv("EMU_WATCH_PC"), int(os.getenv("EMU_WATCH_WAVE", "0"))
+                if wpc is not None and (int(wpc) < 0 or (pc - lib) == int(wpc)) and wi == wwave and gidx == gidy == gidz == 0:
+                  if int(wpc) < 0 and total_inst < 800: print(f"[watch-stream] pc={pc-lib} {inst!r}", flush=True)
+                  for rv in os.getenv("EMU_WATCH_VGPR", "").split(","):
+                    if rv: print(f"[watch pc={pc-lib}] wave{wi} {rv}:", [st._read_vgpr(int(rv[1:]), l) for l in range(8)], flush=True)
+                  for rv in os.getenv("EMU_WATCH_SGPR", "").split(","):
+                    if rv: print(f"[watch pc={pc-lib}] wave{wi} {rv}:", [st._read_sgpr(int(rv[1:]))], flush=True)
                 fxn(*[c_bufs[g] for g in globals_list])
                 if tracing:
                   inst_op = inst.op.value if hasattr(inst, 'op') else 0
@@ -2269,9 +2321,26 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
                 if is_barrier: break  # s_barrier hit: PC already advanced past it, pause this wave
               else: raise RuntimeError("exceeded 1M instructions in single wave, likely infinite loop")
             # All waves have either hit barrier or endpgm — release barrier waves for next round
+            if os.getenv("EMU_DUMP_ROUNDS") and lds_size > 0 and total_inst < int(os.getenv("EMU_DUMP_ROUNDS")):
+              import numpy as _np
+              lds_words = _np.frombuffer((ctypes.c_uint32 * (lds_size//4)).from_address(lds_buf._buf.va_addr), dtype=_np.uint32)
+              nz = _np.argwhere(lds_words != 0)
+              print(f"[emu-dump] round={total_inst} lds nonzero words={len(nz)}",
+                    (f"first16={[f'w{w}:0x{lds_words[w]:08x}' for w in nz[:16].flatten()]}" if len(nz) else ""), flush=True)
+              for wi2, (st2, _) in enumerate(waves):
+                v = _np.frombuffer((ctypes.c_uint32 * (256*st2.wave_size)).from_address(st2.vgpr_buf._buf.va_addr), dtype=_np.uint32)
+                if st2.wave_size == 64:
+                  av = _np.frombuffer((ctypes.c_uint32 * (256*st2.wave_size)).from_address(st2.accvgpr_buf._buf.va_addr), dtype=_np.uint32)
+                else: av = _np.zeros(1, dtype=_np.uint32)
+                print(f"[emu-dump]   wave{wi2} vgpr nonzero={int((v!=0).sum())} accvgpr nonzero={int((av!=0).sum())}", flush=True)
           else: raise RuntimeError("exceeded 10M total scheduling rounds")
           tracing = False  # only trace the first workgroup
 
+          if _trace_on:
+            import pickle
+            tag = os.environ["EMU_TRACE_INST"]
+            with open(f"/tmp/emu_trace_{tag}.pkl", "wb") as f: pickle.dump(dict(_inst_hist), f)
+            os.environ.pop("EMU_TRACE_INST", None)   # only dump for the first matching dispatch
           # Reset LDS for next workgroup
           if lds_size > 0: ctypes.memset(lds_buf._buf.va_addr, 0, max(lds_size, 4))
 

test/mockgpu/amd/pcode.py

@@ -7,7 +7,7 @@ from tinygrad.codegen.decomp.dtype import f2f
 # Type alias for vars dict: stores UOps and tuples for lambda definitions
 VarVal = UOp | tuple[str, list[str], str]
 
-def _const(dt, v): return UOp.const(v, dt)
+def _const(dt, v): return UOp.const(dt, v)
 def _u32(v): return _const(dtypes.uint32, v)
 def _u64(v): return _const(dtypes.uint64, v)
 def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32)
@@ -55,8 +55,8 @@ def _expr_bits(v: UOp) -> int:
   if v.op in (Ops.AND, Ops.XOR):
     widths: list[int] = []
     for src in v.src:
-      if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0:
-        widths.append(src.val.bit_length())
+      if src.op == Ops.CONST and isinstance(src.arg, int) and src.arg > 0 and (src.arg & (src.arg + 1)) == 0:
+        widths.append(src.arg.bit_length())
     if widths: return max(widths)
   return v.dtype.bitsize
 
@@ -144,9 +144,9 @@ def _minmax_reduce(is_max: bool, dt, *args: UOp) -> UOp:
 def _find_two_pi_mul(x):
   if x.op != Ops.MUL or len(x.src) != 2: return None
   for i, s in enumerate(x.src):
-    if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
+    if s.op == Ops.CONST and abs(s.arg - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
     if s.op == Ops.MUL and len(s.src) == 2:
-      vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
+      vals = [ss.arg for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].arg for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
       if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1])
   return None
 
@@ -163,7 +163,7 @@ def _trig_reduce(x, phase=0.0):
 
 def _signext(val: UOp) -> UOp:
   for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
-    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8:
+    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].arg == mask) or val.dtype.itemsize == bits // 8:
       v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
       sb = (v32 >> _u32(bits - 1)) & _u32(1)
       return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
@@ -497,7 +497,7 @@ class Parser:
         if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32)
         return (left >> right) if op == '>>' else (left << right)
       case '+' | '-':
-        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val)
+        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.arg - right.arg)
         return (left + right) if op == '+' else (left - right)
       case '*' | '/':
         # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
@@ -507,7 +507,7 @@ class Parser:
           left, right = left.cast(pdt), right.cast(pdt)
         if op == '*': return left * right
         return (left // right) if dtypes.is_int(left.dtype) else (left / right)
-      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left
+      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.arg == 2.0 else left
 
   _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]
 
@@ -530,7 +530,7 @@ class Parser:
     if self.try_eat_val('-', 'OP'):
       inner = self.unary()
       if inner.op == Ops.CONST:
-        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val)
+        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.arg)
       return inner.neg()
     if self.try_eat_val('+', 'OP'): return self.unary()
     return self.postfix()
@@ -670,14 +670,14 @@ class Parser:
       width = self.parse()
       self.eat('RBRACKET')
       if width.op == Ops.CONST:
-        w = int(width.val)
+        w = int(width.arg)
         return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1)
       return base
     if self.try_eat('COLON'):
       second = self.parse()
       self.eat('RBRACKET')
       if first.op == Ops.CONST and second.op == Ops.CONST:
-        a, b = int(first.val), int(second.val)
+        a, b = int(first.arg), int(second.arg)
         if a < b: return _bitreverse(base, b - a + 1)
         hi, lo = a, b
         if lo >= base.dtype.itemsize * 8:
@@ -699,7 +699,7 @@ class Parser:
     if var_name is None:
       var_name = self._find_var_name(base)
     if first.op == Ops.CONST:
-      idx = int(first.val)
+      idx = int(first.arg)
       # Check for array element (var@idx)
       if var_name and f'{var_name}@{idx}' in self.vars:
         v = self.vars[f'{var_name}@{idx}']
@@ -866,13 +866,13 @@ class Parser:
         idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64)
         safe_idx_hi = is_unaligned.where(idx_hi_native, idx_native)
         hi = mindex(safe_idx_hi)
-        combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
-        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(8, dtypes.uint64))).cast(dtypes.uint32), val)
+        combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
+        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(dtypes.uint64, 8))).cast(dtypes.uint32), val)
     return _cast_to(val, dt)
 
   def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
     if l.dtype != r.dtype:
-      if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int)
+      if r.dtype == dtypes.int and r.op == Ops.CONST and r.arg < 0: l = l.cast(dtypes.int)
       else: r = r.cast(l.dtype)
     return l, r
 
@@ -970,7 +970,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
         expr = p.parse().simplify()
         assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
-        return int(expr.val)
+        return int(expr.arg)
       start_val = parse_bound()
       p.eat('COLON')
       end_val = parse_bound()
@@ -1258,7 +1258,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       def parse_cond(s, kw):
         ll = s.lower()
         return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), env, funcs))
-      def is_const(c, v): return c.op == Ops.CONST and c.val is v
+      def is_const(c, v): return c.op == Ops.CONST and c.arg is v
       cond = parse_cond(line, 'if')
       conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
       branch_assigns: list[tuple[UOp, list]] = []  # (cond, assigns_list) for side-effect merging
@@ -1339,7 +1339,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
               return (dest, (val[0], cnd.where(val[1], val[1])))
             return (dest, val)
           # Build combined condition: each branch fires when its cond is true AND no earlier cond was true
-          remaining = UOp.const(True)
+          remaining = UOp.const(dtypes.bool, True)
           for bc, bse in branch_assigns:
             effective = remaining & bc if remaining.op != Ops.CONST else bc
             for dest, val in bse: assigns.append(_cond_side_effect(effective, dest, val))

test/null/test_const_folding.py

@@ -1,6 +1,6 @@
 import unittest, itertools, math
 from tinygrad import Tensor, dtypes, Context
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType, ConstType, Invalid
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -36,20 +36,21 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
 
 class TestWeakConstFolding(unittest.TestCase):
   def test_weakint_math(self):
-    out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
-    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakint, 2**41))
+    out = (UOp.const(dtypes.weakint, 2**40) + UOp.const(dtypes.weakint, 2**40)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41))
 
   def test_float_unaries(self):
-    for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
-      out = UOp.const(4.0).alu(op).simplify()
-      self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
+    for dtype in (dtypes.weakfloat,):
+      for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
+        out = UOp.const(dtype, 4).alu(op).simplify()
+        self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
 
   def test_weakfloat_math(self):
-    out = (UOp.const(1.25) + UOp.const(2.5)).simplify()
-    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakfloat, 3.75))
+    out = (UOp.const(dtypes.weakfloat, 1.25) + UOp.const(dtypes.weakfloat, 2.5)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75))
 
   def test_invalid_poison(self):
-    self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
+    self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid)
 
 class TestBinaryOpsConstFolding(unittest.TestCase):
   def test_add_literal_zero(self):
@@ -120,10 +121,10 @@ class TestBitcastConstFolding(unittest.TestCase):
     def t(cases: dict[DType, ConstType]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
-          r = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0]
+          r = full_rewrite(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
           self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
           self.assertEqual(r.dtype, to_dt, msg)
-          np.testing.assert_equal(r.val, to_v, msg)
+          np.testing.assert_equal(r.arg, to_v, msg)
 
     t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
     t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})
@@ -144,9 +145,9 @@ class TestBitcastConstFolding(unittest.TestCase):
 
   def test_vec_bitcast(self):
     with Context(SPEC=0):
-      srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src
+      srcs = full_rewrite(UOp.const(dtypes.int32, (-1, -2**31, 75)).bitcast(dtypes.uint32).sink()).src
     self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
-    self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75))
+    self.assertEqual(tuple(x.arg for x in srcs), (2**32-1, 2**31, 75))
 
 # folds advance indexing into basic indexing
 class TestIndexingConstFolding(unittest.TestCase):

test/null/test_gpudims.py

@@ -12,7 +12,7 @@ class TestGroupedDims(unittest.TestCase):
     idxs = get_grouped_dims(prefix, dims, max_sizes, reverse)
     loop_idxs = dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs]))
     loop_idxs = sorted(loop_idxs, key=lambda uop: uop.arg)
-    sizes = [x.src[0].val for x in loop_idxs]
+    sizes = [x.src[0].arg for x in loop_idxs]
     assert len(idxs) == len(dims), f"expected idxs to have same length as dims {len(dims)}, got {len(idxs)}"
     if assert_same_length:
       assert len(loop_idxs) == min(len(sizes), len(dims)), f"expected idxs to have length {min(len(sizes), len(dims))}, got {len(loop_idxs)}"
@@ -24,7 +24,7 @@ class TestGroupedDims(unittest.TestCase):
     total = math.prod(dims)
     specials = sorted(dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])), key=lambda u: u.arg)
     # build flat index and primed flat (same expression with renamed SPECIALs)
-    flat = UOp.const(0)
+    flat = UOp.const(dtypes.weakint, 0)
     for i, idx in enumerate(idxs):
       flat = flat + idx * int(math.prod(dims[i+1:]))
     flat_p = flat.substitute({s: UOp(Ops.SPECIAL, src=s.src, arg=s.arg+"_p") for s in specials})
@@ -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(dtypes.float, 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_gradient.py

@@ -14,10 +14,10 @@ class TestGradient(unittest.TestCase):
   def _test_one_input_function(self, f:Callable, jf:Callable|None=None):
     if jf is None: jf = f
     x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
-    gx = compute_gradient(f(x), UOp.const(1.0), set([x]))[x]
+    gx = compute_gradient(f(x), UOp.const(dtypes.float, 1.0), set([x]))[x]
 
     for val in [-5., -2.0, 0.0, 2.0, 5.]:
-      tg_out = gx.substitute({x: UOp.const(val)}).ssimplify()
+      tg_out = gx.substitute({x: x.const_like(val)}).ssimplify()
       tx = torch.tensor([val], dtype=torch.float, requires_grad=True)
       torch_out = torch.autograd.grad(jf(tx), tx)[0].item()
       self._cmp_nan_okay(tg_out, torch_out)
@@ -26,13 +26,13 @@ class TestGradient(unittest.TestCase):
     if jf is None: jf = f
     x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
     y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float)
-    grads = compute_gradient(f(x, y), UOp.const(1.0), set([x, y]))
+    grads = compute_gradient(f(x, y), UOp.const(dtypes.float, 1.0), set([x, y]))
     gx, gy = grads[x], grads[y]
 
     for valx in [-5., -2.0, 0.0, 2.0, 5.]:
       for valy in [-5., -2.0, 0.0, 2.0, 5.]:
         # Substitute the values into the gradient expressions
-        substitutions = {x: UOp.const(valx), y: UOp.const(valy)}
+        substitutions = {x: x.const_like(valx), y: y.const_like(valy)}
         tg_out_x = gx.substitute(substitutions).ssimplify()
         tg_out_y = gy.substitute(substitutions).ssimplify()
 

test/null/test_graph_rewrite.py

@@ -15,13 +15,13 @@ def apply_rewrite(expr):
 def apply_rewrite_values(expr):
   srcs = full_rewrite(expr.sink()).src
   if len(srcs) == 1:
-    if srcs[0].op is Ops.CONST: return (srcs[0].val,)
-    if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src)
-  return tuple(s.val for s in srcs)
+    if srcs[0].op is Ops.CONST: return (srcs[0].arg,) if not isinstance(srcs[0].arg, tuple) else srcs[0].arg
+    if srcs[0].op is Ops.STACK: return tuple(s.arg for s in srcs[0].src)
+  return tuple(s.arg for s in srcs)
 
 def evaluate_uop(uop, variables):
   if uop.op == Ops.CONST:
-    return uop.val
+    return uop.arg
   elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU:
     return variables[uop.expr]
   elif uop.op in GroupOp.ALU:
@@ -32,67 +32,67 @@ def evaluate_uop(uop, variables):
 
 class TestArithmeticSimplifications(unittest.TestCase):
   def test_full_graph_rewrite_division_by_zero(self):
-    optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0))
+    optimized_div_uop = apply_rewrite(UOp.const(dtypes.float32, 10.0) / UOp.const(dtypes.float32, 0.0))
     self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val))
+    self.assertTrue(math.isinf(optimized_div_uop.arg) or math.isnan(optimized_div_uop.arg))
 
   def test_full_graph_rewrite_redundant_operations(self):
-    optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0))
+    optimized_uop = apply_rewrite((UOp.const(dtypes.float32, 10.0) + UOp.const(dtypes.float32, 0.0)) * UOp.const(dtypes.float32, 1.0))
     self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 10.0)
+    self.assertEqual(optimized_uop.arg, 10.0)
 
   def test_full_graph_rewrite_large_graph(self):
-    prev_uop = UOp.const(0)
+    prev_uop = UOp.const(dtypes.int32, 0)
     for i in range(1, 101):
-      prev_uop += UOp.const(i)
+      prev_uop += UOp.const(dtypes.int32, i)
     optimized_uop = apply_rewrite(prev_uop)
     self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, sum(range(1, 101)))
+    self.assertEqual(optimized_uop.arg, sum(range(1, 101)))
 
   def test_full_graph_rewrite_division_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0))
+    optimized_uop = apply_rewrite(UOp.const(dtypes.float32, 42.0) / UOp.const(dtypes.float32, 1.0))
     self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 42.0)
+    self.assertEqual(optimized_uop.arg, 42.0)
 
   def test_full_graph_rewrite_modulo_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
+    optimized_uop = apply_rewrite(UOp.const(dtypes.int32, 42) % UOp.const(dtypes.int32, 1))
     self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 0)
+    self.assertEqual(optimized_uop.arg, 0)
 
 
 class TestFoldingAndReduction(unittest.TestCase):
   @unittest.skip("reduce is removed now")
   def test_full_graph_rewrite_constant_reduction_folding(self):
-    const1 = UOp.const(5)
-    const2 = UOp.const(10)
-    const3 = UOp.const(20)
+    const1 = UOp.const(dtypes.int32, 5)
+    const2 = UOp.const(dtypes.int32, 10)
+    const3 = UOp.const(dtypes.int32, 20)
     optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
     expected_sum = 5 + 10 + 20
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)
 
   @unittest.skip("reduce is removed now")
   def test_full_graph_rewrite_reduction_with_unused_range(self):
-    const1 = UOp.const(15)
-    const2 = UOp.const(25)
+    const1 = UOp.const(dtypes.int32, 15)
+    const2 = UOp.const(dtypes.int32, 25)
     rng = UOp.range(10, idx=0)
     optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
     expected_sum = 10 * (15 + 25)
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_range_reduction(self):
     simple_range = UOp.range(5, idx=0)
     optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
     expected_sum = sum(range(5))
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_simple_reduction_folding(self):
     simple_range = UOp.range(4, idx=0)
-    add_uop = simple_range + UOp.const(1)
+    add_uop = simple_range + UOp.const(dtypes.int32, 1)
     optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
     expected_sum = sum(i + 1 for i in range(4))
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_nested_loop_collapse(self):
@@ -101,7 +101,7 @@ class TestFoldingAndReduction(unittest.TestCase):
     expr = (outer_range * 10) + inner_range
     optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
     self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
-    self.assertEqual(optimized_reduce_uop.val, sum((i * 10) + j for i in range(8) for j in range(4)))
+    self.assertEqual(optimized_reduce_uop.arg, sum((i * 10) + j for i in range(8) for j in range(4)))
 
 
 class TestModuloAndDivisionFolding(unittest.TestCase):
@@ -110,27 +110,27 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
     x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
     optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
     self.assertEqual(optimized_mod_uop.op, Ops.CONST)
-    self.assertEqual(optimized_mod_uop.val, 2)
+    self.assertEqual(optimized_mod_uop.arg, 2)
 
   def test_full_graph_rewrite_division_folding_with_define_var(self):
     # index dtype because div-mod rules only work on index
     n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
     optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
     self.assertEqual(optimized_div_uop.op, Ops.MUL)
-    self.assertEqual(optimized_div_uop.src[1].val, 2)
+    self.assertEqual(optimized_div_uop.src[1].arg, 2)
 
   def test_full_graph_rewrite_complex_mod_div_folding(self):
     # index dtype because div-mod rules only work on index
     k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
     optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
     self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertEqual(optimized_div_uop.val, 1)
+    self.assertEqual(optimized_div_uop.arg, 1)
 
   def test_graph_rewrite_div_folding_bug(self):
     lhs = UOp(Ops.ADD, src=(
-      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4),
-      UOp.const((0, 256, 512, 768))))
-    rhs = UOp.const((2,)*4)
+      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 32),), arg='lidx0'),)*4),
+      UOp.const(dtypes.int, (0, 256, 512, 768))))
+    rhs = UOp.const(dtypes.int, (2,)*4)
     unopt = lhs<rhs
     opt = apply_rewrite(unopt)
     print(unopt)
@@ -159,11 +159,11 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
 
 class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
   def test_full_graph_rewrite_transcendental_edge_cases(self):
-    optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal()))
+    optimized_sink = full_rewrite(UOp.const(dtypes.float32, -1.0).log2().sink(UOp.const(dtypes.float32, 0.0).reciprocal()))
     optimized_log2_neg, optimized_recip_zero = optimized_sink.src
-    self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}")
-    self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0,
-                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}")
+    self.assertTrue(math.isnan(optimized_log2_neg.arg), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.arg}")
+    self.assertTrue(math.isinf(optimized_recip_zero.arg) and optimized_recip_zero.arg > 0,
+                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.arg}")
 
   @unittest.skip("broken")
   def test_full_graph_rewrite_modulo_negative_dividend(self):
@@ -182,27 +182,27 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
 class TestGEPAndVectorizeRewrite(unittest.TestCase):
   def test_gep_single_element_extraction(self):
     # GEP on a vector dtype to extract a single element
-    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0)
+    base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0)
 
   def test_gep_tuple_extraction(self):
     # GEP on a vector dtype to extract multiple elements as a vector
-    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
+    base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
     self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
 
   def test_gep_on_const_stack(self):
     # GEP on a const STACK to extract a single element
-    const_stack = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(const_stack.index(2)).val, 3.0)
+    const_stack = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0)
 
   def test_gep_tuple_on_const_stack(self):
     # GEP on a const STACK using a tuple to extract multiple elements
-    const_stack = UOp.const((7.0, 8.0, 9.0, 10.0))
+    const_stack = UOp.const(dtypes.float32, (7.0, 8.0, 9.0, 10.0))
     self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
 
   def test_vectorize_multiple_elements(self):
     # Vectorizing multiple elements using GEP
-    base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
+    base_vector = UOp.const(dtypes.float32, (5.0, 10.0, 15.0, 20.0))
     vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4)))
     self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])
 
@@ -213,7 +213,7 @@ from tinygrad.uop.symbolic import symbolic_simple
 
 class TestBottomUpRewrite(unittest.TestCase):
   def test_const_folding(self):
-    a = UOp.const(5)
+    a = UOp.const(dtypes.int, 5)
     ret = (a*3) + (a*7)
     gt = graph_rewrite(ret, symbolic_simple)
     ret = graph_rewrite(ret, symbolic_simple, bottom_up=True)
@@ -305,7 +305,7 @@ class TestRecurse(unittest.TestCase):
     graph_rewrite(a, pm, bottom_up=True)
 
   def test_inf_loop(self):
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
     pm = PatternMatcher([
       (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
       (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
@@ -314,7 +314,7 @@ class TestRecurse(unittest.TestCase):
       graph_rewrite(a, pm)
 
   def test_inf_loop_bottom_up(self):
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
     pm = PatternMatcher([
       (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
       (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
@@ -322,11 +322,11 @@ class TestRecurse(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       graph_rewrite(a, pm, bottom_up=True)
 
-def bidir_append(ctx, x, b): ctx.append((x.val if x.op is Ops.CONST else "+", b))
+def bidir_append(ctx, x, b): ctx.append((x.arg if x.op is Ops.CONST else "+", b))
 class TestBidirectional(unittest.TestCase):
   def test_simple(self):
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
     c = a + b
     pm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, False)) ])
     bpm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, True)) ])
@@ -336,14 +336,14 @@ class TestBidirectional(unittest.TestCase):
 
 class TestStopEarly(unittest.TestCase):
   def test_stop_early(self):
-    a = UOp.const(3)
-    b = UOp.const(4)
+    a = UOp.const(dtypes.int, 3)
+    b = UOp.const(dtypes.int, 4)
     c = a+b
-    cn = UOp.const(7)
-    d = UOp.const(2)
+    cn = UOp.const(dtypes.int, 7)
+    d = UOp.const(dtypes.int, 2)
     def visit_const(c:UOp):
-      print(f"visit {c.val}")
-      assert c.val not in (3,4)
+      print(f"visit {c.arg}")
+      assert c.arg not in (3,4)
     pm_cvisit = PatternMatcher([(UPat(Ops.CONST, name="c"), visit_const),])
     ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
     assert ret == cn+d
@@ -376,7 +376,7 @@ class TestWalkRewrite(unittest.TestCase):
 
   def test_walk_topdown_no_fixed_point(self):
     """A bouncing pattern applies once and stops instead of looping."""
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
     pm = PatternMatcher([
       (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
       (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
@@ -384,7 +384,7 @@ class TestWalkRewrite(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       graph_rewrite(a, pm, bottom_up=True)
     ret = graph_rewrite(a, pm, walk=True)
-    self.assertIs(ret, UOp.const(4))
+    self.assertIs(ret, UOp.const(dtypes.int, 4))
 
   def test_walk_topdown_rewrites_children(self):
     a = UOp.variable('a', 0, 10)
@@ -418,11 +418,11 @@ class TestWalkRewrite(unittest.TestCase):
     """Top-down walk fires pm after children are processed (post-order)."""
     visited = []
     def track_visit(ctx, x):
-      ctx.append(x.val if x.op is Ops.CONST else x.op)
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
       return None
     pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
     graph_rewrite(a + b, pm, ctx=visited, walk=True)
     self.assertEqual(visited, [1, 2, Ops.ADD])
 
@@ -454,23 +454,23 @@ class TestWalkRewrite(unittest.TestCase):
 
   def test_walk_bottomup_no_fixed_point(self):
     """Bottom-up walk also applies once per node, no fixed-point iteration."""
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
     pm = PatternMatcher([
       (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
       (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
     ])
     ret = graph_rewrite(a, pm, bottom_up=True, walk=True)
-    self.assertIs(ret, UOp.const(4))
+    self.assertIs(ret, UOp.const(dtypes.int, 4))
 
   def test_walk_bottomup_visit_order(self):
     """Bottom-up walk fires bpm before descending (pre-order)."""
     visited = []
     def track_visit(ctx, x):
-      ctx.append(x.val if x.op is Ops.CONST else x.op)
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
       return None
     pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
     graph_rewrite(a + b, pm, ctx=visited, bottom_up=True, walk=True)
     # bpm fires on each node before children: +, 1, 2
     self.assertEqual(visited, [Ops.ADD, 1, 2])
@@ -490,15 +490,15 @@ class TestWalkRewrite(unittest.TestCase):
     """Bidirectional walk: bpm fires pre-order, pm fires post-order."""
     visited = []
     def bpm_visit(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
       return None
     def pm_visit(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "pm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
       return None
     bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)])
     pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
     graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
     # bpm fires pre-order, pm fires post-order
     self.assertEqual(visited, [
@@ -509,23 +509,23 @@ class TestWalkRewrite(unittest.TestCase):
     """If bpm matches, children are skipped and pm never fires on that node."""
     visited = []
     def bpm_match(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
       # rewrite const(1) -> const(10), short-circuiting its subtree
-      if x.op is Ops.CONST and x.val == 1: return x.replace(arg=10)
+      if x.op is Ops.CONST and x.arg == 1: return x.replace(arg=10)
       return None
     def pm_match(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "pm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
       return None
     bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)])
     pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
     ret = graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
     # bpm matches const(1) and short-circuits it, so pm never fires on const(1)
     self.assertNotIn((1, "pm"), visited)
     # but pm still fires on const(2) and the rebuilt ADD
     self.assertIn((2, "pm"), visited)
-    self.assertIs(ret, UOp.const(10) + b)
+    self.assertIs(ret, UOp.const(dtypes.int, 10) + b)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_helpers.py

@@ -297,10 +297,10 @@ class TestPolyN(unittest.TestCase):
     from tinygrad.dtype import dtypes
     from tinygrad.uop.ops import UOp
     from test.helpers import eval_uop
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(1.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 0.0)
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(2.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 1.0)
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(3.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 4.0)
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(4.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 9.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 1.0), [1.0, -2.0, 1.0])), 0.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 2.0), [1.0, -2.0, 1.0])), 1.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 3.0), [1.0, -2.0, 1.0])), 4.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 4.0), [1.0, -2.0, 1.0])), 9.0)
 
 class TestTimeToStr(unittest.TestCase):
   def test_seconds(self):           self.assertEqual("   10.01s ", time_to_str(10.01))

test/null/test_linearizer_failures.py

@@ -8,15 +8,15 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
     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)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.WEAK)
+    c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
     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))))
+    c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)
+    c6 = c4.index(((((((c5//UOp.const(dtypes.weakint, 8))%UOp.const(dtypes.weakint, 8))*UOp.const(dtypes.weakint, 8))+(c5%UOp.const(dtypes.weakint, 8)))+(((c2*UOp.const(dtypes.weakint, 40))+(c5//UOp.const(dtypes.weakint, 64)))*UOp.const(dtypes.weakint, 64)))+(c1*UOp.const(dtypes.weakint, 81920))))
     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()
+    c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)
     ast = c10.sink(arg=KernelInfo())
     to_program(ast, renderer=Device[Device.DEFAULT].renderer)

test/null/test_memory_planner.py

@@ -42,7 +42,7 @@ def check_assign(buffer_lists, copies=None):
   for orig_si, new_si in zip(linear.src, result.src):
     for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
       if new.op is Ops.SLICE and id(orig) not in replace_map:
-        replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
+        replace_map[id(orig)] = (new.src[0], new.src[1].arg * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
 
   # verify pinned buffers are not planned
   for buf in held_bufs:

test/null/test_microbenchmarks.py

@@ -1,5 +1,5 @@
 import unittest, time
-from tinygrad import Tensor, UOp, getenv
+from tinygrad import dtypes, Tensor, UOp, getenv
 from tinygrad.helpers import Profiling
 
 PYPROFILE = getenv("PYPROFILE")
@@ -27,29 +27,29 @@ class TestBench(unittest.TestCase):
     print(f"{self._testMethodName:30s} {et*1e6/self.N:.2f} us")
 
   def test_uop_instant_creation(self):
-    for i in range(self.N): UOp.const(100+i)
+    for i in range(self.N): UOp.const(dtypes.int, 100+i)
 
   def test_uop_list_creation(self):
-    [UOp.const(100+i) for i in range(self.N)]
+    [UOp.const(dtypes.int, 100+i) for i in range(self.N)]
 
   def test_uop_add_2n(self):
-    a = UOp.const(2)
+    a = UOp.const(dtypes.int, 2)
     for _ in range(self.N): a = a + a
 
   def test_uop_toposort(self):
-    a = UOp.const(0)
-    for i in range(self.N): a = a + UOp.const(100+i)
+    a = UOp.const(dtypes.int, 0)
+    for i in range(self.N): a = a + UOp.const(dtypes.int, 100+i)
     self.start_time()
     self.assertEqual(len(a.toposort()), 2*self.N+1)
 
   def test_uop_toposort_2n(self):
-    a = UOp.const(0)
+    a = UOp.const(dtypes.int, 0)
     for _ in range(self.N): a = a + a
     self.start_time()
     self.assertEqual(len(a.toposort()), self.N+1)
 
   def test_uop_simplify(self):
-    a = UOp.const(2)
+    a = UOp.const(dtypes.int, 2)
     for _ in range(self.N): (a+a).simplify()
 
   def test_uop_simplify_complex(self):
@@ -68,7 +68,7 @@ class TestBench(unittest.TestCase):
     for _ in range(self.N): expr.simplify()
 
   def test_uop_chain_free(self):
-    a = UOp.const(2)
+    a = UOp.const(dtypes.int, 2)
     for _ in range(self.N): a = a + a
     self.start_time()
     del a

test/null/test_pattern_matcher.py

@@ -5,9 +5,9 @@ from tinygrad.uop.ops import PatternMatcher, UPat
 
 class TestPatternMatcher(unittest.TestCase):
   def test_simple_match(self):
-    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.weakfloat), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(1)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.int, 1)
     self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
 
@@ -61,16 +61,16 @@ class TestPatternMatcher(unittest.TestCase):
 
   def test_uop(self):
     matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
+    c1 = UOp.const(dtypes.float, 1.0)
     c2 = UOp(Ops.ADD, src=(c1, c1))
     self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
 
   def test_uop_set(self):
     matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())])
-    c1 = UOp.const(False)
+    c1 = UOp.const(dtypes.bool, False)
     c2 = UOp(Ops.CAST, arg=dtypes.int, src=(c1,))
-    c3 = UOp.const(1.0)
+    c3 = UOp.const(dtypes.float, 1.0)
     c4 = UOp(Ops.ADD, src=(c3, c3))
     self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), c2.rtag())
@@ -82,11 +82,11 @@ class TestPatternMatcher(unittest.TestCase):
       (UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()),
       (UPat(Ops.MAX, name="x"), lambda x: x.rtag()),
     ])
-    c1 = UOp.const(0.0)
-    c2 = UOp.const(False)
+    c1 = UOp.const(dtypes.float, 0.0)
+    c2 = UOp.const(dtypes.bool, False)
     c3 = UOp(Ops.MAX, src=(c1, c1))
     c4 = UOp(Ops.MUL, src=(c1, c1))
-    c5 = UOp.const(-1)
+    c5 = UOp.const(dtypes.int, -1)
     self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), c2.rtag())
     self.assertEqual(matcher.rewrite(c3), c3.rtag())
@@ -96,11 +96,11 @@ class TestPatternMatcher(unittest.TestCase):
   def test_filter_arg(self):
     matcher = PatternMatcher([
       (UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"),
-       lambda x,c: x.rtag() if c.val in {1, -1} else None)
+       lambda x,c: x.rtag() if c.arg in {1, -1} else None)
     ])
-    y1 = UOp.const(1)
-    y2 = UOp.const(2)
-    y3 = UOp.const(-1)
+    y1 = UOp.const(dtypes.int, 1)
+    y2 = UOp.const(dtypes.int, 2)
+    y3 = UOp.const(dtypes.int, -1)
     c1 = UOp(Ops.MUL, src=(y1, y2))
     c2 = UOp(Ops.MUL, src=(y2, y2))
     c3 = UOp(Ops.MUL, src=(y3, y2))
@@ -114,27 +114,26 @@ class TestPatternMatcher(unittest.TestCase):
 
   def test_dup_name(self):
     matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())])
-    y1 = UOp.const(1.0)
-    y2 = UOp.const(1.0)
+    y1 = UOp.const(dtypes.float, 1.0)
+    y2 = UOp.const(dtypes.float, 1.0)
     c1 = UOp(Ops.ADD, src=(y1, y1))
     c2 = UOp(Ops.ADD, src=(y1, y2))
     self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), c1.rtag())
 
   def test_dtype(self):
-    # a concrete const dtype lives on the pair's CAST
-    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
-    c1 = UOp.const(1.0).cast(dtypes.float32)
-    c2 = UOp.const(1.0).cast(dtypes.float64)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float64, 1.0)
     self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
 
   def test_dtype_set(self):
-    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
-    c1 = UOp.const(1.0).cast(dtypes.float32)
-    c2 = UOp.const(1.0).cast(dtypes.float64)
-    c3 = UOp.const(1.0).cast(dtypes.float16)
-    c4 = UOp.const(1).cast(dtypes.int)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float64, 1.0)
+    c3 = UOp.const(dtypes.float16, 1.0)
+    c4 = UOp.const(dtypes.int, 1)
     self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), c2.rtag())
     self.assertEqual(matcher.rewrite(c3), None)
@@ -142,8 +141,8 @@ class TestPatternMatcher(unittest.TestCase):
 
   def test_src_one(self):
     matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
     c3 = UOp(Ops.ADD, src=(c1,c2))
     self.assertEqual(matcher.rewrite(c3), c3.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
@@ -159,8 +158,8 @@ class TestPatternMatcher(unittest.TestCase):
 
   def test_src_permutations(self):
     matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
     c3 = UOp(Ops.ADD, src=(c1,c2))
     c4 = UOp(Ops.ADD, src=(c3,c2))
     c5 = UOp(Ops.ADD, src=(c2,c3))
@@ -172,8 +171,8 @@ class TestPatternMatcher(unittest.TestCase):
 
   def test_src_repeat(self):
     matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
     c3 = UOp(Ops.ADD, src=(c1,c2))
     c4 = UOp(Ops.ADD, src=(c2,c3))
     self.assertEqual(matcher.rewrite(c3), c3.rtag())
@@ -181,9 +180,9 @@ class TestPatternMatcher(unittest.TestCase):
 
   def test_allow_len(self):
     matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
     c4 = UOp(Ops.EXP2, src=(c1,))
     c5 = UOp(Ops.ADD, src=(c1,c2))
     c6 = UOp(Ops.MULACC, src=(c1,c2,c3))
@@ -192,8 +191,8 @@ class TestPatternMatcher(unittest.TestCase):
     self.assertEqual(matcher.rewrite(c6), c6.rtag())
 
   def test_deep_src_permutations(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
     u1 = (c1 + c2) + c1
     u2 = (c2 + c1) + c1
     matcher = PatternMatcher([

test/null/test_schedule.py

@@ -1640,11 +1640,11 @@ class TestSchedule(unittest.TestCase):
     self.assertEqual(GlobalCounters.mem_used-base, 0)
 
   def test_const_schedule(self):
-    constv = Tensor.empty(2, 2).const_like(10).uop
+    constv = Tensor.empty(2, 2).uop.const_like(10)
     check_schedule(constv, 0)
 
   def test_const_schedule_contig(self):
-    constv = Tensor.empty(2, 2).const_like(10).uop.contiguous()
+    constv = Tensor.empty(2, 2).uop.const_like(10).contiguous()
     check_schedule(constv, 0)
 
   def test_advanced_simple_indexing_combined(self):

test/null/test_simplify_valid_idx.py

@@ -2,7 +2,7 @@ import unittest, itertools
 
 from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite, pm_lower_index_dtype
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
 from test.helpers import full_rewrite
@@ -23,7 +23,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
     UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)),
   ))
 
-def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr)
+def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(dtypes.weakint, nmax),), arg=expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
 def Range(n, nmax): return UOp.range(nmax, n)
 
@@ -207,7 +207,7 @@ class TestImageSimplification(unittest.TestCase):
     if svalid is not None:
       check_uop_against_string(self, off.src[1].get_valid(), svalid)
     else:
-      self.assertEqual(off.src[1].get_valid(), UOp.const(True), "svalid is None but valid is not True")
+      self.assertEqual(off.src[1].get_valid(), UOp.const(dtypes.bool, True), "svalid is None but valid is not True")
 
   def test_idx_gt_c(self):
     # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid
@@ -455,7 +455,7 @@ class TestImageSimplification(unittest.TestCase):
     A1 = lidx0*32 + r0*32 + lidx1*4 - 99
     valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 19)
     alu0 = gidx0 + (A1 % 32)*32 + (A1 // 32 % 16)*1024
-    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(0)))
+    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
     try:
       self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0")
     except AssertionError:
@@ -474,7 +474,7 @@ class TestImageSimplification(unittest.TestCase):
     A1 = lidx0*16 + r0*16 + lidx1*4 - 51
     valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 11)
     alu0 = lidx2 + gidx0*4 + (A1 % 16)*64 + (A1 // 16 % 8)*1024
-    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(0)))
+    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
     try:
       self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
     except AssertionError:
@@ -488,20 +488,10 @@ class TestImageSimplification(unittest.TestCase):
     gidx0 = Special("gidx0", 1064)
     r12 = Range(12, 3)
     valid = ((gidx0 < 645).ne(True)) & (gidx0 < 653)
-    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(0))
+    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(dtypes.weakint, 0))
     load = get_load_image_uop((1, 48, 4), valid, idx)
     self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0")
 
-  def test_drop_gate_committed_in_the_index_pass(self):
-    # the fused index pass runs without symbolic, so committing a weak src must not leave a CAST that
-    # symbolic later folds inside the index only: the gate's copy of the expression has to stay the same node
-    f = UOp.variable("f", 0.0, 9.0, dtypes.float)
-    idx_y = (f + UOp.const(1.0)).cast(dtypes.int)
-    load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)),
-                              (Special("gidx0", 10), idx_y))
-    off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0]
-    self.assertEqual(off.src[1].get_valid(), UOp.const(True))
-
 class TestDropTrueGate(unittest.TestCase):
   def test_drop_true_gate_on_index(self):
     # test that INDEX with a constant True valid gets simplified to drop the valid
@@ -509,8 +499,8 @@ class TestDropTrueGate(unittest.TestCase):
     from tinygrad.uop.ops import graph_rewrite
     from tinygrad.uop.symbolic import sym
     buf = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(0)
-    true_gate = UOp.const(True)
+    idx = UOp.const(dtypes.weakint, 0)
+    true_gate = UOp.const(dtypes.bool, True)
     index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
     # apply the optimization
     result = graph_rewrite(index_with_gate, sym+indexing_simplify)
@@ -526,50 +516,50 @@ class TestRangeShrink(unittest.TestCase):
   def test_range_shrink_single_guard(self):
     # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
     r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(4), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assertEqual(ranges[0].src[0].arg, 4)
 
   def test_range_shrink_picks_max_guard(self):
     # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
     r = Range(0, 204)
-    load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = get_gated_load_uop(r < UOp.const(8), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r)
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 8)
+    self.assertEqual(ranges[0].src[0].arg, 8)
 
   def test_range_no_shrink_guard_ge_max(self):
     # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
     r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(300), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 300), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assertEqual(ranges[0].src[0].arg, 204)
 
   def test_range_no_shrink_when_unguarded_elsewhere(self):
     # one load guards r < 4, but another load uses r without a gate -> no shrink
     r = Range(0, 204)
-    load1 = get_gated_load_uop(r < UOp.const(4), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assertEqual(ranges[0].src[0].arg, 204)
 
   def test_range_no_shrink_when_used_in_reduce(self):
     # range used in both a gated load AND directly in the reduce expression -> no shrink
     r = Range(0, 204)
-    gated_load = get_gated_load_uop(r < UOp.const(4), r)
+    gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
     ranges = self.get_ranges(red.sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assertEqual(ranges[0].src[0].arg, 204)
 
   def test_range_shrink_to_single_iteration(self):
     # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
     r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(1), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 1), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 0)
 
@@ -577,19 +567,19 @@ class TestRangeShrink(unittest.TestCase):
     # emulates mask.where(x.pad_to(mask.shape), Invalid): range should shrink accordingly
     from tinygrad.dtype import Invalid
     r = Range(0, 204)
-    x = (r < 4).where(UOp.const(1.0), Invalid)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assertEqual(ranges[0].src[0].arg, 4)
 
   def test_range_shrink_store_where_invalid_flipped(self):
     # above, but flipped
     from tinygrad.dtype import Invalid
     r = Range(0, 204)
-    x = (r < 4).where(UOp.const(1.0), Invalid)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assertEqual(ranges[0].src[0].arg, 4)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_symbolic_failures.py

@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Variable
-from tinygrad.uop.ops import UOp
 
 
 class TestFuzzFailure(unittest.TestCase):
@@ -9,7 +8,7 @@ class TestFuzzFailure(unittest.TestCase):
     v2=Variable('v2', 0, 2)
     v3=Variable('v3', 0, 1)
     expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3)
-    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(0), v3.const_like(0)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -20,7 +19,7 @@ class TestFuzzFailure(unittest.TestCase):
     v3=Variable('v3', 0, 3)
     expr = (((((((((((((((((((((((((0*4)//5)*2)*-1)*-2)+-4)*4)*2)*3)*4)+-4)*4)+v2)+v2)+v3)//3)+v2)+v1)//9)+3)+1)//1)+-4)//4)*2)
     expr = (((((v1+(v2+(((v3+(v2*2))+1)//3)))+4)//9)+-57)//(9*4))
-    v1_val, v2_val, v3_val = UOp.const(6), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(6), v2.const_like(0), v3.const_like(0)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -30,7 +29,7 @@ class TestFuzzFailure(unittest.TestCase):
     v2=Variable('v2', 0, 1)
     v3=Variable('v3', 0, 2)
     expr = (((((((((((((((((((0//2)//3)+v3)+0)+-4)*-2)*-2)+-1)+2)+3)+v3)+0)//8)*-3)+0)*-2)*-4)*-2)//5)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -40,7 +39,7 @@ class TestFuzzFailure(unittest.TestCase):
     v2=Variable('v2', 0, 3)
     v3=Variable('v3', 0, 4)
     expr = (((((((((((((((((((((((((((((0*-2)+0)*-1)//9)//6)//8)+v1)*-4)+v2)//4)//8)+4)*3)+v1)+v3)//8)//7)+4)+v3)*-4)+1)+v1)*3)+4)*2)//5)//2)//3)*-4)
-    v1_val, v2_val, v3_val = UOp.const(2), UOp.const(0), UOp.const(2)
+    v1_val, v2_val, v3_val = v1.const_like(2), v2.const_like(0), v3.const_like(2)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -50,7 +49,7 @@ class TestFuzzFailure(unittest.TestCase):
     v2=Variable('v2', 0, 1)
     v3=Variable('v3', 0, 3)
     expr = ((((((((((((((0+v2)+v1)*0)+v2)//1)//7)+-2)+v2)+v1)*4)+-3)//5)+v2)+1)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -61,7 +60,7 @@ class TestFuzzFailure(unittest.TestCase):
     v3=Variable('v3', 0, 128)
     expr = (((((((((((((((((((((((((((((0//3)+4)+v1)//2)+-1)//1)*1)*-1)*4)//5)+v1)//6)+v1)*-1)+-4)+v2)+-2)*-3)+v3)+-4)+-2)*-1)//8)//4)*-4)+3)+v3)*
             -2)+v2)
-    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(3), UOp.const(2)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(3), v3.const_like(2)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -71,7 +70,7 @@ class TestFuzzFailure(unittest.TestCase):
     v2=Variable('v2', 0, 5)
     v3=Variable('v3', 0, 128)
     expr = (((((((((((((((((((((((((((((0+v2)*-4)+0)//9)+-4)*-2)*3)*4)//9)+v3)+v1)//4)+v1)+v3)+-1)*4)//4)+v2)//7)//3)+v1)+v2)+v3)+1)*2)//4)*3)+-1)*1)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(2), UOp.const(65)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(2), v3.const_like(65)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -81,7 +80,7 @@ class TestFuzzFailure(unittest.TestCase):
     v2=Variable('v2', 0, 8)
     v3=Variable('v3', 0, 9)
     expr = (((((((0+-1)+2)+v1)*-2)//3)+v1)*-4)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -91,7 +90,7 @@ class TestFuzzFailure(unittest.TestCase):
     v2=Variable('v2', 0, 1)
     v3=Variable('v3', 0, 8)
     expr = (((((((((((((((((((((((((((((0*-2)//1)+3)*-2)+-3)*-4)*1)+v1)+0)%2)%8)%9)+v2)%9)+-4)//4)+-1)*-2)+0)+v1)+v1)+3)+v1)+4)+-4)+0)*2)+-3)%6)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(1), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(1), v3.const_like(0)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -102,7 +101,7 @@ class TestFuzzFailure(unittest.TestCase):
     v3=Variable("v3", 0, 32)
     x5 = (v1 <= 9).where(v1 * -4 - 4, v1 // 9) // 9
     expr = ((x5 >= -4).where(x5, (v2 % 3 + v2) // 5) * -1).maximum(((v1 * -2) % 6 + v3 % 1) * -1) * -1
-    v1_val, v2_val, v3_val = UOp.const(9), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(9), v2.const_like(0), v3.const_like(0)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
@@ -112,7 +111,7 @@ class TestFuzzFailure(unittest.TestCase):
     v2=Variable("v2", 0, 128)
     v3=Variable("v3", 0, 5)
     expr = (((v2 * 0).maximum(8) - v2 * 2) % 5 + v1 // 6 + v1 + 5) % 5
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(7), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(7), v3.const_like(0)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)

test/null/test_tensor_uop_mixin.py

@@ -60,7 +60,7 @@ class TestTensorUOpClone(unittest.TestCase):
     t = _t(3, 4).float()
     self.assertIs(_strip_unique(t.clone().uop), _strip_unique(t.uop.clone()))
   def test_clone_deviceless_const(self):
-    u = UOp.const(2.0)
+    u = UOp.const(dtypes.float, 2.0)
     self.assertIs(_strip_unique(Tensor(u).clone().uop), _strip_unique(u.clone()))
 
 class TestTensorUOpGradient(unittest.TestCase):
@@ -382,7 +382,7 @@ class TestTensorUOpStack(unittest.TestCase):
     self.assertIs(_t(2, 3).uop.stack(w.uop).dtype, dtypes.float32)
   def test_stack_index_dtype(self):
     # index is outside the promotion lattice, equal dtypes bypass promotion
-    self.assertEqual(UOp.const(1).stack(UOp.const(2)).shape, (2,))
+    self.assertEqual(UOp.const(dtypes.weakint, 1).stack(UOp.const(dtypes.weakint, 2)).shape, (2,))
 
 class TestTensorUOpConv2d(unittest.TestCase):
   def test_conv2d_basic(self):

test/null/test_transcendental_helpers.py

@@ -10,7 +10,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
     # 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,))
-    loaded_value = input_buf.index(UOp.const(0)).load()
+    loaded_value = input_buf.index(UOp.const(dtypes.int, 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))
 
@@ -27,48 +27,48 @@ class TestTranscendentalFunctions(unittest.TestCase):
     np.testing.assert_equal(q, 4)
 
   def test_cody_waite_reduction(self):
-    r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(12 * math.pi + 0.1).cast(dtypes.float64)))
+    r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(dtypes.float64, 12 * math.pi + 0.1)))
     np.testing.assert_allclose(r, 0.1)
     np.testing.assert_equal(q, 12)
 
   def test_frexp(self):
     for x in (1, -1):
-      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(float(x)).cast(dtypes.float64)))
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, x)))
       np.testing.assert_equal(mantissa, 0.5)
       np.testing.assert_equal(exponent, 1)
 
     for x in (2, -2):
-      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(2.0).cast(dtypes.float64)))
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 2.0)))
       np.testing.assert_equal(mantissa, 0.5)
       np.testing.assert_equal(exponent, 2)
 
-    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(5.0).cast(dtypes.float64)))
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 5.0)))
     np.testing.assert_equal(mantissa, 0.625)
     np.testing.assert_equal(exponent, 3)
 
-    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(1000.0).cast(dtypes.float64)))
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 1000.0)))
     np.testing.assert_allclose(mantissa, 0.9765625)
     np.testing.assert_equal(exponent, 10)
 
   def test_rintk(self):
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(0.0).cast(dtypes.float))), 0)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.0).cast(dtypes.float))), 5)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.5).cast(dtypes.float))), 6)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.999).cast(dtypes.float))), 6)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.0).cast(dtypes.float))), -5)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.5).cast(dtypes.float))), -6)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.999).cast(dtypes.float))), -6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 0.0))), 0)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.0))), 5)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.5))), 6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.999))), 6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.0))), -5)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.5))), -6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.999))), -6)
 
   def test_pow2if(self):
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(0).cast(dtypes.int), dtypes.float)), 1.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(1).cast(dtypes.int), dtypes.float)), 2.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(2).cast(dtypes.int), dtypes.float)), 4.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(10).cast(dtypes.int), dtypes.float)), 1024.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(63).cast(dtypes.int), dtypes.float)), 2**63)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-1).cast(dtypes.int), dtypes.float)), 0.5)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-2).cast(dtypes.int), dtypes.float)), 0.25)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-10).cast(dtypes.int), dtypes.float)), 2**-10)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-63).cast(dtypes.int), dtypes.float)), 2**-63)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 0), dtypes.float)), 1.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 1), dtypes.float)), 2.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 2), dtypes.float)), 4.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 10), dtypes.float)), 1024.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 63), dtypes.float)), 2**63)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -1), dtypes.float)), 0.5)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -2), dtypes.float)), 0.25)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -10), dtypes.float)), 2**-10)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -63), dtypes.float)), 2**-63)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_uop_graph.py

@@ -7,35 +7,35 @@ from tinygrad.uop.symbolic import sym
 from test.helpers import to_uops_list
 
 simple_pm = PatternMatcher([
-  (UPat.cvar('x', dtypes.weakint), lambda x: UOp.const(1.0) + UOp.const(2.0)),
-  (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(x.val+y.val)),
-  (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(x.val*y.val*z.val)),
-  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.val+c2.val)),
+  (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),
+  (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(dtypes.float, x.arg+y.arg)),
+  (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(dtypes.float, x.arg*y.arg*z.arg)),
+  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)),
 ])
 
 def const_values(u:UOp):
-  if u.op is Ops.CONST: return (u.val,)
-  if u.op is Ops.STACK: return tuple(x.val for x in u.src)
+  if u.op is Ops.CONST: return (u.arg,) if not isinstance(u.arg, tuple) else u.arg
+  if u.op is Ops.STACK: return tuple(x.arg for x in u.src)
   raise AssertionError(f"expected const-like UOp, got {u.op}")
 
 class TestGraphRewriteConst(unittest.TestCase):
   def test_gep_const(self):
-    v1 = UOp.const((0,1,2), dtypes.int)
+    v1 = UOp.const(dtypes.int, (0,1,2))
     v2 = v1.index(1)
     ret = graph_rewrite(v2, sym)
     self.assertEqual(ret.dtype, dtypes.int)
-    self.assertEqual(ret.val, 1)
+    self.assertEqual(ret.arg, 1)
 
   def test_add_const(self):
-    v1 = UOp.const((0,1,2))
-    v2 = UOp.const((5,6,7))
+    v1 = UOp.const(dtypes.int, (0,1,2))
+    v2 = UOp.const(dtypes.int, (5,6,7))
     ret = graph_rewrite(v1+v2, sym)
     self.assertEqual(ret.op, Ops.STACK)
     self.assertEqual(const_values(ret), (5,7,9))
 
   def test_add_const_lose_v(self):
-    v1 = UOp.const((0,1,2))
-    v2 = UOp.const((2,1,0))
+    v1 = UOp.const(dtypes.int, (0,1,2))
+    v2 = UOp.const(dtypes.int, (2,1,0))
     ret = graph_rewrite(v1+v2, sym)
     self.assertEqual(ret.op, Ops.STACK)
     self.assertEqual(const_values(ret), (2,2,2))
@@ -49,47 +49,47 @@ class TestModularWraparound(unittest.TestCase):
     self.assertEqual(len(results), 2)  # +1 for SINK
     self.assertEqual(results[0].op, Ops.CONST)
     self.assertEqual(results[0].dtype, uop.dtype)
-    self.assertEqual(results[0].val, expected)
+    self.assertEqual(results[0].arg, expected)
 
   @xfail_broken_const_wraparound
   def test_cast(self):
     t = self._test
-    t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8), 0xD6)
-    t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)
+    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8), 0xD6)
+    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)
 
   @xfail_broken_const_wraparound
   def test_mul(self):
     t = self._test
-    t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD, 1147018174)
-    t(UOp.const(0xABCD17D6, dtypes.int) * 10, -1241321892)
+    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD, 1147018174)
+    t(UOp.const(dtypes.int, 0xABCD17D6) * 10, -1241321892)
 
   @xfail_broken_const_wraparound
   def test_div(self):
     t = self._test
-    t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD // 11, 104274379)
-    t(UOp.const(0xABCD17D6, dtypes.int) * 10 // 11, -112847444)
+    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD // 11, 104274379)
+    t(UOp.const(dtypes.int, 0xABCD17D6) * 10 // 11, -112847444)
 
   @xfail_broken_const_wraparound
   def test_neg(self):
     t = self._test
-    t(-UOp.const(1, dtypes.uint8), 0xFF)
-    t(-UOp.const(1, dtypes.uint16), 0xFFFF)
-    t(-UOp.const(1, dtypes.uint32), 0xFFFFFFFF)
-    t(-UOp.const(1, dtypes.uint64), 0xFFFFFFFFFFFFFFFF)
+    t(-UOp.const(dtypes.uint8, 1), 0xFF)
+    t(-UOp.const(dtypes.uint16, 1), 0xFFFF)
+    t(-UOp.const(dtypes.uint32, 1), 0xFFFFFFFF)
+    t(-UOp.const(dtypes.uint64, 1), 0xFFFFFFFFFFFFFFFF)
 
   @xfail_broken_const_wraparound
   def test_neg_min_int(self):
     t = self._test
-    t(-UOp.const(-2**7, dtypes.int8), -2**7)
-    t(-UOp.const(-2**15, dtypes.int16), -2**15)
-    t(-UOp.const(-2**31, dtypes.int32), -2**31)
-    t(-UOp.const(-2**63, dtypes.int64), -2**63)
+    t(-UOp.const(dtypes.int8, -2**7), -2**7)
+    t(-UOp.const(dtypes.int16, -2**15), -2**15)
+    t(-UOp.const(dtypes.int32, -2**31), -2**31)
+    t(-UOp.const(dtypes.int64, -2**63), -2**63)
 
   @xfail_broken_const_wraparound
   def test_payne_hanek_reduction_bug(self):
     t = self._test
-    a = (UOp.const(43748177600, dtypes.uint).cast(dtypes.uint) | 36).cast(dtypes.ulong)
-    b = 2536655455 * a + 4294967296 * UOp.const(25366554550, dtypes.ulong)
+    a = (UOp.const(dtypes.uint, 43748177600).cast(dtypes.uint) | 36).cast(dtypes.ulong)
+    b = 2536655455 * a + 4294967296 * UOp.const(dtypes.ulong, 25366554550)
     c = (b + 2261737165) // 4611686018427387904
     t(c, 0)
 
@@ -103,67 +103,67 @@ class TestGraphRewrite(unittest.TestCase):
   # NOTE: this shows why we can't have a UOp in arg
   @unittest.expectedFailure
   def test_no_dedup_args(self):
-    a1 = UOp.variable("a1", UOp.const(0), UOp.const(11), dtypes.int)
-    a2 = UOp.variable("a2", UOp.const(0), UOp.const(11), dtypes.int)
+    a1 = UOp.variable("a1", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11), dtypes.int)
+    a2 = UOp.variable("a2", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11), dtypes.int)
     sink = a1.sink(a2)
     variables = [x for x in graph_rewrite(sink, PatternMatcher([])).toposort() if x.op is Ops.PARAM and x.addrspace is AddrSpace.ALU]
     self.assertEqual(len(variables), 1)
 
   def test_simple(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
     nout = graph_rewrite(c1+c2, simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 3.0)
+    self.assertEqual(nout.arg, 3.0)
 
   def test_depth_2_late(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
     nout = graph_rewrite(c1*c2*(c3+c3), simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 12.0)
+    self.assertEqual(nout.arg, 12.0)
 
   def test_double(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
     nout = graph_rewrite(c1+c2+c3, simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 6.0)
+    self.assertEqual(nout.arg, 6.0)
 
   def test_triple(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
-    c4 = UOp.const(4.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    c4 = UOp.const(dtypes.float, 4.0)
     nout = graph_rewrite(c1+c2+c3+c4, simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 10.0)
+    self.assertEqual(nout.arg, 10.0)
 
   def test_diamond(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
     nout = graph_rewrite((c1+c2)+(c1+c3), simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 7.0)
+    self.assertEqual(nout.arg, 7.0)
 
   def test_magic_4(self):
-    c1 = UOp.const(4)
+    c1 = UOp.const(dtypes.int, 4.0)
     nout = graph_rewrite(c1, simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 3.0)
+    self.assertEqual(nout.arg, 3.0)
 
   def test_depth_2_fold(self):
     v = UOp.variable("v", 0, 1, dtypes.float)
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
     nout = graph_rewrite(v+c1+c2, simple_pm)
     self.assertEqual(nout.op, Ops.ADD)
     self.assertEqual(nout.src[0].op, Ops.PARAM)
     self.assertEqual(nout.src[1].op, Ops.CONST)
-    self.assertEqual(nout.src[1].val, 3.0)
+    self.assertEqual(nout.src[1].arg, 3.0)
 
   def test_commutative_work(self):
     a = UOp.variable('a', 0, 1)
@@ -174,14 +174,14 @@ class TestGraphRewrite(unittest.TestCase):
     a = UOp.variable('a', 0, 1)
     tst = (2+a).simplify()
     self.assertIs(tst.src[0], a)
-    self.assertIs(tst.src[1], UOp.const(2))
+    self.assertIs(tst.src[1], a.const_like(2))
 
   def test_consts_go_last(self):
     a = UOp.variable('a', 0, 1)
     b = UOp.variable('b', 0, 1)
     c = UOp.variable('c', 0, 1)
     d = UOp.variable('d', 0, 1)
-    outs = [2+a, 2+a+d+3+b+c+4, UOp.const(2)+a, (4+d)+c+(2+a)+b]
+    outs = [2+a, 2+a+d+3+b+c+4, a.const_like(2)+a, (4+d)+c+(2+a)+b]
     for out in outs:
       sink = graph_rewrite(out, sym)
       print(sink.render())
@@ -191,67 +191,85 @@ class TestGraphRewrite(unittest.TestCase):
 
 class TestUOpGraph(unittest.TestCase):
   def test_add_constant_fold(self):
-    c1 = UOp.const(1.0, dtypes.float)
-    c2 = UOp.const(2.0, dtypes.float)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
     out = c1+c2
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 3.0)
+    self.assertEqual(out.arg, 3.0)
 
   def test_where_same_fold(self):
     v = UOp.variable('tmp', 0, 1)
-    c0 = UOp.const(0)
+    c0 = UOp.const(dtypes.weakint, 0)
     vc = v != c0
-    c1 = UOp.const(1.0, dtypes.float)
+    c1 = UOp.const(dtypes.float, 1.0)
     out = vc.where(c1, c1)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 1.0)
+    self.assertEqual(out.arg, 1.0)
 
   def test_where_const_fold(self):
-    bf = UOp.const(False)
-    c1 = UOp.const(1.0, dtypes.float)
-    c2 = UOp.const(2.0, dtypes.float)
+    bf = UOp.const(dtypes.bool, False)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
     out = bf.where(c1, c2)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 2.0)
+    self.assertEqual(out.arg, 2.0)
 
   def test_const_cast(self):
-    bf = UOp.const(False)
+    bf = UOp.const(dtypes.bool, False)
     out = bf.cast(dtypes.int)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0)
+    self.assertEqual(out.arg, 0)
 
   def test_const_bitcast(self):
-    bf = UOp.const(1.0, dtypes.float)
+    bf = UOp.const(dtypes.float, 1.0)
     out = bf.bitcast(dtypes.uint32)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0x3F800000)
+    self.assertEqual(out.arg, 0x3F800000)
 
   @unittest.expectedFailure
   def test_const_shape_change_bitcast(self):
-    bf = UOp.const(0x3F).cast(dtypes.uint8)
+    bf = UOp.const(dtypes.uint8, 0x3F)
     out = bf.bitcast(dtypes.half)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
 
+  def test_coalesce_aliased_stores(self):
+    from tinygrad.codegen.late.coalesce import memory_coalescing
+    from tinygrad import Device
+    # two distinct stores to the same INDEX (e.g. a double-buffered LDS slot written in a
+    # prologue and a loop body) can't be merged: merging would drop one write
+    lbuf = UOp.placeholder((64,), dtypes.half, 0, AddrSpace.LOCAL)
+    r = UOp.range(4, 0, AxisType.LOOP)
+    gbuf = UOp.placeholder((64,), dtypes.half, 1, AddrSpace.GLOBAL)
+    s_a = lbuf.index(r*2).store(gbuf.index(r*2).load())
+    s_b = lbuf.index(r*2).store(gbuf.index(r*2+64).load())   # aliases s_a
+    s_c = lbuf.index(r*2+1).store(gbuf.index(r*2+32).load()) # adjacent, unique -> coalesceable with nothing
+    out = memory_coalescing(UOp.sink(s_a, s_b, s_c).end(r), Device["NULL"].renderer)
+    stores = [u for u in out.toposort() if u.op is Ops.STORE]
+    # both aliased stores survive (both datas preserved), nothing merged across them
+    self.assertEqual(len([u for u in stores if u.src[0].src[1] is not None and u.src[0].op is Ops.INDEX]), 3)
+    datas = sorted([u.src[1] for u in stores], key=str)
+    self.assertEqual(len(datas), 3)
+
   def test_devectorize_derives_lane_dtype(self):
     from tinygrad.codegen import do_devectorize
     # an Invalid lane derives bool while the value lane derives float: the lane rebuild must derive, not inherit
-    lhs = UOp.stack(UOp.invalid(), UOp.const(1.0).cast(dtypes.float))
+    lhs = UOp.stack(UOp.invalid(), UOp.const(None, 1.0).cast(dtypes.float))
     out = do_devectorize(lhs * lhs)
     invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
     self.assertIs(invalid_lane_mul.dtype, dtypes.bool)
@@ -259,7 +277,7 @@ class TestUOpGraph(unittest.TestCase):
   @unittest.skip("this test isn't valid uops")
   def test_noop_vectorize_fold(self):
     d0 = UOp.param(0, dtypes.float, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
     ld = d0.load(idx, dtype=dtypes.float)
     vec = UOp(Ops.STACK, dtypes.float, (ld,))
     x = vec.index(0)
@@ -273,7 +291,7 @@ class TestUOpGraph(unittest.TestCase):
     d0 = UOp.param(0, dtypes.float, (1,))
     d1 = UOp.param(1, dtypes.float, (1,))
     d2 = UOp.param(2, dtypes.float, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
     def _test_vec(geps, count=4):
       vec = UOp(Ops.STACK, dtypes.float, geps)
       out = d0.index(idx).store(vec)
@@ -310,7 +328,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_gep_vec_const_fold(self):
     for vec_size in [2, 4, 8]:
-      consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
+      consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
       vec = UOp(Ops.STACK, src=tuple(consts))
       with Context(SPEC=0):
         uops = to_uops_list([vec.index(i) for i in range(vec_size)])
@@ -320,7 +338,7 @@ class TestUOpGraph(unittest.TestCase):
   def test_cast_alu_fold(self):
     d0 = UOp.param(0, dtypes.bool, (1,))
     d1 = UOp.param(1, dtypes.int, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
     ld = d1.index(idx)
     alu = (ld<1).cast(dtypes.bool)
     out = d0.index(idx).store(alu)
@@ -330,7 +348,7 @@ class TestUOpGraph(unittest.TestCase):
   def test_double_cast_fold(self):
     d0 = UOp.param(0, dtypes.float, (1,))
     d1 = UOp.param(1, dtypes.int, (1,))
-    idx = UOp.const(0, dtypes.int)
+    idx = UOp.const(dtypes.int, 0)
     ld = d1.index(idx)
     alu = ld.cast(dtypes.float).cast(dtypes.float)
     out = d0.index(idx).store(alu)
@@ -339,8 +357,8 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_depth_2_const_fold(self):
     v = UOp.variable("tmp", 0, 1, dtypes.int)
-    c2 = UOp.const(2, dtypes.int)
-    c4 = UOp.const(4, dtypes.int)
+    c2 = UOp.const(dtypes.int, 2)
+    c4 = UOp.const(dtypes.int, 4)
     vc = v+c2
     out = vc+c4
     uops = to_uops_list([out])
@@ -348,19 +366,19 @@ class TestUOpGraph(unittest.TestCase):
     out = uops[-2]  # -2 to skip SINK
     self.assertEqual(out.op, Ops.ADD)
     self.assertEqual(out.src[1].op, Ops.CONST)
-    self.assertEqual(out.src[1].val, 6)
+    self.assertEqual(out.src[1].arg, 6)
 
   def test_bitcast_to_same_dtype_fold(self):
     for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
       d0 = UOp.param(0, dt, (1,))
-      v = d0.index(UOp.const(0))
+      v = d0.index(UOp.const(dtypes.int, 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}")
 
   def test_sub_with_cast_folds(self):
     a = Variable("a", 0, 5)
     uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
-    assert uops[0] == UOp.const(0, dtypes.int)
+    assert uops[0] == UOp.const(dtypes.int, 0)
     assert uops[-1].op == Ops.SINK
 
   def test_where_on_gated_load_fold(self):
@@ -372,7 +390,7 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([out.index(ridx0).store(w)])
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg==5
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
@@ -382,7 +400,7 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([w])
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: assert u.src[1].val==5
+      if u.op is Ops.LOAD: assert u.src[1].arg==5
 
   def test_where_on_gated_load_with_cast(self):
     ridx0 = UOp.range(100, 0)
@@ -394,13 +412,13 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([out.index(ridx0).store(w)])
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg == 5
 
   def test_where_on_casted_gated_load_extra_cond(self):
     ridx0 = UOp.range(100, 0)
     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)
+    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(dtypes.float, 0)).cast(dtypes.half)
     out = UOp.param(1, dtypes.half, (100,))
     uops = to_uops_list([out.index(ridx0).store(w)])
     for u in uops:
@@ -410,7 +428,7 @@ class TestUOpGraph(unittest.TestCase):
     ridx0 = UOp.range(100, 0)
     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)
+    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(dtypes.float, 0), ld).cast(dtypes.half)
     out = UOp.param(1, dtypes.half, (100,))
     uops = to_uops_list([out.index(ridx0).store(w)])
     for u in uops:
@@ -426,22 +444,22 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([st])
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.STORE: assert u.src[1].val==5
+      if u.op is Ops.STORE: assert u.src[1].arg==5
 
   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,))
-    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
     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)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
     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)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -449,19 +467,19 @@ 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,))
-    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
     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)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
     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)
-    comparison = (loaded_idx + c5) < UOp.const(60000)
+    comparison = (loaded_idx + c5) < UOp.const(dtypes.weakint, 60000)
     c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -470,9 +488,9 @@ class TestUOpGraph(unittest.TestCase):
     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)
+    idx = UOp.const(dtypes.int, 0)
     ld0 = glbl1.index(UOp.invalid())
-    ld1 = glbl2.index(idx.valid(UOp.const(True)))
+    ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True)))
     uops = to_uops_list([glbl0.index(idx).store(ld1+ld0)])
     # the gate and invalid value are deleted from ld1
     self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)
@@ -484,7 +502,7 @@ class TestUOpGraph(unittest.TestCase):
     st = smem.index(lidx).store(glbl0.index(lidx).load())
     barrier = st.barrier()
     ld0 = smem.after(barrier).index(UOp.invalid())
-    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(True)))
+    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(dtypes.bool, True)))
     uops = to_uops_list([glbl0.index(lidx).store(ld1+ld0)])
 
     # the gate and invalid value are deleted from ld1
@@ -492,10 +510,10 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_fold_gated_store(self):
     glbl = UOp.param(0, dtypes.int, (1,))
-    idx0 = UOp.const(0)
-    val = UOp.const(42)
+    idx0 = UOp.const(dtypes.int, 0)
+    val = UOp.const(dtypes.int, 42)
     st0 = glbl.index(UOp.invalid()).store(val)
-    st1 = glbl.index(idx0.valid(UOp.const(True))).store(val)
+    st1 = glbl.index(idx0.valid(UOp.const(dtypes.bool, True))).store(val)
     uops = to_uops_list([st0, st1])
     # only the second store happens
     self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1)
@@ -503,9 +521,9 @@ class TestUOpGraph(unittest.TestCase):
   @unittest.skip("this is a uop type error")
   def test_asserts_bad_gate(self):
     glbl0 = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(0)
-    bad_gate = UOp.const(1)
-    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))])
+    idx = UOp.const(dtypes.int, 0)
+    bad_gate = UOp.const(dtypes.int, 1)
+    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])
 
   def test_after_end(self):
     r = UOp.range(10, 0)
@@ -513,7 +531,7 @@ class TestUOpGraph(unittest.TestCase):
     c = r + 1
     self.assertIn(r, c.ranges)
 
-    e = UOp.const(1).end(r)
+    e = UOp.const(dtypes.int, 1).end(r)
     self.assertNotIn(r, e.ranges)
 
     a = c.after(e)
@@ -563,7 +581,7 @@ class TestConstBufferize(unittest.TestCase):
     CONST doesn't depend on ranges (constant is same value everywhere).
     """
     from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
-    c = UOp.const(42.0)
+    c = UOp.const(dtypes.float, 42.0)
     r1 = UOp.range(3, 0)
     bufferize_with_range = UOp(Ops.STAGE, src=(c, r1), arg=BufferizeOpts(device="CPU"))
     self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range
@@ -571,13 +589,13 @@ class TestConstBufferize(unittest.TestCase):
     result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
     # BUFFERIZE should be removed, result is const broadcast to shape
     self.assertNotEqual(result.op, Ops.STAGE)
-    const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
     self.assertIn(42.0, const_vals)
 
   def test_const_bufferize_with_multiple_ranges(self):
     """Test CONST.BUFFERIZE with multiple ranges is also folded."""
     from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
-    c = UOp.const(3.14)
+    c = UOp.const(dtypes.float, 3.14)
     r1 = UOp.range(3, 0)
     r2 = UOp.range(4, 1)
     bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU"))
@@ -586,14 +604,14 @@ class TestConstBufferize(unittest.TestCase):
     result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')
     # BUFFERIZE should be removed
     self.assertNotEqual(result.op, Ops.STAGE)
-    const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
     self.assertIn(3.14, const_vals)
 
 class TestUOpTags(unittest.TestCase):
   def test_inc_by_one(self):
-    g = UOp.const(1) + UOp.const(1)
+    g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)
     assert g.ssimplify() == 2
-    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.val+1, tag=1) if x.tag is None else None)])
+    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.arg+1, tag=1) if x.tag is None else None)])
     pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
     g = graph_rewrite(g, pm_plus_1)
     assert g.ssimplify() == 4
@@ -682,36 +700,36 @@ class TestUOpGetItem(unittest.TestCase):
 
 class TestUOpBroadcast(unittest.TestCase):
   def test_broadcast_row(self):
-    a = UOp.const(1.0).expand((4, 8))
-    b = UOp.const(2.0).expand((4, 1))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(4, 1))
     c = a + b
     self.assertEqual(c.shape, (4, 8))
     self.assertEqual(c.op, Ops.ADD)
 
   def test_broadcast_col(self):
-    a = UOp.const(1.0).expand((4, 8))
-    b = UOp.const(2.0).expand((1, 8))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(1, 8))
     c = a + b
     self.assertEqual(c.shape, (4, 8))
     self.assertEqual(c.op, Ops.ADD)
 
   def test_broadcast_lower_dim(self):
-    a = UOp.const(1.0).expand((4, 8))
-    b = UOp.const(2.0).expand((8,))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(8,))
     c = a * b
     self.assertEqual(c.shape, (4, 8))
     self.assertEqual(c.op, Ops.MUL)
 
   def test_broadcast_scalar(self):
-    a = UOp.const(1.0).expand((4, 8))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
     c = a * 2
     self.assertEqual(c.shape, (4, 8))
     self.assertEqual(c.op, Ops.MUL)
 
   def test_broadcast_symbolic_same_shape(self):
     t = Variable("t", 1, 10)
-    a = UOp.const(1.0).expand((1, 1, t))
-    b = UOp.const(2.0).expand((1, 1, t))
+    a = UOp.const(dtypes.float, 1, shape=(1, 1, t))
+    b = UOp.const(dtypes.float, 2, shape=(1, 1, t))
     c = a + b
     self.assertEqual(c.op, Ops.ADD)
 

test/null/test_uop_repr.py

@@ -1,34 +1,34 @@
 import unittest
-from tinygrad import UOp
+from tinygrad import UOp, dtypes
 
 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=())")
+    a = UOp.const(dtypes.int, 42)
+    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.int, arg=42, src=())")
   def test_different_consts(self):
-    a, b = UOp.const(42), UOp.const(3)
+    a, b = UOp.const(dtypes.int, 42), UOp.const(dtypes.int, 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, dtypes.int, arg=None, src=(\n" +
+      "  UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
+      "  UOp(Ops.CONST, dtypes.int, 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)
+    a = UOp.const(dtypes.int, 42)
     expected = (
-      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
+      "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
       "  x0,))"
     )
     self.assertEqual(repr(a+a), expected)
   def test_nested_walrus_indentation(self):
     # Ensure indentation is consistent at multiple levels
-    b = (a:=UOp.const(1)) + a
+    b = (a:=UOp.const(dtypes.int, 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, dtypes.int, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "    x1:=UOp(Ops.CONST, dtypes.int, arg=1, src=()),\n" +
       "    x1,)),\n" +
       "  x0,))"
     )

test/null/test_uop_resolve.py

@@ -4,48 +4,44 @@ from tinygrad.uop.ops import UOp, resolve
 
 class TestUOpResolve(unittest.TestCase):
   def test_simple_int(self):
-    u = UOp.const(4, dtypes.int)
+    u = UOp.const(dtypes.int, 4)
     self.assertEqual(int(u), 4)
 
-  def test_weak_const(self):
-    self.assertEqual(int(UOp.const(5)), 5)
-    self.assertEqual(float(UOp.const(1.5)), 1.5)
-
   def test_int_add(self):
-    u = UOp.const(4, dtypes.int) + 7
+    u = UOp.const(dtypes.int, 4) + 7
     self.assertEqual(int(u), 11)
 
   def test_lt(self):
-    u = UOp.const(4) < 7
+    u = UOp.const(dtypes.int, 4) < 7
     self.assertTrue(u)
 
   def test_rfloordiv(self):
-    u = 8 // UOp.const(4, dtypes.int)
+    u = 8 // UOp.const(dtypes.int, 4)
     self.assertEqual(int(u), 2)
 
   def test_rtruediv(self):
-    u = 9 / UOp.const(4, dtypes.float)
+    u = 9 / UOp.const(dtypes.float, 4)
     self.assertEqual(float(u), 2.25)
 
   def test_leq(self):
-    u = UOp.const(4) <= 4
+    u = UOp.const(dtypes.int, 4) <= 4
     self.assertTrue(u)
 
   def test_ne(self):
-    u = UOp.const(4) != 7
+    u = UOp.const(dtypes.int, 4) != 7
     self.assertTrue(u)
 
   def test_ne_f(self):
-    u = UOp.const(4) != 4
+    u = UOp.const(dtypes.int, 4) != 4
     self.assertFalse(u)
 
   def test_ngt(self):
-    u = UOp.const(4) > 7
+    u = UOp.const(dtypes.int, 4) > 7
     self.assertFalse(u)
 
   def test_ssimplify(self):
-    self.assertEqual((8 % UOp.const(4)).ssimplify(), 0)
-    self.assertEqual((8 * UOp.const(4)).ssimplify(), 32)
+    self.assertEqual((8 % UOp.const(dtypes.int, 4)).ssimplify(), 0)
+    self.assertEqual((8 * UOp.const(dtypes.int, 4)).ssimplify(), 32)
 
   def test_ambiguous_less_than(self):
     u = UOp.variable("i", 1, 10)
@@ -56,7 +52,7 @@ class TestUOpResolve(unittest.TestCase):
     self.assertFalse(resolve(u < -1, True))
 
   def test_float_direct(self):
-    u = UOp.const(4.5, dtypes.float) + 7
+    u = UOp.const(dtypes.float, 4.5) + 7
     self.assertEqual(float(u), 11.5)
 
   def test_var_cmp_t(self):

test/null/test_uop_symbolic.py

@@ -12,12 +12,13 @@ from tinygrad.uop.validate import uops_to_z3
 def check_uop_against_string(self, v:UOp, s:str):
   sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.RANGE, Ops.SPECIAL, Ops.PARAM)}
   s_eval = eval(s, sym_vars)
-  if isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(s_eval)
+  if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(dtypes.weakint, s_eval)
+  elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(dtypes.from_py(s_eval), s_eval)
   s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
   self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")
 
 def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
-def uconst(val): return UOp.const(val)
+def uconst(val): return UOp.const(dtypes.weakint, val)
 def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)
 
@@ -103,16 +104,16 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual(UOp.gcd(a, a*b, a*3).simplify(), a)
     self.assertEqual(UOp.gcd(a*a*a, a*b*a, a*3*a).simplify(), a*a)
     self.assertEqual(UOp.gcd(a*a*10, b*a*5, a*a*5).simplify(), a*5)
-    self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), uconst(5))
-    self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), uconst(1))
+    self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), a.const_like(5))
+    self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), a.const_like(1))
 
   def test_divides_exact(self):
     a = Variable("a", 1, 8)
     b = Variable("b", 1, 8)
     self.assertEqual((a*a*3).divide_exact(a).simplify(), a*3)
-    self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), uconst(1))
-    self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), uconst(2))
-    self.assertEqual((a*b*3).divide_exact(uconst(3)).simplify(), a*b)
+    self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), a.const_like(1))
+    self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), a.const_like(2))
+    self.assertEqual((a*b*3).divide_exact(a.const_like(3)).simplify(), a*b)
     self.assertEqual((a*a*3).divide_exact(a*(-3)).simplify(), a*-1)
     self.assertEqual((a*a*b*3).divide_exact(a*b).simplify(), a*3)
     self.assertEqual((a*3+a*b).divide_exact(a).simplify(), b+3)
@@ -882,7 +883,7 @@ class TestSymbolic(unittest.TestCase):
     idx = Variable("idx", 0, 24)
     self.helper_test_variable(idx//4, 0, 6, "(idx//4)")
     # TODO: simplify the true branch
-    self.helper_test_variable((idx<4).where(idx//4, uconst(-1)), -1, 6, "(idx<4).where((idx//4), -1)")
+    self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)")
 
   def test_floordiv_lt(self):
     # x//d<c <=> x<c*d for d>0, and <=> c*d<x for d<0
@@ -941,7 +942,7 @@ class TestSymbolic(unittest.TestCase):
 
   def test_where_removal(self):
     cond = Variable("a", 0, 3) < 2
-    u1, u0 = UOp.const(True), UOp.const(False)
+    u1, u0 = cond.const_like(True), cond.const_like(False)
     self.helper_test_variable(cond, 0, 1, "(a<2)")
     self.helper_test_variable(cond.where(u1, u0), 0, 1, "(a<2)")
     self.helper_test_variable(cond.where(u1, u0).where(u1, u0), 0, 1, "(a<2)")
@@ -1016,12 +1017,12 @@ class TestSymbolic(unittest.TestCase):
 
     # TODO: copied from render, render does not support cast
     glbl = UOp.param(0, dtypes.int, (1,))
-    uops = get_uops(UOp(Ops.STORE, src=(glbl.index(UOp.const(0, dtypes.int)), expr)).sink())
+    uops = get_uops(UOp(Ops.STORE, src=(glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
     rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
 
     # the vars are now scalar PARAMs
     pvar = {u.expr: u for u in rewritten_uop.toposort() if u.op is Ops.PARAM}
-    self.assertEqual(rewritten_uop, (pvar['s']<UOp.const(2, dtypes.int)).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))
+    self.assertEqual(rewritten_uop, (pvar['s']<2).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))
 
   def test_where_merge_branches(self):
     cond1 = Variable("s", 0, 10) < 6
@@ -1178,7 +1179,7 @@ class TestSymbolicVariables(unittest.TestCase):
     b = Variable("x", 1, 1).bind(1)
     s = b.simplify()
     self.assertEqual(s.op, Ops.CONST)
-    self.assertEqual(s.val, 1)
+    self.assertEqual(s.arg, 1)
 
 class TestSymInfer(unittest.TestCase):
   def test_sym_infer(self):
@@ -1202,13 +1203,9 @@ class TestSymInfer(unittest.TestCase):
     # floor: 1 % -1000 = -999, 1 // -1000 = -1
     assert sym_infer(a%b, var_vals) == -999
     assert sym_infer(a//b, var_vals) == -1
-  def test_sym_infer_with_cast(self):
-    a = Variable("a", 0, 100, dtypes.int)
-    assert sym_infer(a.cast(dtypes.long) + 1, {a.expr: 5}) == 6
-    assert sym_infer(a.cast(dtypes.float) * 0.5, {a.expr: 5}) == 2.5
   def test_sym_infer_with_bitcast(self):
     a = Variable("a", 1, 10, dtypes.int)
-    expr = ((a.bitcast(dtypes.uint) << UOp.const(1)).bitcast(dtypes.int) + 2)
+    expr = ((a.bitcast(dtypes.uint) << UOp.const(dtypes.uint, 1)).bitcast(dtypes.int) + 2)
     ret = sym_infer(expr, {a.expr: 2})
     assert isinstance(ret, int)
     assert ret == 6
@@ -1219,7 +1216,7 @@ class TestSymInfer(unittest.TestCase):
     c = Variable("c", 0, 0xFFFFFFFF, dtypes.uint)
     assert sym_infer(c.bitcast(dtypes.int), {c.expr: 0xFFFFFFFF}) == -1
 
-    assert sym_infer(UOp.const(1.5).cast(dtypes.float).bitcast(dtypes.uint), {}) == 1069547520
+    assert sym_infer(UOp.const(dtypes.float, 1.5).bitcast(dtypes.uint), {}) == 1069547520
 
   def test_sym_infer_deeply_nested(self):
     # build an expression that exceeds Python's nested parentheses limit for eval
@@ -1340,7 +1337,7 @@ class TestInvalidIndex(unittest.TestCase):
   def test_invalid_times_0(self):
     ridx = Variable("ridx", 0, 10)
     idx = (ridx<5).where(ridx, UOp.invalid())*0
-    self.assertIs(idx.simplify(), (ridx<5).where(uconst(0), UOp.invalid()),
+    self.assertIs(idx.simplify(), (ridx<5).where(UOp.const(dtypes.weakint, 0), UOp.invalid()),
                   "multiplying an index by 0 should preserve the invalid")
 
   def test_alu_moves_inside_invalid(self):
@@ -1365,24 +1362,24 @@ class TestInvalidIndex(unittest.TestCase):
     self.assertIs((UOp.invalid()<Variable("a",0,10)).simplify().dtype, dtypes.bool)
 
   def test_alu_invalid_vconst(self):
-    c1 = UOp.const((1, 1, Invalid, Invalid))
-    c2 = UOp.const((1, Invalid, 1, 1))
-    self.assertIs((c1+c2).simplify(), UOp.const((2, Invalid, Invalid, Invalid)))
+    c1 = UOp.const(dtypes.weakint, (1, 1, Invalid, Invalid))
+    c2 = UOp.const(dtypes.weakint, (1, Invalid, 1, 1))
+    self.assertIs((c1+c2).simplify(), UOp.const(dtypes.weakint, (2, Invalid, Invalid, Invalid)))
 
 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,))
-    index = buf.index(UOp.const(0))
+    index = buf.index(UOp.const(dtypes.weakint, 0))
     # Direct: store(idx, load(idx)) -> NOOP
     self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
     # Emergent: store(idx, load(idx) + 0) -> store(idx, load(idx)) -> NOOP
-    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(0)), sym).op, Ops.NOOP)
+    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(dtypes.int, 0)), sym).op, Ops.NOOP)
     # Emergent: store(idx, load(idx) * 1) -> store(idx, load(idx)) -> NOOP
-    self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(1)), sym).op, Ops.NOOP)
+    self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(dtypes.int, 1)), sym).op, Ops.NOOP)
     # Negative: store(idx, load(idx) + 1) should NOT fold
-    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(1)), sym).op, Ops.STORE)
+    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(dtypes.int, 1)), sym).op, Ops.STORE)
 
 class TestMoveWhereOnLoad(unittest.TestCase):
   def test_bool_index_preserves_dtype(self):
@@ -1393,7 +1390,7 @@ class TestMoveWhereOnLoad(unittest.TestCase):
     cond = (a < 4) & (r < 2)
     valid = (a < 2)  # pre-existing valid on the load (to pass can_move check for the r-only clause)
     idx = buf.index(a.valid(valid))
-    expr = cond.where(idx, UOp.const(0))
+    expr = cond.where(idx, idx.const_like(0))
     out = graph_rewrite(expr, pm_move_where_on_load)
     type_verify(out, spec_shared)  # Invalid matches any dtype
 
@@ -1478,7 +1475,7 @@ class TestFuzzFailure(unittest.TestCase):
     v2=Variable('v2', 0, 2)
     v3=Variable('v3', 0, 1)
     expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3)
-    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(0), v3.const_like(0)
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     assert num==rn, f"{num} != {rn}"

test/null/test_uop_vmin_vmax.py

@@ -5,12 +5,12 @@ from tinygrad.dtype import dtypes, Invalid
 class TestVminVmaxProperties(unittest.TestCase):
   def test_vmin_vmax_constant(self):
     # vmin and vmax for a constant
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
     self.assertEqual(uop.vmin, 42)
     self.assertEqual(uop.vmax, 42)
 
   def test_vmin_vmax_cmpne(self):
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
     def test_bool(u, x):
       self.assertEqual(u.vmin, x)
       self.assertEqual(u.vmax, x)
@@ -81,8 +81,8 @@ 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), dtype=dtypes.float)
+    x = UOp.const(dtypes.float, 0.0)
+    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(dtypes.int, 0), dtype=dtypes.float)
     uop = x * y
     # TODO: these should be 0, but definitely should not be nan
     self.assertEqual(uop.vmin, -math.inf)
@@ -167,7 +167,7 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertNotEqual(i.vmin, i.vmax)
 
   def test_vmin_vmax_invalid_vconst(self):
-    x = UOp.const((0, 4, Invalid, Invalid))
+    x = UOp.const(dtypes.weakint, (0, 4, Invalid, Invalid))
     self.assertEqual((x.vmin, x.vmax), (0, 4))
 
 class TestVminVmaxDivMod(unittest.TestCase):
@@ -198,14 +198,14 @@ class TestVminVmaxDivMod(unittest.TestCase):
 
   def test_vmin_vmax_floordiv_floormod(self):
     x = UOp.variable('x', -7, 7)
-    floordiv = x.alu(Ops.FLOORDIV, UOp.const(3))
+    floordiv = x.alu(Ops.FLOORDIV, x.const_like(3))
     self.assertEqual(floordiv.vmin, -3)
     self.assertEqual(floordiv.vmax, 2)
-    floormod = x.alu(Ops.FLOORMOD, UOp.const(3))
+    floormod = x.alu(Ops.FLOORMOD, x.const_like(3))
     self.assertEqual(floormod.vmin, 0)
     self.assertEqual(floormod.vmax, 2)
     # negative const divisor: floormod range is [c+1, 0]
-    floormod_neg = x.alu(Ops.FLOORMOD, UOp.const(-3))
+    floormod_neg = x.alu(Ops.FLOORMOD, x.const_like(-3))
     self.assertEqual(floormod_neg.vmin, -2)
     self.assertEqual(floormod_neg.vmax, 0)
 
@@ -286,44 +286,44 @@ class TestVminVmaxDivMod(unittest.TestCase):
 class TestVminVmaxVConst(unittest.TestCase):
   def test_vmin_vmax_vconst_single_element(self):
     # vmin and vmax for a single-element vector constant
-    uop = UOp.const((42,))
+    uop = UOp.const(dtypes.int32, (42,))
     self.assertEqual(uop.vmin, 42)
     self.assertEqual(uop.vmax, 42)
 
   def test_vmin_vmax_vconst_multiple_elements(self):
     # vmin and vmax for a multi-element vector constant
-    uop = UOp.const((10, 20, -5, 7))
+    uop = UOp.const(dtypes.int32, (10, 20, -5, 7))
     self.assertEqual(uop.vmin, -5)
     self.assertEqual(uop.vmax, 20)
 
   def test_vmin_vmax_vconst_all_equal(self):
     # vmin and vmax for a vector where all elements are equal
-    uop = UOp.const((7, 7, 7))
+    uop = UOp.const(dtypes.int32, (7, 7, 7))
     self.assertEqual(uop.vmin, 7)
     self.assertEqual(uop.vmax, 7)
 
   def test_vmin_vmax_vconst_with_negative_values(self):
     # vmin and vmax for a vector constant containing negative values
-    uop = UOp.const((-10, -20, -5, -15))
+    uop = UOp.const(dtypes.int32, (-10, -20, -5, -15))
     self.assertEqual(uop.vmin, -20)
     self.assertEqual(uop.vmax, -5)
 
   def test_vmin_vmax_vconst_with_floats(self):
     # vmin and vmax for a vector constant of float values
-    uop = UOp.const((1.5, -3.2, 0.0))
+    uop = UOp.const(dtypes.float32, (1.5, -3.2, 0.0))
     self.assertEqual(uop.vmin, -3.2)
     self.assertEqual(uop.vmax, 1.5)
 
   def test_vmin_vmax_vconst_with_bools(self):
     # vmin and vmax for a vector constant of bool values
-    uop = UOp.const((True, False, False))
+    uop = UOp.const(dtypes.bool, (True, False, False))
     self.assertIs(uop.vmin, False)
     self.assertIs(uop.vmax, True)
 
   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,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
     val = UOp(Ops.LOAD, src=(d1.index(idx),))
     uop = (val // 32)
     self.assertEqual(uop.vmin, -67108864)
@@ -332,17 +332,17 @@ class TestVminVmaxVConst(unittest.TestCase):
 class TestConstFactor(unittest.TestCase):
   def test_const_factor_constant(self):
     # const_factor for a constant
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
     self.assertEqual(uop.const_factor(), 42)
 
   def test_const_factor_addition(self):
     # const_factor for an addition of constants
-    uop = UOp.const(30) + UOp.const(12)
+    uop = UOp.const(dtypes.int32, 30) + UOp.const(dtypes.int32, 12)
     self.assertEqual(uop.const_factor(), 6)  # GCD(30, 12) = 6
 
   def test_const_factor_multiplication(self):
     # const_factor for a multiplication of constants
-    uop = UOp.const(5) * UOp.const(7)
+    uop = UOp.const(dtypes.int32, 5) * UOp.const(dtypes.int32, 7)
     self.assertEqual(uop.const_factor(), 5)  # For multiplication, it's one of the factors
 
   def test_const_factor_with_variable(self):
@@ -377,14 +377,14 @@ class TestConstFactor(unittest.TestCase):
 class TestDivides(unittest.TestCase):
   def test_divides_constant_exact(self):
     # Divides a constant by an exact divisor
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
     result = uop.divides(7)
     self.assertIsNotNone(result)
     self.assertEqual(result.const_factor(), 6)  # 42 / 7 = 6
 
   def test_divides_constant_inexact(self):
     # Try to divide a constant by a non-exact divisor
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
     result = uop.divides(5)
     self.assertIsNone(result)  # 42 is not divisible by 5
 

test/null/test_uops.py

@@ -5,24 +5,23 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
 from test.helpers import eval_uop, to_uops_list
 
 class TestDTypeFromUOp(unittest.TestCase):
   def test_broadcastable_promotion(self):
-    self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(1.0).cast(dtypes.float32), UOp.const(1.0).cast(dtypes.float16)), None), dtypes.float32)
-    self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(1).cast(dtypes.int8), UOp.const(1).cast(dtypes.int32)), None), dtypes.int32)
+    self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
+    self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.int32, 1)), None), dtypes.int32)
 
   def test_same_dtype_fast_path(self):
-    src = (UOp.const(1), UOp.const(2))
+    src = (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.weakint, 2))
     self.assertEqual(dtype_from_uop(Ops.ADD, src, None), dtypes.weakint)
 
   def test_where_promotion(self):
-    cond = UOp.const(True)
-    srcs = (cond, UOp.const(1.0).cast(dtypes.float32), UOp.const(1.0).cast(dtypes.float16))
-    self.assertEqual(dtype_from_uop(Ops.WHERE, srcs, None), dtypes.float32)
+    cond = UOp.const(dtypes.bool, True)
+    self.assertEqual(dtype_from_uop(Ops.WHERE, (cond, UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
     idx = UOp.range(4, 0)
     self.assertEqual(idx.valid(idx < 4).dtype, dtypes.weakint)
 
@@ -38,39 +37,33 @@ class TestDTypeFromUOp(unittest.TestCase):
     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
-    self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)
+    self.assertEqual(UOp.const(dtypes.int32, 3).dtype, dtypes.int32)
 
   def test_weak_dtype_rejected_by_program_spec(self):
     for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
-      with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program)
-      type_verify(UOp.const(value, concrete).sink(), spec_program)
-
-  def test_invalid_stated_dtype(self):
-    # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
-    # and the spec is what rejects a non-bool Invalid
-    self.assertIs(UOp.const(Invalid, dtypes.float32), UOp.invalid())
-    with self.assertRaises(RuntimeError): type_verify(UOp(Ops.CONST, dtypes.float32, arg=Invalid), spec_shared)
+      with self.assertRaises(RuntimeError): type_verify(UOp.const(weak, value).sink(), spec_program)
+      type_verify(UOp.const(concrete, value).sink(), spec_program)
 
   def test_invalid_dtype_and_consumers(self):
     invalid = UOp.invalid()
     self.assertIs(invalid.dtype, dtypes.bool)
-    self.assertIs(UOp.const(Invalid, dtypes.float32), invalid)
+    self.assertIs(UOp.const(dtypes.float32, Invalid), invalid)
     self.assertIs((moved:=invalid.reshape((1,))).cast(dtypes.float32), moved)
     scratch = Tensor.invalids(4, dtype=dtypes.float32)
     self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort()
-      if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
-    invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32)
+      if u.arg is Invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
+    invalid, value = UOp.invalid(), UOp.const(dtypes.float32, 1)
     for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)),
-              UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
+              UOp.const(dtypes.bool, 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)
     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))
+    self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(dtypes.float, 0))
     type_verify(out.sink(), spec_program)
 
   def test_remove_invalid_stack_lanes(self):
-    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(1, dtypes.half), UOp.invalid()))
+    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(dtypes.half, 1), UOp.invalid()))
     out = graph_rewrite(stack, pm_remove_invalid)
-    self.assertEqual(out.src, (UOp.const(1, dtypes.half), UOp.const(0, dtypes.half)))
+    self.assertEqual(out.src, (UOp.const(dtypes.half, 1), UOp.const(dtypes.half, 0)))
     type_verify(out.sink(), spec_program)
 
 class TestLowerIndexDtype(unittest.TestCase):
@@ -79,7 +72,7 @@ class TestLowerIndexDtype(unittest.TestCase):
     # width the offset bounds select (this one needs long)
     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)))
+    shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(dtypes.weakint, 4)))
     lowered = graph_rewrite(shrink.sink(), pm_lower_index_dtype)
     self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
     sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK)
@@ -118,20 +111,6 @@ class TestSafeCast(unittest.TestCase):
     self.assertEqual(a.cast(dtypes.int8).cast(dtypes.int64).simplify(), a.cast(dtypes.int64))
     self.assertEqual(a.cast(dtypes.int8).cast(dtypes.float).simplify(), a.cast(dtypes.float))
 
-class TestConstFloatEq(unittest.TestCase):
-  def test_nan_eq_ne_agree(self):
-    nan = dtypes.float32.const(math.nan)
-    self.assertTrue(nan == math.nan)
-    self.assertFalse(nan != math.nan)  # float.__ne__ would say True here
-    self.assertFalse(nan == Invalid)
-    self.assertTrue(nan != Invalid)    # __ne__ must defer to the reflected eq, not swallow NotImplemented
-
-  def test_matchers_agree_on_nan(self):
-    n = UOp.const(math.nan, dtypes.float32)
-    for compiled in (False, True):
-      pm = PatternMatcher([(UPat(Ops.CONST, arg=math.nan), lambda: True)], compiled=compiled)
-      self.assertTrue(pm.rewrite(n), f"{compiled=}")
-
 class TestExecALU(unittest.TestCase):
   def test_sqrt(self):
     self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
@@ -219,9 +198,9 @@ class TestGatedStoreRewrite(unittest.TestCase):
   def test_tiny_gate_store(self):
     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)))
-    val = UOp.const(42.0).cast(dtypes.float)
+    gate = gidx0<UOp.const(dtypes.weakint, 1)
+    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(dtypes.weakint, 2)).valid(gate)))
+    val = UOp.const(dtypes.float, 42.0)
     store = UOp(Ops.STORE, src=(idx, val))
     uops = to_uops_list([store])
     if_uop = next(u for u in uops if u.op is Ops.IF)
@@ -236,10 +215,10 @@ class TestGatedStoreRewrite(unittest.TestCase):
     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))))
+    idx = gidx0 * UOp.const(dtypes.weakint, 2)
+    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(dtypes.weakint, 1))))
     idx1 = UOp(Ops.INDEX, src=(gmem1, idx))
-    val = UOp.const(42.0).cast(dtypes.float)
+    val = UOp.const(dtypes.float, 42.0)
     stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
     uops = to_uops_list(stores)
     if_uop = next(u for u in uops if u.op is Ops.IF)
@@ -256,11 +235,11 @@ class TestGatedStoreRewrite(unittest.TestCase):
     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)
+    idx = gidx0*UOp.const(dtypes.weakint, 2)
+    gate = gidx0<UOp.const(dtypes.weakint, 1)
     idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gate)))
     idx1 = UOp(Ops.INDEX, src=(gmem1, idx.valid(gate)))
-    val = UOp.const(42.0).cast(dtypes.float)
+    val = UOp.const(dtypes.float, 42.0)
     stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
     uops = to_uops_list(stores)
     ifs = [u for u in uops if u.op is Ops.IF]
@@ -278,7 +257,7 @@ class TestFastIdiv(unittest.TestCase):
   def test_division_power_of_two(self):
     for dt in (dtypes.int32, dtypes.uint32):
       g = UOp.param(0, dt, (3,))
-      c = UOp.const(2).cast(dt)
+      c = UOp.const(dt, 2)
       l = g.index(c)
       a = UOp(Ops.CDIV, dt, (l, c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -291,7 +270,7 @@ class TestFastIdiv(unittest.TestCase):
     # 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,))
-      c = UOp.const(8).cast(dt)
+      c = UOp.const(dt, 8)
       a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
       ops = [x.op for x in uops]
@@ -303,7 +282,7 @@ class TestFastIdiv(unittest.TestCase):
     # uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
     for dt in (dtypes.uint32, dtypes.uint64):
       g = UOp.param(0, dt, (3,))
-      c = UOp.const(2).cast(dt)
+      c = UOp.const(dt, 2)
       a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
       ops = [x.op for x in uops]
@@ -315,7 +294,7 @@ class TestFastIdiv(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
   def test_fast_idiv_and_mod(self):
     g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(3).cast(dtypes.uint)
+    c = UOp.const(dtypes.uint, 3)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
     uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -335,7 +314,7 @@ class TestFastIdiv(unittest.TestCase):
   def test_fast_idiv_bounded_numerator_zero(self):
     x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
     for val in range(2):
-      self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3).cast(x.dtype)), vals=(val,)), cdiv(val, 3))
+      self.assertEqual(eval_uop(x.alu(Ops.CDIV, x.const_like(3)), vals=(val,)), cdiv(val, 3))
 
   @Context(DISABLE_FAST_IDIV=0)
   def test_fast_idiv_remove_powers_of_two(self):
@@ -350,7 +329,7 @@ class TestFastIdiv(unittest.TestCase):
   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,))
-    c = UOp.const(7).cast(dtypes.uint)
+    c = UOp.const(dtypes.uint, 7)
     l = UOp(Ops.LOAD, src=(g.index(c),))
     a = UOp(Ops.CDIV, src=(l, c))
     uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -361,7 +340,7 @@ class TestFastIdiv(unittest.TestCase):
 
   def test_disable_fast_idiv(self):
     g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(3).cast(dtypes.uint)
+    c = UOp.const(dtypes.uint, 3)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
     with Context(DISABLE_FAST_IDIV=1):
@@ -373,8 +352,8 @@ class TestFastIdiv(unittest.TestCase):
 class TestUOpMethod(unittest.TestCase):
   @unittest.skip("uops lt no longer ordered")
   def test_compare_alu_same_src_different_arg(self):
-    a = UOp.const(2.0)
-    b = UOp.const(3.0)
+    a = UOp.const(dtypes.float, 2.0)
+    b = UOp.const(dtypes.float, 3.0)
 
     add = UOp(Ops.ADD, src=(a, b))
     mul = UOp(Ops.MUL, src=(a, b))
@@ -389,8 +368,8 @@ class TestUOpMethod(unittest.TestCase):
     self.assertEqual(list(var_vals)[0], a.expr)
 
   def test_const_factor(self):
-    gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(8),), arg='gidx0')
-    self.assertEqual(UOp.const(17).const_factor(), 17)
+    gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 8),), arg='gidx0')
+    self.assertEqual(UOp.const(dtypes.int, 17).const_factor(), 17)
     self.assertEqual(gidx0.const_factor(), 1)
     self.assertEqual((gidx0*3).const_factor(), 3)
     self.assertEqual((gidx0*3+6).const_factor(), 3)
@@ -409,26 +388,26 @@ class TestUOpMethod(unittest.TestCase):
 
   def test_const_zero_neg_zero_different(self):
     # -0.0 and 0.0 must be different UOps (for IEEE754 correctness, e.g. 1/-0.0 = -inf)
-    pos_zero = UOp.const(0.0)
-    neg_zero = UOp.const(-0.0)
+    pos_zero = UOp.const(dtypes.float, 0.0)
+    neg_zero = UOp.const(dtypes.float, -0.0)
     self.assertIsNot(pos_zero, neg_zero)
     self.assertNotEqual(hash(pos_zero.arg), hash(neg_zero.arg))
 
   def test_const_nan_same(self):
     # nan constants should be deduplicated
-    nan1 = UOp.const(float('nan'))
-    nan2 = UOp.const(float('nan'))
+    nan1 = UOp.const(dtypes.float, float('nan'))
+    nan2 = UOp.const(dtypes.float, float('nan'))
     self.assertIs(nan1, nan2)
 
 class TestUOpStr(unittest.TestCase):
   def test_uop_str(self):
-    a = UOp.const(2.0) + UOp.const(3.0)
+    a = UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 3.0)
     for _ in range(20): a = a + a
     assert len(str(a)) < 10_000, "exponential string growth"
     assert str(eval(str(a))) == str(a)
 
   def test_vectorized_str(self):
-    vec = UOp(Ops.STACK, src=tuple(UOp.const(x) for x in range(4)))
+    vec = UOp(Ops.STACK, src=tuple(UOp.const(dtypes.int, x) for x in range(4)))
     assert str(eval(str(vec))) == str(vec)
 
   def test_reduceop_arg(self):
@@ -463,16 +442,16 @@ class TestUOpRender(unittest.TestCase):
     u = UOp(Ops.STACK, dtype=dtypes.void, src=())
     self.assertEqual(u.render(), "{}")
   def test_render_vectorize_same(self):
-    u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
+    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
     self.assertEqual(u.render(simplify=False), "{0,0,0}")
   def test_render_vectorize_different(self):
-    u = UOp(Ops.STACK, src=tuple(UOp.const(i) for i in range(3)))
+    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
     self.assertEqual(u.render(simplify=False), "{0,1,2}")
   def test_render_vectorize_same_simplified(self):
-    u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
+    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
     self.assertEqual(u.render(), "{0,0,0}")
   def test_render_vectorize_different_simplified(self):
-    u = UOp(Ops.STACK, src=tuple(UOp.const(i) for i in range(3)))
+    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
     self.assertEqual(u.render(), "{0,1,2}")
 
 class TestContiguousViewOffset(unittest.TestCase):

test/null/test_validate_oob.py

@@ -12,12 +12,12 @@ class TestValidateOOB(unittest.TestCase):
   def test_const_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.const(0)).load(dtype=dtypes.int)])  # valid
-      to_uops_list([buf.index(UOp.const(15)).load(dtype=dtypes.int)])  # valid (last element)
+      to_uops_list([buf.index(UOp.const(dtypes.int, 0)).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(UOp.const(dtypes.int, 15)).load(dtype=dtypes.int)])  # valid (last element)
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(16)).load(dtype=dtypes.int)])  # off by one
+        to_uops_list([buf.index(UOp.const(dtypes.int, 16)).load(dtype=dtypes.int)])  # off by one
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(42)).load(dtype=dtypes.int)])  # way out
+        to_uops_list([buf.index(UOp.const(dtypes.int, 42)).load(dtype=dtypes.int)])  # way out
 
   def test_variable_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
@@ -138,19 +138,11 @@ class TestValidateOOB(unittest.TestCase):
       with self.assertRaises(RuntimeError):
         to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])  # oob
 
-  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,))
-      shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4)))
-      ld0 = shrink.load(dtype=dtypes.int).index(0)
-      to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])
-
   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,))
-      gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0")
+      gidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.weakint, 16),), arg="gidx0")
       ld_bool = buf_bool.index(gidx).load()
       with self.assertRaises(RuntimeError):
         to_uops_list([buf_int.index(gidx.valid(ld_bool)).load()])  # gidx 0..15, buf_int size 8
@@ -164,12 +156,12 @@ class TestValidateOOB(unittest.TestCase):
       sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)
 
       # Define indices, valids and barrier
-      gidx = UOp(Ops.SPECIAL, src=(UOp.const(416),), arg="gidx0")
-      lidx = UOp(Ops.SPECIAL, src=(UOp.const(10),), arg="lidx0")
+      gidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 416),), arg="gidx0")
+      lidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 10),), arg="lidx0")
 
       gate = (gidx<400) & (lidx<8)
 
-      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(1))
+      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(dtypes.uint, 1))
 
       barrier = UOp(Ops.BARRIER, src=(local_store,))
       if_barrier = UOp(Ops.IF, src=(gate, barrier))
@@ -187,7 +179,7 @@ class TestValidateOOB(unittest.TestCase):
       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))))
+      ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx, ridx<16&mask))))
       to_uops_list([ld0])
 
 if __name__ == "__main__":

test/null/test_viz.py

@@ -96,10 +96,10 @@ class TestViz(unittest.TestCase):
   def test_exceptions(self):
     # VIZ tracks rewrites up to and including the error
     def count_3(x:UOp):
-      assert x.val <= 3
-      return x.replace(arg=x.val+1)
+      assert x.arg <= 3
+      return x.replace(arg=x.arg+1)
     err_pm = PatternMatcher([(UPat.cvar("x"), count_3),])
-    a = UOp.const(1)
+    a = UOp.const(dtypes.int, 1)
     with save_viz() as viz:
       with self.assertRaises(AssertionError): exec_rewrite(a, [err_pm])
     lst = viz.list_items()
@@ -199,8 +199,8 @@ class TestViz(unittest.TestCase):
     self.assertEqual(ansistrip(a2["label"]), "CUSTOM\nx\nyzww\nw")
 
   def test_inf_loop(self):
-    a = UOp.const(3)
-    b = UOp.const(4)
+    a = UOp.const(dtypes.int, 3)
+    b = UOp.const(dtypes.int, 4)
     pm = PatternMatcher([
       (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
       (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
@@ -226,7 +226,7 @@ class TestViz(unittest.TestCase):
   def test_enter_calls_rewrite(self):
     pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))])
     with save_viz() as viz:
-      inner = UOp.const(3)
+      inner = UOp.const(dtypes.int, 3)
       call = UOp(Ops.CALL, src=(UOp(Ops.SINK, src=(inner,)),))
       func = UOp(Ops.FUNCTION, src=(UOp(Ops.TUPLE, src=(call,)),))
       graph_rewrite(func, TrackedPatternMatcher(pm.patterns), enter_calls=True)
@@ -236,8 +236,8 @@ class TestViz(unittest.TestCase):
   def test_const_node_visibility(self):
     with save_viz() as viz:
       a = UOp.variable("a", 0, 10, dtype=dtypes.int)
-      z = UOp.const(0, a.dtype)
-      y = UOp.const(math.pi, dtypes.float)
+      z = UOp.const(a.dtype, 0)
+      y = UOp.const(dtypes.float, math.pi)
       alu = a*z
       ret = exec_rewrite(sink:=UOp.sink(alu, y), [sym])
     lst = viz.list_items()
@@ -253,7 +253,7 @@ class TestViz(unittest.TestCase):
 
   def test_const_reshape_expand_folded(self):
     # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes
-    c = UOp.const(1.0).expand((3,4))  # creates CONST->EXPAND chain
+    c = UOp.const(dtypes.float, 1.0, shape=(3,4))  # creates CONST->EXPAND chain
     a = UOp.variable("a", 0.0, 10.0, dtypes.float)
     alu = a + c
     with save_viz() as viz:
@@ -267,13 +267,13 @@ class TestViz(unittest.TestCase):
 
   def test_stack_movement_not_folded_unless_all_const(self):
     a = UOp.variable("a", 0, 10, dtype=dtypes.int)
-    c = UOp.const(1)
+    c = UOp.const(dtypes.int, 1)
     stack = a.stack(c)
     reshaped = stack.reshape((1, 2))
     graph = uop_to_json(VizData(), reshaped)
     self.assertFalse(graph[id(stack)]["exclude"])
 
-    const_stack = c.stack(UOp.const(2))
+    const_stack = c.stack(UOp.const(dtypes.int, 2))
     const_reshaped = const_stack.reshape((1, 2))
     const_graph = uop_to_json(VizData(), const_reshaped)
     self.assertTrue(const_graph[id(const_stack)]["exclude"])
@@ -401,7 +401,7 @@ class TestVizIntegration(unittest.TestCase):
     with save_viz() as viz:
       def test(root):
         return graph_rewrite(root, sym)
-      test(c:=UOp.const(1))
+      test(c:=UOp.const(dtypes.int, 1))
       test(c+1)
     ls = viz.list_items()
     self.assertEqual(len(ls), 1)
@@ -414,7 +414,7 @@ class TestVizIntegration(unittest.TestCase):
       @track_rewrites()
       def test(root):
         return graph_rewrite(root, sym)
-      test(c:=UOp.const(1))
+      test(c:=UOp.const(dtypes.int, 1))
       test(c+1)
     ls = viz.list_items()
     self.assertEqual(len(ls), 2)
@@ -425,22 +425,21 @@ class TestVizIntegration(unittest.TestCase):
     with save_viz() as viz:
       def default_test(root): return graph_rewrite(root, sym)
       tracked_test = track_rewrites()(default_test)
-      c = UOp.const(1)
+      c = UOp.const(dtypes.int, 1)
       default_test(c+1) # goes to the default group
       tracked_test(c)   # all rewrites after this go inside the second group.
       default_test(c+2)
     ls = viz.list_items()
     self.assertEqual(len(ls), 2)
     graph = next(viz.get_details(0, 0))["graph"]
-    # both operands of c+1 are the same bare weak CONST, so the graph has two nodes
     self.assertEqual(list(graph), [id(c), id(c+1)])
     self.assertTrue(graph[id(c)]["exclude"])
     self.assertFalse(graph[id(c+1)]["exclude"])
     self.assertEqual(list(next(viz.get_details(1, 0))["graph"]), [id(c)])
     graph = next(viz.get_details(1, 1))["graph"]
-    self.assertEqual(list(graph), [id(c), id((c+2).src[1]), id(c+2)])
+    self.assertEqual(list(graph), [id(c), id(c.const_like(2)), id(c+2)])
     self.assertTrue(graph[id(c)]["exclude"])
-    self.assertTrue(graph[id((c+2).src[1])]["exclude"])
+    self.assertTrue(graph[id(c.const_like(2))]["exclude"])
     self.assertFalse(graph[id(c+2)]["exclude"])
 
   def test_recurse(self):

test/opt/test_tensor_cores.py

@@ -81,13 +81,6 @@ class TestTensorCores(unittest.TestCase):
     for tc in Device[Device.DEFAULT].renderer.tensor_cores:
       helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
 
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
-  def test_tensor_cores_nested_reduce(self):
-    tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
-    a, b = Tensor.empty(tc.dims[1]*2, tc.dims[2], dtype=tc.dtype_in), Tensor.empty(tc.dims[2], tc.dims[0], dtype=tc.dtype_in)
-    ast = replace_opts(a.matmul(b, dtype=tc.dtype_out).sum(0).schedule_linear().src[-1].src[0], [Opt(OptOps.TC, 0, (-1, 0, 1))])
-    with self.assertRaises(KernelOptError): to_program(ast, Device[Device.DEFAULT].renderer)
-
   @Context(ALLOW_TF32=1)
   @unittest.skipIf(Device.DEFAULT == "PYTHON", "not generated on EMULATED device")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")

test/unit/test_assign.py

@@ -661,7 +661,7 @@ class TestAssign(unittest.TestCase):
 
   def test_assign_deviceless_const(self):
     s = Tensor.empty(4, device="CPU:1", dtype=dtypes.float)
-    s.assign(Tensor(UOp.const(2.0).cast(dtypes.float)))
+    s.assign(Tensor(UOp.const(dtypes.float, 2.0)))
     np.testing.assert_equal(s.numpy(), [2, 2, 2, 2])
 
   def test_nested_after_contiguous_store(self):

test/unit/test_dtype_weak.py

@@ -3,8 +3,7 @@ import tempfile, unittest, math
 from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
 from tinygrad.dtype import least_upper_float
-from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite, pm_lower_index_dtype, pm_commit_weak
-from tinygrad.uop.symbolic import symbolic_simple
+from tinygrad.uop.ops import UOp, Ops, dtype_from_uop
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.engine.jit import JitError
 
@@ -12,18 +11,18 @@ from tinygrad.engine.jit import JitError
 class TestWeakPromotion(unittest.TestCase):
   def test_rand_requires_concrete(self):
     with self.assertRaises(ValueError): Tensor.rand(2, dtype=dtypes.weakfloat)
-    with self.assertRaises(ValueError): Tensor.const(1.0).rand_like()
-    with self.assertRaises(ValueError): Tensor.const(1.0).randn_like()
+    with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).rand_like()
+    with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).randn_like()
 
   def test_reduce_strips_weakness(self):
     for weak, value, strong in ((dtypes.weakint, 1, dtypes.default_int), (dtypes.weakfloat, 1.0, dtypes.default_float)):
-      t = Tensor.const(value, weak).expand(3)
+      t = Tensor.const(weak, value).expand(3)
       for out in (t.sum(), t.max(), t.prod(), t.cumsum(0), t.cummax(0)[0]): self.assertEqual(out.dtype, strong)
-    self.assertEqual((Tensor.const(1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float32)
+    self.assertEqual((Tensor.const(dtypes.weakfloat, 1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float32)
 
   def test_materialize_at_default_dtype(self):
     for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),):
-      t = Tensor.const(value, weak)
+      t = Tensor.const(weak, value)
       self.assertEqual(t.dtype, weak)
       self.assertEqual(t.data().itemsize, strong.itemsize)
       self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize)
@@ -33,7 +32,7 @@ class TestWeakPromotion(unittest.TestCase):
       self.assertEqual(t.contiguous().dtype, weak)
 
   def test_assign_into_weak_commits(self):
-    t = Tensor.const(0.5)
+    t = Tensor.const(dtypes.weakfloat, 0.5)
     t.assign(Tensor(1.0, dtype=dtypes.default_float))
     self.assertEqual((t.dtype, t.item()), (dtypes.default_float, 1.0))
 
@@ -44,59 +43,10 @@ class TestWeakPromotion(unittest.TestCase):
     r = Tensor([2], dtype=dtypes.uint8, device="CPU").copysign(Tensor([1], dtype=dtypes.uint32, device="CPU"))
     self.assertEqual((r.dtype, r.tolist()), (dtypes.uint32, [2]))
 
-  def test_minimum_commits_both_operands(self):
-    r = Tensor(1).minimum(Tensor([2], dtype=dtypes.uint8, device="CPU"))
-    self.assertEqual((r.dtype, r.tolist()), (dtypes.uint8, [1]))
-
-  def test_broadcasted_keeps_const_weak(self):
-    # a python scalar stays a bare weak CONST through _broadcasted, lifted only to the KIND of the lub
-    x, y = Tensor([1], dtype=dtypes.int8)._broadcasted(3)
-    self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.CONST, dtypes.weakint, dtypes.int8))
-    x, y = Tensor([1], dtype=dtypes.int8)._broadcasted(0.5)
-    self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.CONST, dtypes.weakfloat, dtypes.weakfloat))
-    x, y = Tensor.const(1).reshape(1)._broadcasted(Tensor([1.0], dtype=dtypes.float32))
-    self.assertEqual((x._uop.base.op, x._uop.base.val, x.dtype, x.shape, y.dtype),
-                     (Ops.CONST, 1, dtypes.weakfloat, (1,), dtypes.float32))
-
-  def test_uop_scalar_const_lifts_kind(self):
-    for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
-                                                 (dtypes.int32, 1, dtypes.int32, dtypes.weakint),
-                                                 (dtypes.int32, 0.5, dtypes.weakfloat, dtypes.weakfloat),
-                                                 (dtypes.float32, 1, dtypes.float32, dtypes.weakfloat)):
+  def test_uop_scalar_const_unchanged(self):
+    for dtype, value in ((dtypes.weakint, 1), (dtypes.int32, 1), (dtypes.float32, 0.5)):
       out = UOp.variable("x", 0.0 if dtype == dtypes.float32 else 0, 10.0 if dtype == dtypes.float32 else 10, dtype) + value
-      self.assertEqual((out.dtype, out.src[1].op, out.src[1].dtype), (out_dtype, Ops.CONST, const_dtype))
-    # the kind lift converts the VALUE too (the arg is the only dtype carrier once UOp.const loses its dtype arg),
-    # and a bare weak const UOp is the same spelling as the python scalar: both lift to the same node
-    x = UOp.variable("x", 0.0, 1.0, dtypes.float32)
-    self.assertIsInstance((x + 2).src[1].val, float)
-    self.assertIs(x + UOp.const(2), x + 2)
-
-  def test_index_dtype_ignores_weakness(self):
-    with Context(SPEC=2):
-      idx = UOp.const(0).cast(dtypes.int32)
-      weak = UOp.const(1.0).expand((1,))
-      self.assertEqual(UOp(Ops.INDEX, dtypes.float32, (weak, idx)).dtype, dtypes.float32)
-      with self.assertRaisesRegex(RuntimeError, "bad dtype"): UOp(Ops.INDEX, dtypes.int32, (weak, idx))
-
-  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))
-      gate = UOp.const(True)
-      out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_lower_index_dtype, ctx={})
-    # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it
-    self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0, dtypes.bfloat16), gate))
-
-  def test_weak_srcs_commit_only_at_a_concrete_lub(self):
-    weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0)))
-    self.assertIs(graph_rewrite(weak_lub, pm_lower_index_dtype, ctx={}), weak_lub)
-    concrete = UOp.const(2.0).cast(dtypes.float16)
-    where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_lower_index_dtype, ctx={})
-    self.assertEqual(tuple(x.dtype for x in where.src), (dtypes.bool, dtypes.float16, dtypes.float16))
-
-  def test_weak_shift_lhs_commits_the_node(self):
-    # a shift derives its lhs's dtype, so committing the lhs restates the root (WGSL's packed store writes `mask << shift_am`)
-    shl = graph_rewrite(UOp.const(0xFFFF) << UOp.variable("x", 0, 16, dtypes.uint), symbolic_simple+pm_commit_weak)
-    self.assertEqual((shl.dtype, shl.src[0]), (dtypes.uint, UOp.const(0xFFFF, dtypes.uint)))
+      self.assertEqual((out.dtype, out.src[1].dtype), (dtype, dtype))
 
   @unittest.expectedFailure  # TODO: a weak const defers to its consumer (JAX): these dtypes change once python scalars are weak consts
   def test_changed_rows(self):
@@ -129,17 +79,17 @@ class TestWeakPromotion(unittest.TestCase):
   def test_weak_int_binop(self):
     v = UOp.variable("i", 0, 10, dtypes.weakint)
     self.assertEqual((v << 1).dtype, dtypes.weakint)
-    self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(1, dtypes.int8), UOp.const(1, dtypes.uint32)), None), dtypes.int8)
-    self.assertEqual(UOp.const(1).alu(Ops.SHL, UOp.const(1, dtypes.uint)).dtype, dtypes.weakint)
+    self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.uint32, 1)), None), dtypes.int8)
+    self.assertEqual(UOp.const(dtypes.weakint, 1).alu(Ops.SHL, UOp.const(dtypes.uint, 1)).dtype, dtypes.weakint)
     self.assertEqual((v & 3).dtype, dtypes.weakint)
-    with self.assertRaises(RuntimeError): Tensor.const(1.0) << Tensor.const(1.0)
-    with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64))
+    with self.assertRaises(RuntimeError): Tensor.const(dtypes.weakfloat, 1.0) << Tensor.const(dtypes.weakfloat, 1.0)
+    with self.assertRaises(RuntimeError): UOp.const(dtypes.int32, 1).alu(Ops.SHL, UOp.const(dtypes.float64, 1))
     for op in (Ops.SHL, Ops.SHR):
       with self.assertRaises(RuntimeError):
-        UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32))
+        UOp.const(dtypes.float32, 1).alu(op, UOp.const(dtypes.int32, 1))
     # float bitwise builds, the spec rejects it
     with Context(SPEC=1):
-      f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0)
+      f32, wf = UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.weakfloat, 1.0)
       for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, dtypes.float32, (f32, f32)), UOp(Ops.AND, dtypes.int32, (wf, wf))):
         with self.assertRaises(RuntimeError): type_verify([bad], spec_shared)
 
@@ -178,7 +128,7 @@ class TestWeakPromotion(unittest.TestCase):
 class TestWeakStorageBoundary(unittest.TestCase):
   # weak has no storage: a weak assignment source casts when it defers to the destination, everything else raises
   def test_weak_source(self):
-    w05 = Tensor.const(0.5).reshape(1)
+    w05 = Tensor.const(dtypes.weakfloat, 0.5).reshape(1)
     dst = Tensor.zeros(2, dtype=dtypes.int8, device="CPU").contiguous().realize()
     with self.assertRaises(RuntimeError): dst.assign(w05.expand(2))                   # weakfloat into int does not defer
     with self.assertRaises(RuntimeError): dst[0:1] = w05
@@ -199,7 +149,7 @@ class TestWeakMaterializationEntries(unittest.TestCase):
   def test_reads_commit_storage_raises(self):
     for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),):
       def weak_val():
-        return Tensor([True], device="CPU").where(Tensor.const(value, weak), Tensor.const(value, weak))
+        return Tensor([True], device="CPU").where(Tensor.const(weak, value), Tensor.const(weak, value))
       self.assertEqual(weak_val().dtype, weak)
       self.assertEqual(weak_val().to("CPU").dtype, weak)
       self.assertEqual(weak_val().data().format, strong.fmt)
@@ -215,7 +165,7 @@ class TestWeakMaterializationEntries(unittest.TestCase):
   def test_weak_is_virtual(self):
     # NOTE: int64 lub uint64 is weakfloat, so this is device-ful weak from promotion, never from a cast to weak
     devful = Tensor([1], dtype=dtypes.int64, device="CPU") + Tensor([1], dtype=dtypes.uint64, device="CPU")
-    for t in (Tensor.const(0.5), devful):
+    for t in (Tensor.const(dtypes.weakfloat, 0.5), devful):
       self.assertTrue(t.uop.is_virtual)
       # realize is a no-op, so a weak input can never become the real buffer TinyJit needs
       with self.assertRaises(JitError): TinyJit(lambda x: (x+1).realize())(t)
@@ -226,7 +176,7 @@ class TestWeakMaterializationEntries(unittest.TestCase):
 
   def test_empty_reads_commit(self):
     for weak, strong in ((dtypes.weakfloat, dtypes.default_float),):
-      empty = Tensor.const(0, weak).reshape(1).shrink(((0, 0),))
+      empty = Tensor.const(weak, 0).reshape(1).shrink(((0, 0),))
       self.assertEqual(empty.data().format, strong.fmt)
       self.assertEqual(empty.numpy().dtype.itemsize, strong.itemsize)
       self.assertEqual(empty.tolist(), [])

test/unit/test_function.py

@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from tinygrad.function import function
 from tinygrad import Tensor, GlobalCounters, Device
-from tinygrad.dtype import Invalid
+from tinygrad.dtype import dtypes, Invalid
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
 
 class TestFunction(unittest.TestCase):
@@ -593,7 +593,7 @@ class TestFunctionTuple(unittest.TestCase):
     state = Tensor([10., 20., 30., 40.], device="CPU").contiguous().realize()
     @function(precompile=True, allow_implicit=True)
     def f(a:Tensor):
-      after = state.uop.after(state.uop.shrink(((0, 2),)).store(Invalid))
+      after = state.uop.after(state.uop.shrink(((0, 2),)).store(UOp.const(dtypes.float32, Invalid, shape=(2,))))
       return Tensor(after).contiguous() + a
     out = f(Tensor([1., 1., 1., 1.], device="CPU").contiguous().realize())
     np.testing.assert_allclose(out.numpy(), [11., 21., 31., 41.])

test/unit/test_invalid_tensor.py

@@ -134,14 +134,14 @@ class TestInvalidTensor(unittest.TestCase):
     self._invalid_test_helper(out, [1.0, 2.0, None, None])
 
   def test_uop_where_keeps_invalid_bare(self):
-    cond = UOp.const(0) < UOp.const(1)
-    idx = UOp(Ops.STACK, src=tuple(UOp.const(x) for x in range(3)))
+    cond = UOp.const(dtypes.weakint, 0) < UOp.const(dtypes.weakint, 1)
+    idx = UOp(Ops.STACK, src=tuple(UOp.const(dtypes.weakint, x) for x in range(3)))
     out = cond.where(idx, UOp.invalid())
     self.assertIs(cond.op, Ops.CMPLT)
     self.assertIs(idx.op, Ops.STACK)
     self.assertIs(out.op, Ops.WHERE)
     self.assertIs(out.src[2].op, Ops.CONST)
-    self.assertTrue(out.src[2].is_invalid)
+    self.assertIs(out.src[2].arg, Invalid)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_jit.py

@@ -351,13 +351,13 @@ class TestJit(unittest.TestCase):
     @TinyJit
     def f(x:Tensor) -> Tensor: return (x + 1).realize()
     with self.assertRaises(JitError):
-      f(Tensor(UOp.const(2.0).cast(dtypes.float))).item()
+      f(Tensor(UOp.const(dtypes.float, 2.0))).item()
 
   def test_jit_deviceless_compute_input(self):
     @TinyJit
     def f(x:Tensor) -> Tensor: return (x + 1).realize()
     with self.assertRaises(JitError):
-      f(Tensor(UOp.const(2.0).cast(dtypes.float) + UOp.const(1.0).cast(dtypes.float))).item()
+      f(Tensor(UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 1.0))).item()
 
   def test_jit_init_empty_alt(self):
     @TinyJit

test/unit/test_llm_server.py

@@ -6,8 +6,6 @@ from tinygrad.llm.model import Transformer, TransformerConfig
 
 TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
                            norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32)
-V_START_POS = UOp.variable("start_pos", 0, TEST_CONFIG.max_context-1)
-V_TOKS = UOp.variable("toks", 1, 32)  # 32 is the default chunk_size in generate
 
 class TestTransformerGenerate(unittest.TestCase):
   def test_kv_cache_reuse(self):
@@ -16,7 +14,7 @@ class TestTransformerGenerate(unittest.TestCase):
 
     captured_inputs = []
     def mock_call(self, tokens, start_pos, temperature):
-      captured_inputs.append((tokens.shape, start_pos))
+      captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val))
       return Tensor([[42]])
 
     with patch.object(Transformer, '__call__', mock_call):
@@ -33,7 +31,9 @@ class TestTransformerGenerate(unittest.TestCase):
       next(gen)
 
     # should process tokens[6:] = [42, 10, 11, 12] since first 6 have cached k/v
-    self.assertEqual(captured_inputs, [((1, V_TOKS.bind(4)), V_START_POS.bind(6))])
+    toks_shape = captured_inputs[0][0][-1]
+    self.assertEqual(toks_shape.val if isinstance(toks_shape, UOp) else toks_shape, 4)
+    self.assertEqual(captured_inputs[0][1], 6)
 
   def test_kv_cache_invalidation(self):
     """Test that generate invalidates the KV cache when tokens diverge from the cached prefix."""
@@ -41,7 +41,7 @@ class TestTransformerGenerate(unittest.TestCase):
 
     captured_inputs = []
     def mock_call(self, tokens, start_pos, temperature):
-      captured_inputs.append((tokens.shape, start_pos))
+      captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val))
       return Tensor([[42]])
 
     with patch.object(Transformer, '__call__', mock_call):
@@ -55,7 +55,9 @@ class TestTransformerGenerate(unittest.TestCase):
       next(gen)
 
     # should process all 3 tokens from start
-    self.assertEqual(captured_inputs, [((1, V_TOKS.bind(3)), V_START_POS.bind(0))])
+    toks_shape = captured_inputs[0][0][-1]
+    self.assertEqual(toks_shape.val if isinstance(toks_shape, UOp) else toks_shape, 3)
+    self.assertEqual(captured_inputs[0][1], 0)
 
   def test_two_prompts_schedule_cache(self):
     """Third prompt should hit the schedule cache, not miss (first two warm up both jits: prefill + decode)."""

test/unit/test_metal_graph.py

@@ -17,7 +17,7 @@ class TestMetalGraph(unittest.TestCase):
       buf.op = Ops.SLICE
       src = MagicMock()
       src.dtype = dtypes.uint8
-      buf.src = (src, UOp.const(offset))
+      buf.src = (src, UOp.const(dtypes.weakint, offset))
       buf.dtype = dtypes.uint8
     else:
       buf.op = Ops.BUFFER

test/unit/test_multitensor.py

@@ -60,8 +60,8 @@ class TestMultiTensor(unittest.TestCase):
   def test_shard_elementwise(self): self._test_shard_op(lambda t:(t+t).reshape(2, 2), [[2.,2.],[2.,2.]])
   def test_alu_deviceless_const(self):
     s = Tensor([1.0, 2, 3, 4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0)
-    np.testing.assert_equal((s + Tensor(UOp.const(1.0).cast(dtypes.float))).numpy(), [2, 3, 4, 5])
-    np.testing.assert_equal((s + Tensor(UOp.const(1.0).cast(dtypes.float)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5])
+    np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0))).numpy(), [2, 3, 4, 5])
+    np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5])
 
   def test_add_rank_expand_shard(self):
     # a sharded src keeps its own rank under implicit broadcast, its shard axis right-aligns into the output

test/unit/test_tensor_data.py

@@ -66,8 +66,8 @@ class TestTensorData(unittest.TestCase):
     assert dat.shape == ()
 
   def test_const_dtype_for_uop(self):
-    self.assertEqual(Tensor.const(UOp.const(1.0).cast(dtypes.float32), dtypes.int8).dtype, dtypes.int8)
-    self.assertEqual(Tensor.const(UOp.variable("x", 1, 10).bind(5), dtypes.int32).item(), 5)
+    self.assertEqual(Tensor.const(dtypes.int8, UOp.const(dtypes.float32, 1.0)).dtype, dtypes.int8)
+    self.assertEqual(Tensor.const(dtypes.int32, UOp.variable("x", 1, 10).bind(5)).item(), 5)
 
   def test_data_float32(self):
     a = Tensor([[1,2.5],[3,4]], dtype=dtypes.float32)

tinygrad/callify.py

@@ -57,11 +57,11 @@ def _make_buffer_view(src:UOp) -> UOp|None:
   if (offset := src.contiguous_view_offset()) is None: return None
   buf = src.base
   if buf.op is Ops.SLICE:
-    byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize
+    byte_offset = buf.src[1].arg * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize
     buf = buf.src[0]
     if byte_offset % buf.dtype.itemsize != 0: return None
     offset = byte_offset // buf.dtype.itemsize
-  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel())
+  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(None, offset)), src.numel())
 
 def contiguous_mops_to_view(c:UOp, src:UOp):
   """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""

tinygrad/codegen/__init__.py

@@ -3,12 +3,12 @@ import itertools, functools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
 from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
-from tinygrad.uop.ops import AxisType, pm_commit_weak, pm_cast_weak
+from tinygrad.uop.ops import AxisType
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
-from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.dtype import dtypes, AddrSpace, Invalid
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
@@ -84,7 +84,7 @@ def expand_wmma(ctx:dict[int, int], u:UOp):
 expander2 = PatternMatcher([
   (UPat(Ops.REDUCE, name="r"), expand_reduce),
   (UPat(Ops.RANGE, name="r"),
-   lambda ctx, r: UOp.const(tuple(range(r.vmax+1)), r.dtype) \
+   lambda ctx, r: UOp.const(r.dtype, tuple(range(r.vmax+1))) \
     .reshape(tuple([r.vmax+1 if i == ctx[r.arg[0]] else 1 for i in range(len(ctx))])) if r.arg[0] in ctx else None),
   (UPat(Ops.WMMA, name="u"), expand_wmma),
 ])+pm_flatten_range+mop_cleanup
@@ -123,11 +123,11 @@ pm_expand_broadcast = pm_wmma_add+PatternMatcher([
 def do_devectorize(b:UOp):
   if b.shape == (): return None
   # broadcasting needs to be already unpacked, Invalid matches any dtype and shape
-  if not all(x.shape == b.shape or x.base.is_invalid for x in b.src): return None
+  if not all(x.shape == b.shape or x.base.arg is Invalid for x in b.src): return None
   src = []
   for idx in itertools.product(*[range(x) for x in b.shape]):
-    idx_c = [UOp.const(i) for i in idx]
-    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)))
+    idx_c = [UOp.const(None, i) for i in idx]
+    src.append(b.replace(dtype=None, src=tuple(x.base if x.base.arg 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):
@@ -362,7 +362,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, pm_decomp, name="early decompositions")
 
   # late decomps + move gates from unrenderable INVALID where
-  sink = graph_rewrite(sink, pm_dtype_decomps+pm_commit_weak, ctx=(set(), ren), name="decomp dtypes")
+  sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes")
   pm_decomp = pm_decomp+\
     get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\
     get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
@@ -371,7 +371,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # final rules for the renderer (without sym)
   extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
-  pm_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends
+  pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends
   sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")
 
   # add implicit barriers (stores/loads through LOCAL memory ordered by AFTER or across loop iterations need workgroup barriers)

tinygrad/codegen/decomp/dtype.py

@@ -18,7 +18,7 @@ def reindex(idx:UOp, off:int, mul=2) -> UOp:
 
 # 4.3.1 is the relevant section in TAOCP
 def l2i(op: Ops, dt: DType, *uops:UOp):
-  zero = UOp.const(0, dt)
+  zero = UOp.const(dt, 0)
   if len(uops) == 2: a0, a1 = uops
   elif len(uops) == 3: a0, a1, b0 = uops  # a shift's count is a single word
   elif len(uops) == 4: a0, a1, b0, b1 = uops
@@ -57,10 +57,10 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
         ua0, ua1, ub0, ub1 = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), b0.bitcast(dtypes.uint), b1.bitcast(dtypes.uint)
         a0, a1 = (a_neg:=a1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ua0, ua1))[0], ua0), a_neg.where(n[1], ua1)
         b0, b1 = (b_neg:=b1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ub0, ub1))[0], ub0), b_neg.where(n[1], ub1)
-      q, r = (z:=UOp.const(0, dtypes.uint), z), (z, z)
+      q, r = (z:=UOp.const(dtypes.uint, 0), z), (z, z)
       for i in range(63, -1, -1):
-        r = l2i(Ops.SHL, dtypes.uint, *r, UOp.const(1, dtypes.uint), z)
-        r = (r[0] | l2i(Ops.SHR, dtypes.uint, a0, a1, UOp.const(i, dtypes.uint), z)[0] & 1), r[1]
+        r = l2i(Ops.SHL, dtypes.uint, *r, UOp.const(dtypes.uint, 1), z)
+        r = (r[0] | l2i(Ops.SHR, dtypes.uint, a0, a1, UOp.const(dtypes.uint, i), z)[0] & 1), r[1]
         cond = l2i(Ops.CMPLT, dtypes.uint, *r, b0, b1).logical_not()
         diff = l2i(Ops.SUB, dtypes.uint, *r, b0, b1)
         q = ((q[0] | shl(cond.cast(dtypes.uint), i % 32), q[1]) if i < 32 else (q[0], q[1] | shl(cond.cast(dtypes.uint), i % 32)))
@@ -158,7 +158,7 @@ pm_long_decomp = PatternMatcher([
   (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
    x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None),
   (UPat(Ops.CONST, tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), lambda x:
-   UOp.const(truncate[x.tag[1]]((x.val >> 32) if x.tag[0] == 1 else (x.val & 0xFFFFFFFF)), x.tag[1]))
+   UOp.const(x.tag[1], truncate[x.tag[1]]((x.arg >> 32) if x.tag[0] == 1 else (x.arg & 0xFFFFFFFF))))
 ])
 
 # float decomposition patterns - ctx is (fr, to) tuple

tinygrad/codegen/decomp/op.py

@@ -77,7 +77,7 @@ def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher:
   # these are rewrites that make things simpler
   pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
   # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
-  if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None))
+  if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None))
   pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
   # no real hardware supports THREEFRY, but NullRenderer does
   if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32))
@@ -91,19 +91,19 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
   if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(),
     lambda x,y: (x | y).logical_not())]
   # rewrite MUL/CDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
-  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.val, 0)) else None)]
+  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
   if Ops.SHR in ops:
     # uint CDIV by 2**v -> x >> v (FLOORDIV is lowered to CDIV by the rule above before reaching here)
     pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.uints), UPat.cvar("c"))),
-      lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None)]
+      lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
     # signed CDIV (trunc) by 2**v -> (x + (x<0 ? c-1 : 0)) >> v
     pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("c"))),
       lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v
-        if (v:=powers_of_two.get(c.val, 0)) else None)]
+        if (v:=powers_of_two.get(c.arg, 0)) else None)]
     if not disable_fast_idiv:
       # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0)
       pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
-        lambda ctx, x, d: fast_idiv(ctx, x, d.val) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
+        lambda ctx, x, d: fast_idiv(ctx, x, d.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
       # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs;
       # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail)
       pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))),
@@ -119,15 +119,15 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
       (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c)<x),
       (UPat.var("x", dtypes.sints)*-1 < UPat.cvar("c"), lambda x,c:-c<x),
       ((UPat.cvar("c1")<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2")),
-        lambda x,c1,c2: x.eq(c1+1) if c1.val+1==c2.val-1 else None),  # (c-1)<x & x<(c+1) -> x==c
+        lambda x,c1,c2: x.eq(c1+1) if c1.arg+1==c2.arg-1 else None),  # (c-1)<x & x<(c+1) -> x==c
     ]
   if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))]
   if Ops.MULACC in ops:
     pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
     # also fuse (x << n) + c → MULACC(x, 2^n, c) since MUL→SHL may run first
-    if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1<<n.val), c))]
+    if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1<<n.arg), c))]
   # some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
   if Ops.FDIV in ops:
     pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
-    pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
+    pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(None, 1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
   return PatternMatcher(pat)

tinygrad/codegen/decomp/transcendental.py

@@ -16,8 +16,8 @@ def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d)[0] - 1)) - (0 i
 def exponent_mask(d:DType) -> int: return (1 << dtypes.finfo(d)[0]) - 1
 
 # **** utils ****
-def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().val) if isinstance(y, UOp) else 2**y)
-def shl(x:UOp|int, y:UOp|int) -> UOp: return x * (2**(y.simplify().val) if isinstance(y, UOp) else 2**y)
+def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y)
+def shl(x:UOp|int, y:UOp|int) -> UOp: return x * (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y)
 
 def rintk(d:UOp) -> UOp:
   """round d:float to int away from 0"""
@@ -93,7 +93,7 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
   def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
   def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
 
-  a = [_take(UOp.const(0, dtypes.uint32), i) for i in range(4)]
+  a = [_take(UOp.const(dtypes.uint32, 0), i) for i in range(4)]
   #  (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e))
   # Note: e >= 1 for all numbers d >= 1.0. assume e != 0
   hi = _shl_lazy(a[0], e) | _shr_lazy(a[1], offset)

tinygrad/codegen/late/coalesce.py

@@ -84,7 +84,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
   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)),))
   shapes[buf.arg.slot] = (h, w)
-  if valid.op is not Ops.CONST or valid.val is not True:
+  if valid.op is not Ops.CONST or valid.arg is not True:
     return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid), dtype=dtypes.float)
   else:
     return buf.index(cidx.src[1], cidx.src[0], dtype=dtypes.float)
@@ -106,15 +106,17 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
     # TODO: this should handle images too, it's just memory coalescing
     if u.op in {Ops.LOAD, Ops.STORE}:
       assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores"
-      assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}"
+      # movement-op-wrapped accesses (e.g. a REG placeholder store through a RESHAPE) aren't
+      # index-addressed; there's nothing to coalesce for them
+      if u.src[0].op is not Ops.INDEX: continue
       buf, idx_u = u.src[0].src
       if buf.addrspace == AddrSpace.REG: continue
       idx, valid = idx_u.get_idx(), idx_u.get_valid()
       root_src: UOp|str
-      if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].val
-      elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].val
-      elif idx.op is Ops.CONST and idx.val is Invalid: root_src, arg = "INVALID", 0
-      elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val
+      if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg
+      elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg
+      elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0
+      elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg
       else: root_src, arg = idx, 0
       memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)
 
@@ -127,7 +129,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
     if ctx is not None and ctx.target.device == "DSP":
       lengths = [128,64,32,16,8,4]
       must_divide = False
-    elif buf.dtype not in (dtypes.float, dtypes.half, dtypes.int, dtypes.uint, *dtypes.fp8s) and not is_image_shape(buf._shape):
+    elif buf.dtype not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not is_image_shape(buf._shape):
       pass
     elif buf.addrspace == AddrSpace.REG:
       pass
@@ -137,21 +139,22 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
       # TODO: a better way to get this than ctx
       lengths = [8,4,2] if buf.dtype == dtypes.half and getenv("ALLOW_HALF8") else [4,2]
     lengths.append(1)  # worst case, it's not folded
+    # stores that alias (same buf+idx+valid from multiple store sites, e.g. a double-buffered LDS
+    # slot written in a prologue and a loop body) can't be merged: merging would drop one write.
+    # keep those scalar and only coalesce the unique ones.
+    keys = [k for k in sorted(offsets.keys()) if op is Ops.LOAD or len(offsets[k]) == 1]
     # do the grouping
-    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
+    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(keys), lambda x: x[1]-x[0])]
     for full_grp in grouped_offsets:
       while len(full_grp):
-        offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(full_grp[0])
+        offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(None, full_grp[0])
         length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0]
         grp = full_grp[:length]
         # NOTE: we apply the valid again after we determine the length
         offset = offset.valid(valid) if valid is not None else offset
-        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset)
+        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(None, len(grp)))) if len(grp) > 1 else buf.index(offset)
         if op == Ops.STORE:
-          datas = []
-          for i,g in enumerate(grp):
-            assert len(offsets[g]) == 1, f"attempting multiple stores: {len(offsets[g])}"
-            datas.append(offsets[g][0].src[1])
+          datas = [offsets[g][0].src[1] for g in grp]
           store = idx.store(UOp.stack(*datas) if len(datas) > 1 else datas[0])
           for i,g in enumerate(grp): replacements[offsets[g][0]] = store
         else:

tinygrad/codegen/late/gater.py

@@ -3,7 +3,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops
 from tinygrad.dtype import Invalid, dtypes
 
 def move_where_load(gate, l, a, w):
-  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else
+  return l.replace(src=(l.src[0], l.vconst_like(0) if a.arg is Invalid else
     a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(w.dtype)
 
 pm_move_gates_from_index = PatternMatcher([

tinygrad/codegen/late/regalloc.py

@@ -52,7 +52,7 @@ class LinearScanRegallocContext:
         # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes
         sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize)
         offset = self.stack_size + (sz - self.stack_size % sz) % sz
-        self.spills[v] = UOp.const(offset, dtypes.int32)
+        self.spills[v] = UOp.const(dtypes.int32, offset)
         self.stack_size = offset + sz
       r = alloc(cons if cons is not None else v.cons, i)
       self.insert_before.setdefault(i, []).append((v, r))
@@ -84,7 +84,7 @@ class LinearScanRegallocContext:
 
       # allocate stack array
       if u.op is Ops.BUFFER:
-        self.locals[u] = UOp.const(self.stack_size, dtypes.int32)
+        self.locals[u] = UOp.const(dtypes.int32, self.stack_size)
         self.stack_size += u.max_numel() * u.dtype.itemsize
 
       # loop prologue, avoid loading inside the loop

tinygrad/codegen/opt/heuristic.py

@@ -26,18 +26,22 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   """
   # NOTE: unless TC_OPT is > 0, we only trigger tensor cores if there's only one reduce axis
   if USE_TC > 0 and (len(k.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (TC_OPT.value >= 1)):
-    for axis in range(3):
-      tk = k.copy()
-      # check TC first and apply hand-coded opts if successful
-      try: rngs = tk.apply_opt(Opt(OptOps.TC, axis, (TC_SELECT.value, TC_OPT.value, USE_TC.value)))
-      except KernelOptError: continue
-      for tc_dim in [1,0]: # attempt to upcast M and N
-        szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
-        if szs:
-          # set it to the replaced range
-          rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
-      if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
-        tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
+    good_tc_opt = False
+    tk = k.copy()
+    try: # check TC first and apply hand-coded opts if successful
+      rngs = tk.apply_opt(Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, USE_TC.value)))
+      good_tc_opt = True
+    except KernelOptError:
+      pass
+    if good_tc_opt:
+      if rngs is not None:
+        for tc_dim in [1,0]: # attempt to upcast M and N
+          szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
+          if szs:
+            # set it to the replaced range
+            rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
+        if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
+          tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
       return tk
 
   # make a copy so it does not mutate the input
@@ -126,8 +130,8 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
           if rng in idx.backward_slice: num_strides += 1
           for c in idx.split_uop(Ops.ADD):
             if c is rng: sum_strides += 1
-            if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].val
-            if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].val
+            if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg
+            if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg
         xb_choices.append((num_strides, sum_strides, axis, upcast_amount))
     if xb_choices:
       xb_choices = sorted(xb_choices)

tinygrad/codegen/opt/postrange.py

@@ -198,7 +198,7 @@ class Scheduler:
       for b in self.bufs:
         if rng in (i:=b.src[1].get_idx()).backward_slice_with_self:
           nb = b.replace(src=(b.src[0], i.valid(valid&b.src[1].get_valid())))
-          replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(Invalid, b.dtype))
+          replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(b.dtype, Invalid))
       self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
     elif opt.op is OptOps.SWAP:
       try:
@@ -245,8 +245,6 @@ class Scheduler:
           if not (axis < len(axis_choices)): continue
           axes = list(axis_choices[axis])
 
-          if any(a.arg[-1] is AxisType.REDUCE for a in axes[:2]): raise KernelOptError("tensor core X/Y axes can't be REDUCE")
-
           # tag the reduceop
           self.ast = self.ast.substitute({reduceop: reduceop.replace(tag="TC")})
 
@@ -304,7 +302,7 @@ class Scheduler:
             # TODO: remove tc_upcast_axes from the arg
             # do the reduce_axes always disappear? i think they don't
             # they need to be moved into the WMMA srcs
-            tc_uop = UOp.wmma(srcs[0], srcs[1], UOp.const((0.0,)*tc.elements_per_thread[2], tc.dtype_out),
+            tc_uop = UOp.wmma(srcs[0], srcs[1], UOp.const(tc.dtype_out, (0.0,)*tc.elements_per_thread[2]),
                               tc.dims, self.ren.target.device, tc.threads, tc_upcast_axes=tc_upcast_axes)
 
             # preserve extra reduces

tinygrad/codegen/simplify.py

@@ -47,7 +47,7 @@ def mark_gated(ctx, idx):
     guards = {r:c for v in cond.split_uop(Ops.AND) if v.op is Ops.CMPLT and (r:=v.src[0]).op is Ops.RANGE and (c:=v.src[1]).op is Ops.CONST}
   else: x, guards = idx, {}
   # ensure that we choose max(c_i) for all i where r < c_i
-  ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].val < c.val)}
+  ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].arg < c.arg)}
   # but if a range is ever ungated, we cannot shrink it
   ctx |= {r:r.src[0] for r in x.ranges if r not in guards}
 
@@ -59,10 +59,14 @@ pm_simplify_ranges = PatternMatcher([
   (UPat(Ops.SINK, name="x"), lambda ctx, x: do_substitute(ctx, x, lambda r,c: r.replace(src=(c,)))),
 ])
 
+SPLITTABLE_TYPES = {AxisType.WEAK, AxisType.REDUCE, AxisType.LOOP}
+
 def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
-  # ranges that aren't looped over can't be split
-  if r not in ctx and r.arg[-1] not in {AxisType.WARP, AxisType.DEVICE} \
-    and r.src[0].op is Ops.CONST and r.src[0].divides(c.val) is not None: ctx[r] = c
+  # ranges that aren't looped over can't be split. ranges with hardware meaning are never split
+  # (LOCAL/WARP/THREAD/GLOBAL/GROUP_REDUCE/DEVICE map to launch dims; UPCAST/UNROLL are vector
+  # widths): splitting them scrambles the logical<->hardware mapping of hand-written kernels.
+  if r not in ctx and r.arg[-1] in SPLITTABLE_TYPES \
+    and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
 
 def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:
   ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None})

tinygrad/dtype.py

@@ -17,7 +17,6 @@ class ConstFloat(float):
     if self is other: return True
     if isinstance(other, float) and math.isnan(self) and math.isnan(other): return True
     return float.__eq__(self, other)
-  def __ne__(self, other): return res if (res:=self.__eq__(other)) is NotImplemented else not res  # float.__ne__ disagrees with __eq__ on nan
   def __hash__(self): return hash(self.bits)
   def __repr__(self): return f"ConstFloat({float.__repr__(self)})"
   def __str__(self): return float.__repr__(self)
@@ -75,7 +74,8 @@ class DType(metaclass=DTypeMetaClass):
   def max(self):
     if dtypes.is_int(self): return 2**(self.bitsize)-1+self.min
     return float("inf") if dtypes.is_float(self) else True
-  def const(self, val: ConstType):
+  def const(self, val: tuple[ConstType, ...]|ConstType):
+    if isinstance(val, tuple): return tuple(map(self.const, val))
     if isinstance(val, InvalidType): return val
     # NOTE: float('nan') != float('nan'), so we canonicalize here
     if isinstance(val, float) and math.isnan(val): val = math.nan
@@ -165,8 +165,6 @@ assert dtypes.is_float(dtypes.default_float), f"{DEFAULT_FLOAT.value} is not a f
 assert dtypes.is_int(dtypes.default_int), f"{DEFAULT_INT.value} is not an int dtype"
 def strong_dtype(dtype:DType) -> DType:
   return {dtypes.weakint: dtypes.default_int, dtypes.weakfloat: dtypes.default_float}.get(dtype, dtype)
-def weak_dtype(dtype:DType) -> DType:
-  return dtypes.weakfloat if dtypes.is_float(dtype) else dtypes.weakint if dtypes.is_int(dtype) else dtype
 
 # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html
 # we don't support complex type

tinygrad/engine/realize.py

@@ -28,7 +28,7 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
   ast, arg_uops = call.src[0], get_call_arg_uops(call)
   if ast.op is Ops.PROGRAM: return ast.arg.name
   if ast.op is Ops.SLICE:
-    offset = ast.src[1].val * arg_uops[1].dtype.itemsize
+    offset = ast.src[1].arg * arg_uops[1].dtype.itemsize
     return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow")
   if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow")
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow")
@@ -156,7 +156,7 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
 def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   resolved = resolve_params(call, ctx.input_uops)
   bufs = [cast(Buffer, b.buffer) for b in resolved]
-  bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].val*bufs[1].dtype.itemsize)
+  bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].arg*bufs[1].dtype.itemsize)
   with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv
   return None
 
@@ -198,7 +198,7 @@ def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
 
 def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)]
-  shape, pos_var = tuple(s.val for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
+  shape, pos_var = tuple(s.arg for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
   with track_stats(ctx, call, bufs[0].device, bufs, ctx.var_vals):
     bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
   return None

tinygrad/function.py

@@ -1,6 +1,7 @@
 import functools, time
 from typing import Generic, TypeVar, Callable, cast, overload
 from tinygrad.helpers import Context, dedup, getenv, DEBUG
+from tinygrad.dtype import Invalid
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import get_state_dict
@@ -21,7 +22,8 @@ def invalid_outputs(uret:UOp) -> set[UOp]:
   # invalids() returns fresh write-only scratch: a clone storing CONST(Invalid)
   # don't capture it as an input; only skip fresh buffers, not realized ones
   return {u.src[0].buf_uop for u in uret.backward_slice_with_self
-          if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}
+          if u.op is Ops.STORE and u.src[1].base.op is Ops.CONST and u.src[1].base.arg is Invalid
+          and not u.src[0].buf_uop.is_realized}
 
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):

tinygrad/helpers.py

@@ -620,5 +620,3 @@ class count:
     cur = self.n
     self.n += self.step
     return cur
-
-# test change for the szdiff bot

tinygrad/llm/gguf.py

@@ -37,7 +37,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
 
   def q_to_uint8(t: Tensor, b: int) -> Tensor:
     # TODO: rewrite with arange?
-    shift_tensor, bitmask = Tensor.const(tuple(2**(i*b) for i in range(8//b)), t.dtype), 0xff >> (8 - b)
+    shift_tensor, bitmask = Tensor.const(t.dtype, tuple(2**(i*b) for i in range(8//b))), 0xff >> (8 - b)
     return t.unsqueeze(-1).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
 
   if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None:
@@ -74,7 +74,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
       scale_words = blocks[:, 66:98].bitcast(dtypes.uint32)
       db = d * (scale_words.rshift(28).cast(dtypes.float32) + 0.5).reshape((-1, 8, 1, 1)) * 0.5
-      sign_idx = scale_words.unsqueeze(-1).rshift(Tensor.const((0, 7, 14, 21), dtypes.uint32)).bitwise_and(0x7F).reshape((-1, 32)).cast(dtypes.int32)
+      sign_idx = scale_words.unsqueeze(-1).rshift(Tensor.const(dtypes.uint32, (0, 7, 14, 21))).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.iq3xxs_grid, (256, 4))[blocks[:, 2:66]].reshape((-1, 8, 4, 8))
@@ -95,7 +95,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       return (db * _ggml_iq_grid(t.device, _ggml.iq2s_grid, (1024, 8))[q].reshape((-1, 16, 2, 8)) * signs).flatten(-3)
     if ggml_type == 23:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1))
-      scale_shifts = Tensor.const((0, 2, 4, 6, 8, 10, 12, 14), dtypes.uint16)
+      scale_shifts = Tensor.const(dtypes.uint16, (0, 2, 4, 6, 8, 10, 12, 14))
       iq4_xs_lut = Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device)
       scales_l = Tensor.stack((sl:=blocks[:, 4:8]).bitwise_and(0xF), sl.rshift(4), dim=2).reshape((-1, 8))
       scales_h = blocks[:, 2:4].bitcast(dtypes.uint16).unsqueeze(-1).rshift(scale_shifts).bitwise_and(0x03).reshape((-1, 8)).cast(dtypes.uint8)

tinygrad/llm/model.py

@@ -432,10 +432,9 @@ class Transformer:
     if start_pos < len(self._cached_tokens) and (resets := [r for b in self.blk for r in b._state_reset_ops()]): Tensor.realize(*resets)
     out, prompt_len = None, len(tokens)
     while len(tokens) < self.max_context:
-      n_toks = min(chunk_size, len(tokens) - start_pos)
-      sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks)
+      sp, nt = v_start_pos.bind(start_pos), v_toks.bind(min(chunk_size, len(tokens) - start_pos))
       out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
-      start_pos += n_toks
+      start_pos += nt.val
       # chunked prefill: keep processing until all prompt tokens are consumed
       if start_pos < len(tokens): continue
       tokens.append(int(out.item()))

tinygrad/mixin/creation.py

@@ -9,7 +9,7 @@ if TYPE_CHECKING:
 
 class CreationMixin(DTypeMixin, MovementMixin):
   @staticmethod
-  def const(b, dtype=None): raise NotImplementedError
+  def const(dtype, b): raise NotImplementedError
 
   def const_like(self, b: ConstType) -> Self: return self._wrap_uop(self._uop.const_like(b))
 
@@ -78,7 +78,7 @@ class CreationMixin(DTypeMixin, MovementMixin):
     from tinygrad.uop.ops import UOp
     new_shape = argfix(shape)
     dt = to_dtype(dtype) if dtype is not None else fill_value.dtype if isinstance(fill_value, UOp) else dtypes.from_py(fill_value)
-    val = cls.const(fill_value, dt)
+    val = cls.const(dt, fill_value)
     val = val.reshape((1,)*len(new_shape)).expand(new_shape)
     if not buffer: return val
     ret = val.empty_like(dt if dtype is not None else None, device)

tinygrad/mixin/dtype.py

@@ -1,5 +1,5 @@
 from typing import TYPE_CHECKING, Self
-from tinygrad.dtype import DType, DTypeLike, dtypes, to_dtype
+from tinygrad.dtype import DType, DTypeLike, dtypes, to_dtype, Invalid
 from tinygrad.uop import Ops
 
 if TYPE_CHECKING:
@@ -30,7 +30,7 @@ class DTypeMixin:
     print(t.dtype, t.numpy())
     ```
     """
-    return self if self.dtype == (dt:=to_dtype(dtype)) or self._uop.base.is_invalid else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
+    return self if self.dtype == (dt:=to_dtype(dtype)) or self._uop.base.arg is Invalid else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
 
   def bitcast(self, dtype:DTypeLike) -> Self:
     """

tinygrad/mixin/elementwise.py

@@ -1,7 +1,7 @@
 import math, functools, operator
 from typing import TYPE_CHECKING, Literal, Self
 from tinygrad.uop import Ops
-from tinygrad.dtype import dtypes, ConstType, PyConst, least_upper_dtype, least_upper_float, weak_dtype
+from tinygrad.dtype import dtypes, ConstType, PyConst, least_upper_dtype, least_upper_float
 from tinygrad.helpers import argfix, polyN
 from tinygrad.mixin.creation import CreationMixin
 
@@ -21,12 +21,7 @@ class ElementwiseMixin(CreationMixin):
   def _broadcasted(self, y: 'Self|ConstType|UOp', reverse: bool = False) -> tuple[Self, Self]:
     y = self.ufix(y)
     x, y = (self, y) if not reverse else (y, self)
-    out_dtype = least_upper_dtype(x.dtype, y.dtype)
-    # keep weak CONST weak, might lift weakint -> weakfloat
-    def promote(t):
-      if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.val, weak_dtype(out_dtype)))
-      return t.cast(out_dtype)
-    return promote(x), promote(y)
+    return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype)
 
   def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self:
     lhs, rhs = self._broadcasted(x, reverse)
@@ -395,8 +390,6 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     t, x = self._broadcasted(x)
-    # ~ is width-dependent: min(a,b) == ~max(~a,~b) only holds at a common width, so a weak operand commits at its sibling's
-    t, x = t.cast(dt:=least_upper_dtype(t.dtype, x.dtype)), x.cast(dt)
     return t._inverse().maximum(x._inverse())._inverse()
 
   def copysign(self, other: Self | ConstType) -> Self:

tinygrad/mixin/op.py

@@ -159,7 +159,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
       per_dim.append((idx >= s) & (idx < e) & (((e-1-idx) if m['stride'] < 0 else (idx-s)) % st == 0))
     vb = vb.flip(tuple(d for d, m in enumerate(mops) if m['stride'] < 0))
     vb = vb.pad(tuple((m['boundary'][0], self.shape[d] - m['boundary'][1]) for d, m in enumerate(mops)))
-    return (type(self).uprod(*per_dim) if per_dim else type(self).const(True)).where(vb, self)
+    return (type(self).uprod(*per_dim) if per_dim else type(self).const(dtypes.bool, True)).where(vb, self)
 
   @classmethod
   def arange(cls, start, stop=None, step=1, dtype:DTypeLike|None=None) -> Self:
@@ -703,28 +703,6 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     m, _, ss = self._softmax(axis, dtype)
     return m - ss.log()
 
-  def softmin(self, axis=-1, dtype:DTypeLike|None=None) -> Self:
-    """
-    Applies the softmin function to the tensor along the specified axis.
-
-    Rescales the elements of the tensor such that they lie in the range [0, 1] and sum to 1.
-
-    You can pass in the `axis` keyword argument to control the axis along which the softmin is computed.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    Tensor.manual_seed(42)
-    t = Tensor.randn(2, 3)
-    print(t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.softmin().numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.softmin(axis=0).numpy())
-    ```
-    """
-    return (-self).softmax(axis, dtype)
-
   def cat(self, *args:Self, dim:int=0) -> Self:
     """
     Concatenates self with other tensors in `args` along an axis specified by `dim`.
@@ -1411,7 +1389,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     ret = self
     for dim in range(dims):
       ret = ret.transpose(0, dim)
-      ret = sum(type(self).const(tuple(float(m[k]) for m in mat), ret.dtype).reshape((len(mat),)+(1,)*(ret.ndim-1)) * ret[k]
+      ret = sum(type(self).const(ret.dtype, tuple(float(m[k]) for m in mat)).reshape((len(mat),)+(1,)*(ret.ndim-1)) * ret[k]
                 for k in range(len(mat[0])))
       assert not isinstance(ret, int), "sum over empty winograd matrix"
       ret = ret.transpose(0, dim)
@@ -1898,7 +1876,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     # https://keccak.team/keccak_specs_summary.html
 
     def ctensor(l: Sequence[PyConst], dtype: DType = dtypes.uint64):
-      return type(self).const(tuple(l), dtype)
+      return type(self).const(dtype, tuple(l))
     rot_offsets = [44, 43, 21, 14, 28, 20, 3, 45, 61, 1, 6, 25, 8, 18, 27, 36, 10, 15, 56, 62, 55, 39, 41, 2]
     rot_offsets_v0, rot_offsets_v1 =  ctensor([0] + [1 << v for v in rot_offsets]), ctensor([1] + [1 << (64 - v) for v in rot_offsets])
 
@@ -1918,7 +1896,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     lbe = (data.shape[1] - 1) * 200 + rate - data_pad
     if data_pad == 1: mb = [(lbe, 0), (1, dsbyte ^ 0x80), (200 - rate, 0)]
     else: mb = [(lbe, 0), (1, dsbyte), (data_pad - 2, 0), (1, 0x80), (200 - rate, 0)]
-    pad_mask = type(self).cat(*(type(self).const(v, dtypes.uint8).expand(l) for l, v in mb if l > 0)).unsqueeze(0)
+    pad_mask = type(self).cat(*(type(self).const(dtypes.uint8, v).expand(l) for l, v in mb if l > 0)).unsqueeze(0)
 
     data = (data.flatten(1) ^ pad_mask).reshape(*data.shape[:2], 200).bitcast(dtypes.uint64)
 

tinygrad/renderer/cstyle.py

@@ -34,18 +34,18 @@ base_rewrite = PatternMatcher([
   # const
   (UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, ctx.infinity)})"),
   (UPat(Ops.CONST, arg=-math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, f'-{ctx.infinity}')})"),
-  (UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx.nan)})" if math.isnan(x.val) else None),
-  (UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.val}f"),
-  (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.val}l"),
-  (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}ul"),
-  (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}u"),
-  (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.val else "0"),
+  (UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx.nan)})" if math.isnan(x.arg) else None),
+  (UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.arg}f"),
+  (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.arg}l"),
+  (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}ul"),
+  (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}u"),
+  (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.arg else "0"),
   # consts are rendered to larger type and casted
-  (UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}f')})"),
-  (UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}u')})"),
-  (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, str(x.val))})"),
+  (UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.arg}f')})"),
+  (UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.arg}u')})"),
+  (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, str(x.arg))})"),
   # default const render
-  (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.val)),
+  (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
 
   # SHRINK/INDEX
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)),
@@ -164,7 +164,7 @@ class CStyleLanguage(Renderer):
     if buf.addrspace == AddrSpace.ALU:
       # this is lane access in C
       if idx.op is not Ops.CONST: return f"({self[buf]})[{self[idx]}]"
-      return self[buf]+(f"[{idx.val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.val]}")
+      return self[buf]+(f"[{idx.arg}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.arg]}")
     return f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})"
 
   def render_buffer(self, x:UOp):
@@ -236,7 +236,7 @@ class CStyleLanguage(Renderer):
       assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}"
 
       if u.op in {Ops.ENDIF, Ops.END}: depth -= 1
-      if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
+      if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK} or \
         (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \
         (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
         (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
@@ -320,7 +320,7 @@ class OpenCLRenderer(CStyleLanguage):
     (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
     # bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort
     (UPat(Ops.CONST, dtypes.bfloat16, name="x"),
-      lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.val)))[0] >> 16)}u"),
+      lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.arg)))[0] >> 16)}u"),
     # load/store image (OpenCL)
     (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"),
     (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))),
@@ -356,7 +356,7 @@ class MetalRenderer(CStyleLanguage):
   code_for_workitem = {"g": lambda x: f"gid.{chr(120+int(x))}", "l": lambda x: f"lid.{chr(120+int(x))}"}
   # uint3 used for gid/lid - TODO: this should probably be `ushort3 lid [[thread_position_in_threadgroup]]`
   extra_args = ['uint3 gid [[threadgroup_position_in_grid]]', 'uint3 lid [[thread_position_in_threadgroup]]']
-  type_map = {dtypes.uint32: "uint", dtypes.bfloat16: "bfloat"}
+  type_map = {dtypes.bfloat16: "bfloat"}
 
   # precise::sin
   code_for_op = {**CStyleLanguage.code_for_op, Ops.SIN: lambda x,dtype: f"precise::sin({x})"}
@@ -420,7 +420,7 @@ class CUDARenderer(CStyleLanguage):
     Ops.EXP2: lambda x,dtype: f"hexp2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"exp2({x})",
     Ops.SQRT: lambda x,dtype: f"hsqrt({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sqrt({x})",
     Ops.RECIPROCAL: lambda x,dtype: f"hrcp({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"(1/{x})" }
-  type_map = {dtypes.uint32: "uint", dtypes.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"}
+  type_map = {dtypes.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"}
   extra_matcher = create_non_native_float_pats(dtypes.fp8s, casting=False) + PatternMatcher([
     (UPat(Ops.CAST, dtypes.fp8s, UPat.var("x", dtypes.fp8s), name='y'), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=y.dtype else None),
   ])
@@ -436,7 +436,7 @@ class CUDARenderer(CStyleLanguage):
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
     # TODO: why is dtypes.bfloat16.name == "__bf16"? would be easier not override dtypes.name
-    prefix = ["typedef unsigned int uint;", "#define INFINITY (__int_as_float(0x7f800000))", "#define NAN (__int_as_float(0x7fffffff))",
+    prefix = ["#define INFINITY (__int_as_float(0x7f800000))", "#define NAN (__int_as_float(0x7fffffff))",
               "template <class T, class F> __device__ __forceinline__ T tg_bitcast(F v) { union U { F f; T t; }; U u; u.f = v; return u.t; }"]
     used_dtypes = uops_to_dtypes(uops)
     if any(dt in dtypes.fp8s for dt, _ in used_dtypes): prefix.append("#include <cuda_fp8.h>")
@@ -474,6 +474,42 @@ class NVCCRenderer(CUDARenderer):
 def fp8_index(dtype: DType): return (dtypes.fp8e4m3, dtypes.fp8e5m2).index(dtype.scalar())
 def _ocml(op): return lambda x,dtype: f"__ocml_{op}_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})"
 
+def _g2l_parts(u:UOp) -> tuple[UOp, UOp, UOp, UOp]|None:
+  """STORE(local[li]) <- LOAD(global[gi]) (scalar INDEX or vec SHRINK): a global->shared copy expressible
+  as one buffer_load_lds (direct-to-LDS) instruction on gfx9.4+. Returns (buf, lidx, gbuf, gidx)."""
+  if u.op is not Ops.STORE or len(u.src) != 2: return None
+  li, ld = u.src
+  if li.op is Ops.INDEX and li.addrspace == AddrSpace.LOCAL and len(li.src) == 2: buf, idx = li.src
+  elif li.op is Ops.SHRINK and li.src[1].dtype is not None and li.src[0].addrspace == AddrSpace.LOCAL: buf, idx = li.src[0], li.src[1]
+  else: return None
+  if ld.op is not Ops.LOAD or len(ld.src) != 1: return None
+  gi = ld.src[0]
+  if gi.op is Ops.INDEX and gi.addrspace == AddrSpace.GLOBAL and len(gi.src) == 2: gbuf, gidx = gi.src
+  elif gi.op is Ops.SHRINK and gi.src[0].addrspace == AddrSpace.GLOBAL: gbuf, gidx = gi.src[0], gi.src[1]
+  else: return None
+  if li.dtype.scalar() != ld.dtype.scalar(): return None
+  return buf, idx, gbuf, gidx
+
+def _g2l_match(u:UOp) -> bool: return _g2l_parts(u) is not None
+
+def _render_g2l_lds(ctx, u:UOp) -> str|None:
+  if (parts := _g2l_parts(u)) is None: return None
+  buf, idx, gbuf, gidx = parts
+  sz = u.src[0].dtype.itemsize   # whole copy size in bytes (16 for an 8xbf16 chunk)
+  esz = u.src[0].dtype.scalar().itemsize
+  return (f"llvm_amdgcn_raw_buffer_load_lds(make_srsrc_((void*){ctx[gbuf]}, {gbuf.max_numel()*gbuf.dtype.itemsize}), "
+          f"(as3_uint32_ptr)(&({ctx[buf]}[({ctx[idx]})])), {sz}, ((unsigned)({ctx[gidx]}))*{esz}U, 0, 0, 0);")
+
+G2L_LDS_DECLS = [
+  "typedef int int32x4_t __attribute__((ext_vector_type(4)));",
+  "typedef __attribute__((address_space(3))) unsigned* as3_uint32_ptr;",
+  ("extern __attribute__((device)) void\n"
+   "llvm_amdgcn_raw_buffer_load_lds(int32x4_t rsrc, as3_uint32_ptr lds_ptr, int size, int voffset, int soffset, int offset, int aux)\n"
+   ' __asm("llvm.amdgcn.raw.buffer.load.lds");'),
+  """static inline __attribute__((device)) int32x4_t make_srsrc_(const void* p, unsigned rb) {
+  int32x4_t r = {(int)(unsigned long)p, (int)(((unsigned long)p)>>32), (int)rb, 0x110000};
+  return r;\n}"""]
+
 class HIPRenderer(CStyleLanguage):
   shared_max = 65536
   # NOTE: this is only really needed on gfx12, even though gfx11 reports the same limitation
@@ -491,13 +527,15 @@ class HIPRenderer(CStyleLanguage):
     if not self.is_cdna4(target.arch): self.extra_matcher += pm_manual_bf16_cast
     if self.is_cdna(target.arch):
       self.string_rewrite = PatternMatcher([
+        # direct global->LDS copies (buffer_load_lds), skipping the register round-trip
+        (UPat(Ops.STORE, name="st"), lambda ctx,st: _render_g2l_lds(ctx, st) if getenv("HK_G2L") else None),
         (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]},"
           f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[0][2] == 128 else None),
         (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"),
-        (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.val) else None),
+        (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.arg) else None),
         (UPat(Ops.CONST, dtypes.fp8s, arg=math.inf, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.infinity}, {fp8_index(x.dtype)})"),
         (UPat(Ops.CONST, dtypes.fp8s, arg=-math.inf, name="x"), lambda ctx,x: f"f32_to_fp8(-{ctx.infinity}, {fp8_index(x.dtype)})"),
-        (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.val}f, {fp8_index(x.dtype)})"),
+        (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.arg}f, {fp8_index(x.dtype)})"),
         (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",),
           lambda ctx,x: f"f32_to_fp8({ctx[x.src[0]]}, {fp8_index(x.dtype)})"),
         (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",),
@@ -522,7 +560,7 @@ class HIPRenderer(CStyleLanguage):
       lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
       if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None),
     # bfloat16 constant casting
-    (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.val, dtypes.float))),
+    (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(dtypes.float, x.arg))),
   ])
 
   def asm(self, prg:UOp, lin:UOp) -> bytes:
@@ -536,9 +574,12 @@ class HIPRenderer(CStyleLanguage):
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
     prefix, ockl = [], []
+    g2l_used = any(_g2l_match(u) for u in uops) or \
+      any(u.op is Ops.CUSTOMI and isinstance(u.arg, str) and u.arg.startswith("llvm_amdgcn_raw_buffer_load_lds") for u in uops)
+    if self.is_cdna(self.target.arch) and g2l_used: prefix += G2L_LDS_DECLS
     type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" }
     used_dtypes = uops_to_dtypes(uops)
-    if any(u.op is Ops.CONST and not math.isfinite(u.val) for u in uops):
+    if any(u.op is Ops.CONST and not math.isfinite(u.arg) for u in uops):
       prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"]
     if any(u.op is Ops.SPECIAL for u in uops):
       prefix.append("typedef long unsigned int size_t;")

tinygrad/renderer/isa/x86.py

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

tinygrad/renderer/llvmir.py

@@ -67,7 +67,7 @@ base_rewrite = PatternMatcher([
    f"  {ctx[x]} = getelementptr inbounds {ldt(x.dtype)}, {ldt(x.dtype, ptr=True)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"),
   # register index
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("idx")), name="x"), lambda ctx,buf,idx,x:
-   f"  {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.val}" if buf.addrspace == AddrSpace.ALU else None),
+   f"  {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.arg}" if buf.addrspace == AddrSpace.ALU else None),
 
   # load/store
   (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"),
@@ -170,7 +170,7 @@ class LLVMRenderer(Renderer):
           kernel.append(f"  {r[u]} = addrspacecast [{size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{size} x {ldt(u.dtype)}]*")
         else:
           kernel.append(f"  {r[u]} = alloca [{size} x {ldt(u.dtype)}], align 16")
-      elif u.op is Ops.CONST: r[u] = lconst(u.val, u.dtype)
+      elif u.op is Ops.CONST: r[u] = lconst(u.arg, u.dtype)
       elif u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype):
         r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast
       else:
@@ -274,7 +274,7 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
           src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
           if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
         (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
-          src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
+          src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(x.src[2].dtype, 0.0)
             for j in range(x.max_numel()*2)))),
           arg=(*x.arg[:4], None)).index(i*2)
           for i in range(x.max_numel()))) if x.max_numel() == 8 else None),

tinygrad/renderer/nir.py

@@ -122,7 +122,7 @@ class NIRRenderer(Renderer):
 
   extra_matcher = PatternMatcher([
     # handle negative unsigned CONST
-    (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype.max+x.val+1, x.dtype) if x.val < 0 else None),
+    (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype, x.dtype.max+x.arg+1) if x.arg < 0 else None),
     # from ptx
     (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
     # load/store bool -> uint8
@@ -144,7 +144,7 @@ class NIRRenderer(Renderer):
   ])
 
   def_rewrite = PatternMatcher([
-    (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.val, x.dtype)),
+    (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.arg, x.dtype)),
     (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, x.dtype.itemsize if x.addrspace is AddrSpace.ALU else 8)),
     (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
     (UPat(Ops.STORE, src=(UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val"))),
@@ -185,7 +185,7 @@ class NIRRenderer(Renderer):
 
   def render(self, uops:list[UOp]):
     self.prerender(uops)
-    for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].val
+    for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].arg
     self.r: dict[UOp, Any] = {}
     self.param_idx = 0
     ranges: list[mesa.nir_def|None] = []
@@ -194,7 +194,7 @@ class NIRRenderer(Renderer):
       if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.STACK and len(u.src) == 0): pass
       elif u.op in {Ops.INDEX, Ops.SHRINK}:
         # INDEX on a register value picks the element, memory INDEX is handled in the LOAD/STORE patterns
-        if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].val)
+        if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].arg)
       elif u.op is Ops.AFTER:
         self.r[u] = self.r[u.src[0]]
       elif u.op == Ops.SINK:

tinygrad/renderer/ptx.py

@@ -79,8 +79,8 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i
   (a.itemsize < b.itemsize or dtypes.is_int(b) or b == dtypes.bool) else ''
 
 string_rewrite = PatternMatcher([
-  (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.val, x.dtype)}, 0;"),
-  (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.val, x.dtype)};"),
+  (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.arg, x.dtype)}, 0;"),
+  (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.arg, x.dtype)};"),
   (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"),
   (UPat(Ops.PARAM, name="x"), lambda ctx, x:
    f"ld.param.{ctx.types[dtypes.ulong] if x.addrspace is AddrSpace.GLOBAL else ctx.mem_types[x.dtype]} {ctx.r[x]}, [data{x.arg.slot}+0];"),
@@ -203,7 +203,7 @@ class PTXRenderer(Renderer):
         # on REG, INDEX/SHRINK pick the register (must be CONST) and LOAD is a noop
         if u.op is not Ops.LOAD and u.src[1].op is not Ops.CONST:
           raise RuntimeError(f"PTX does not support dynamic register indexing: {u}")
-        r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].val]
+        r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].arg]
         continue
       if u.op is Ops.SPECIAL: r[u] = "%" + u.arg
       elif u.op is Ops.LOAD:

tinygrad/renderer/wgsl.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import strip_parens
 def _mask(dt:DType): return 0xFF if dt.itemsize == 1 else 0xFFFF
 
 def sign_extend(val:UOp, sext_am:int):
-  return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(0xffffffff, dtypes.uint32) << sext_am, UOp.const(0, dtypes.uint32)) \
+  return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(dtypes.uint32, 0xffffffff) << sext_am, UOp.const(dtypes.uint32, 0)) \
         | val.bitcast(dtypes.uint32)).bitcast(dtypes.int)
 
 # store for char: buf[idx/4] <- (var << (idx%4)*8))
@@ -17,7 +17,7 @@ def packed_store(bidx:UOp, var:UOp, gate:UOp|None=None):
   if var.dtype == dtypes.bool: var = var.cast(dtypes.int32)
   new_v, wmask = (var & mask).cast(dtypes.uint32) << shift_am, ((mask << shift_am) ^ 0xFFFFFFFF).cast(dtypes.uint32)
   idx = UOp(Ops.INDEX, src=(bidx.src[0], div_idx))
-  buf = UOp.load(idx, *((UOp.const(0, dtypes.uint32), gate) if gate is not None else ()), dtype=dtypes.uint32)
+  buf = UOp.load(idx, *((UOp.const(dtypes.uint32, 0), gate) if gate is not None else ()), dtype=dtypes.uint32)
   return UOp.store(idx, (buf & wmask) | new_v, *((gate,) if gate is not None else ()))
 
 # load for char: sign_extend(buf[idx/4] >> ((idx%4)*8))
@@ -68,10 +68,10 @@ class WGSLRenderer(CStyleLanguage):
 
   string_rewrite = PatternMatcher([
     (UPat(Ops.NEG, dtypes.uints, src=(UPat.var('x'))), lambda ctx,x: f"(0-{ctx[x]})"),
-    (UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.val else "false"),
+    (UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.arg else "false"),
     (UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"),
-     lambda x: f"bitcast<u32>({x.val})" if x.val < 0 else f"{x.val&0xFFFFFFFF}u"),
-    (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}"),
+     lambda x: f"bitcast<u32>({x.arg})" if x.arg < 0 else f"{x.arg&0xFFFFFFFF}u"),
+    (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}"),
     (UPat(Ops.BUFFER, name="x"), lambda ctx,x:
      f"var{'<workgroup>' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"),
     (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),

tinygrad/runtime/graph/metal.py

@@ -113,5 +113,5 @@ class MetalGraph(GraphRunner):
   @staticmethod
   def supports_uop(batch_devs, new_call:UOp) -> bool:
     # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
-    if any(b.op is Ops.SLICE and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
+    if any(b.op is Ops.SLICE and b.src[1].arg * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
     return GraphRunner.supports_uop(batch_devs, new_call)

tinygrad/runtime/ops_amd.py

@@ -909,7 +909,7 @@ class PCIIface(PCIIfaceBase):
 
 class USBIface(PCIIface):
   def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
-    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
+    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001), "AMD")):
       raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
     self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
     self.dev_impl = AMDev(self.pci_dev)