add threading to the http server

* external benchmark schedule in 5 sec (codex slop)

* prune

* real?

* delete

.github/workflows/benchmark.yml

@@ -515,14 +515,12 @@ jobs:
       run: |
         echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - 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: 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: Test copy speeds
       run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 

.github/workflows/test.yml

@@ -353,6 +353,8 @@ 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
 
@@ -641,12 +643,8 @@ 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)
@@ -657,6 +655,25 @@ 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,10 +22,6 @@ 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)]
@@ -38,12 +34,8 @@ class Attention:
     # update the cache
     self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
 
-    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
+    keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
+    values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
 
     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))
@@ -86,15 +78,14 @@ class Transformer:
       seqlen = tokens.shape[1]
       tok_emb = self.wte(tokens)
 
-    # 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)))
+    # 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))))
 
     h = tok_emb + pos_emb
 
     if HALF: h = h.half()
 
-    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
+    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None
 
     for hi in self.h: h = hi(h, start_pos, mask)
 

examples/mlperf/model_train.py

@@ -1711,9 +1711,10 @@ def train_gptoss():
     wandb.init(config=config, **wandb_args, project="MLPerf-gpt-oss")
 
   model_params = GPT_OSS_20B
-  model_params['vocab_size'] = 128256
+  model_params['vocab_size'] = getenv("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)
@@ -1748,7 +1749,10 @@ def train_gptoss():
 
   from extra.gemm.cdna_asm_gemm import _mx_block_scale
   model_state = get_state_dict(model)
-  fp8_scale_names = {n: f"{n}_scale" for n, t in model_state.items() if t.dtype == FP8_DTYPE}
+  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_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,6 +25,7 @@ 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
@@ -44,6 +45,11 @@ 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
@@ -79,7 +85,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:
+  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
     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,
@@ -92,7 +98,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:
+  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
     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,
@@ -108,7 +114,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:
+  if FUSED_SILU_W13 and not MXFP4:
     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)
@@ -179,6 +185,9 @@ 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,10 +13,14 @@ 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.008
+INIT_STD = 0.02
+ASM_GEMM = getenv("ASM_GEMM", 0)
+
 
 def _quant_dequant_fwd(x:Tensor) -> Tensor:
   # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
@@ -59,10 +63,34 @@ 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)
@@ -99,9 +127,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)
+    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_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)
+    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_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
 
     # output
@@ -111,10 +139,15 @@ 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):
-    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 _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 _attn_mask(self, seqlen:int, dtype) -> Tensor:
     i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
@@ -173,17 +206,32 @@ 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()
-    thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
-    weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
+    dim, inter = self.dim, self.intermediate_size
 
-    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
+    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
     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
+from tinygrad.uop.ops import UOp, Ops, AxisType
 
 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.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, len(self.device)).unshard(0)).clone()
+    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()
 
   def _zero_gather(self, t:Tensor) -> Tensor:
     if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
@@ -123,6 +123,9 @@ 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
+import os, sys, pickle, time, re, tempfile, struct, shutil, io
 import numpy as np
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 
@@ -9,6 +9,39 @@ 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)
@@ -65,8 +98,7 @@ 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:
-    pickle.dump(run_onnx_jit, f)
+  with open(OUTPUT, "wb") as f: dump_pickle(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")
@@ -136,7 +168,7 @@ def bench(run, inputs):
 
 if __name__ == "__main__":
   if getenv("RUN_PICKLE"):
-    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(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)
@@ -144,7 +176,7 @@ if __name__ == "__main__":
     onnx_file = fetch(OPENPILOT_MODEL)
     inputs, outputs = compile(onnx_file)
 
-    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
 
     test_vs_compile(pickle_loaded, inputs, outputs)
     if getenv("SELFTEST"):

examples/openpilot/load_pickle.py

@@ -1,5 +1,6 @@
-import sys, pickle
+import sys
 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"
@@ -7,7 +8,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: pickle.load(open(PKL, 'rb'))
+  with WallTimeEvent(BenchEvent.STEP) as wte: load_pickle(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.arg}"
+          global_size[j] = f"_{name.expr}[0] + {val.val}"
 
   prg = ""
   if target == "clang":

extra/gemm/cdna_asm_gemm.py

@@ -1,8 +1,9 @@
-import atexit, functools, pathlib
+import atexit, functools, math, 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
+from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
 
@@ -107,6 +108,23 @@ 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
@@ -119,6 +137,50 @@ 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()
@@ -174,10 +236,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(dtypes.weakint, K)+k))*
-         B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
+  mul = (A.flatten().index((m*UOp.const(K)+k))*
+         B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
   red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
-  store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n)).store(red).end(m, n)
+  store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
   return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
 
 # ** bf16 A @ B.T kernel in C
@@ -244,7 +306,6 @@ 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):
@@ -341,13 +402,29 @@ 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) -> Tensor:
+             a_pretranspose:Tensor|None=None, mxfp4:bool=False) -> 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:
@@ -355,7 +432,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 else a.dtype
+  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE or mxfp4 else a.dtype
 
   batch, M, K = a.shape
   N = b.shape[1]
@@ -378,7 +455,21 @@ 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 mx:
+    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:
       # mxfp8 1x32 block scaling
       if mx_scales is not None:
         a_si, a_e8, b_si, b_e8 = mx_scales

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(dtypes.float, (0.0,)*4), end=init_l)
+  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const((0.0,)*4, dtypes.float), 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(out.dtype, 0.0)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
+  return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).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,55 +1,25 @@
 """
 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
-
-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.
+  @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
 """
 
 from tinygrad.dtype import dtypes, AddrSpace, DType
@@ -61,25 +31,17 @@ from tinygrad.tensor import Tensor
 # tilelang builtins, expressed with tinygrad UOp APIs
 # ---------------------------------------------------------------------------
 
-def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp:
+def alloc_shared(shape:tuple[int, ...], dtype:DType, slot:int) -> UOp:
   """T.alloc_shared: one LOCAL buffer shared by all threads in the block."""
-  return UOp.placeholder(tuple(shape), dtype, next(UOp.unique_num), AddrSpace.LOCAL)
+  return UOp.placeholder(tuple(shape), dtype, slot, AddrSpace.LOCAL)
 
-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.
-  """
+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."""
   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, next(UOp.unique_num), AddrSpace.REG)
+  fragment = UOp.placeholder(shard_shape, dtype, slot, AddrSpace.REG)
   return fragment.unshard(axes, rngs)
 
 # ---------------------------------------------------------------------------
@@ -106,64 +68,45 @@ 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)
-  # 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.
+
+  # 16*8 threads = 128 threads
   tx = UOp.range(TX, 2, AxisType.LOCAL)
   ty = UOp.range(TY, 3, AxisType.LOCAL)
 
-  # 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)
+  # 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))
 
-  # 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))
+  # zero out the regs to start. this is expanded by the devectorizer
+  C_local = C_local.after(C_local.store(0.0))
 
   # 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)
 
-  # 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)
+  # 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(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.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)))
 
-  # 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)
+  # 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)
   jj = UOp.range(TN, 13, AxisType.UPCAST)
-  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)
+  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)
   # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
-  C_loc = C_loc[ir*TM + ty, tx*TN + jj].set(acc, end=(kk, ir, jj, ko))
+  C_local = C_local[ir, ty, tx, jj].set(acc, end=(kk, ir, jj, ko))
 
-  # 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))
+  # 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))
 
-  # 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=()))
+  # close the locals and globals
+  return c_st.end(tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
 
 # ---------------------------------------------------------------------------
 # python wrapper: same signature as the tilelang function
@@ -188,7 +131,8 @@ if __name__ == "__main__":
   b = Tensor.randn(K, N, dtype=dtype_in).contiguous()
   ref = (a @ b).relu().realize()
 
-  out = matmul_relu(a, b).realize()
+  for _ in range(10):
+    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(dtypes.uint32, x)
+  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
     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,18 +157,19 @@ 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(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()
+  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()
   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(dtypes.weakint, size_dw)), 2), (submitted + 1).cast(dtypes.uint64))
+  bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(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(dtypes.int, 0))
+  put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
   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(dtypes.uint32, x)) for off,x in enumerate(pkt)])
+  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)])
 
   # advance the put/write pointers past the packet
   bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
@@ -186,26 +187,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(dtypes.uint32, x) for off in range(0, sz, ctx.max_copy_size) for x in (
+  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 (
     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(dtypes.uint32, x) for x in (
+  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) 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(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)))))
+    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)))))
 
 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(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
+  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) 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),
@@ -218,7 +219,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(dtypes.int, 0)
+  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int)
 
   # the sdma queue's ring and its host-side ring/write/put pointers
   for d in devs: q = Device[d].sdma_queue(0)
@@ -234,8 +235,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(dtypes.uint32, 0)).end(zi)
-  i = UOp.range(UOp.const(dtypes.int, size_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,))
   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(dtypes.weakint, 0))
+  amax_idx = amax_out.reshape((1,)).index(UOp.const(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(dtypes.uint32, v)
-  return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(dtypes.uint64, 0), 'laneId': u32()}
+  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()}
 
 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(dtypes.uint32, 0x3f800000)  # 1.0f
-    s1 = UOp.const(dtypes.uint32, 0x40000000)  # 2.0f
+    s0 = UOp.const(0x3f800000, dtypes.uint32)  # 1.0f
+    s1 = UOp.const(0x40000000, dtypes.uint32)  # 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(dtypes.uint32, 0x40000000)  # 2.0f
-    s1 = UOp.const(dtypes.uint32, 0x40400000)  # 3.0f
+    s0 = UOp.const(0x40000000, dtypes.uint32)  # 2.0f
+    s1 = UOp.const(0x40400000, dtypes.uint32)  # 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', {}).arg, 0)
-    self.assertEqual(parse_expr('42', {}).arg, 42)
-    self.assertEqual(parse_expr('42U', {}).arg, 42)
+    self.assertEqual(parse_expr('0', {}).val, 0)
+    self.assertEqual(parse_expr('42', {}).val, 42)
+    self.assertEqual(parse_expr('42U', {}).val, 42)
 
   def test_negative_integers(self):
     """Test parsing negative integer literals."""
     result = parse_expr('-1', {})
-    self.assertEqual(result.arg, -1)
+    self.assertEqual(result.val, -1)
     self.assertEqual(result.dtype, dtypes.int)
 
   def test_float_literals(self):
     """Test parsing float literals."""
     result = parse_expr('1.0F', {})
-    self.assertEqual(result.arg, 1.0)
+    self.assertEqual(result.val, 1.0)
     self.assertEqual(result.dtype, dtypes.float32)
 
   def test_hex_literals(self):
     """Test parsing hex literals."""
     result = parse_expr('0xFF', {})
-    self.assertEqual(result.arg, 255)
+    self.assertEqual(result.val, 255)
 
   def test_variable_lookup(self):
     """Test variable lookup in parse_expr."""
-    vrs = {'x': UOp.const(dtypes.uint32, 42)}
+    vrs = {'x': UOp.const(42, dtypes.uint32)}
     result = parse_expr('x', vrs)
-    self.assertEqual(result.arg, 42)
+    self.assertEqual(result.val, 42)
 
   def test_binary_ops(self):
     """Test parsing binary operations."""
-    vrs = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)}
+    vrs = {'a': UOp.const(10, dtypes.uint32), 'b': UOp.const(5, dtypes.uint32)}
 
     # 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.arg, 5)
+    self.assertEqual(result.val, 5)
 
   def test_ternary(self):
     """Test parsing ternary expressions."""
-    vrs = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)}
+    vrs = {'cond': UOp.const(True), 'a': UOp.const(1, dtypes.uint32), 'b': UOp.const(0, dtypes.uint32)}
     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(dtypes.uint32, 0xFFFFFFFF)  # All ones - CLZ should be 0
+    S0 = UOp.const(0xFFFFFFFF, dtypes.uint32)  # 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(dtypes.uint32, 0)
+    S0 = UOp.const(0, dtypes.uint32)
     _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.arg, -1)
+    self.assertEqual(val.val, -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(dtypes.uint32, 1)  # LSB set - CTZ should be 0
+    S0 = UOp.const(1, dtypes.uint32)  # 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(dtypes.uint64, 0),
-      'DATA': UOp.const(dtypes.uint32, 0x3f800000),
+      'ADDR': UOp.const(0, dtypes.uint64),
+      'DATA': UOp.const(0x3f800000, dtypes.uint32),
       '_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(dtypes.uint32, 0)
-    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}
+    addr = UOp.const(0, dtypes.uint32)
+    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)}
 
     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(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),
+      '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),
     }
-    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
+    srcs['laneId'] = UOp.const(0, dtypes.uint32)
     _, 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(dtypes.uint32, 0),
-      'OFFSET0': UOp.const(dtypes.uint32, 0),
-      'OFFSET1': UOp.const(dtypes.uint32, 1),
+      'ADDR': UOp.const(0, dtypes.uint32),
+      'OFFSET0': UOp.const(0, dtypes.uint32),
+      'OFFSET1': UOp.const(1, dtypes.uint32),
       '_lds': lds,
     }
-    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
+    srcs['laneId'] = UOp.const(0, dtypes.uint32)
     _, 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(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),
+      '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),
     }
-    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
+    srcs['laneId'] = UOp.const(0, dtypes.uint32)
     _, 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().arg, 108)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][0].simplify().arg, 120)  # type: ignore[index]
+    self.assertEqual(assigns[0][1][0].simplify().val, 108)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][0].simplify().val, 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(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),
+      '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),
     }
-    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
+    srcs['laneId'] = UOp.const(0, dtypes.uint32)
     _, 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().arg, 0xAAAAAAAA)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][1].simplify().arg, 0xBBBBBBBB)  # type: ignore[index]
+    self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][1].simplify().val, 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(dtypes.uint32, 10)
-    s1 = UOp.const(dtypes.uint32, 20)
-    scc = UOp.const(dtypes.uint32, 1)
+    s0 = UOp.const(10, dtypes.uint32)
+    s1 = UOp.const(20, dtypes.uint32)
+    scc = UOp.const(1, dtypes.uint32)
     _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(dtypes.uint32, row)})
-      self.assertEqual(parsed.simplify().arg, expected)
+      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)})
+      self.assertEqual(parsed.simplify().val, 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(dtypes.uint32, lane)})
-      self.assertEqual(parsed.simplify().arg, expected)
+      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)})
+      self.assertEqual(parsed.simplify().val, expected)
 
   def test_permlane64_wave64_pcode_indices(self):
     vgpr = UOp.param(0, dtypes.uint32, (256,))
     srcs = {
-      '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),
+      '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),
       '_vgpr': vgpr,
       '_wave_size': 64,
-      'S0': UOp.const(dtypes.uint32, 0),
-      'S1': UOp.const(dtypes.uint32, 0),
-      'S2': UOp.const(dtypes.uint32, 0),
+      'S0': UOp.const(0, dtypes.uint32),
+      'S1': UOp.const(0, dtypes.uint32),
+      'S2': UOp.const(0, dtypes.uint32),
     }
 
     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.arg
+      return idx.val
 
     _, 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().arg, dst_idx)  # type: ignore[index]
+      self.assertEqual(assigns[lane][1][0].simplify().val, 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(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
+    u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64)
     lds = UOp.param(3, dtypes.uint32, (16384,))
     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(dtypes.bool, True)}
+            '_lds': lds, '_vmem': lds, '_active': UOp.const(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
+from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm, quantize_mxfp4
 from test.helpers import needs_second_gpu
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
 
@@ -150,6 +150,45 @@ 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(dtypes.int, 3), C[0], ret_dtype=dtypes.void)
+  call = F[0].load().call(UOp.const(3).cast(dtypes.int), 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(dtypes.int, 21), ret_dtype=dtypes.int)
+  val = F[0].load().call(UOp.const(21).cast(dtypes.int), 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(dtypes.float, 2.0)).contiguous()
+    t = Tensor(UOp.const(2.0, dtypes.float)).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(dtypes.uint64, 5), UOp.const(dtypes.uint64, 10))
+    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
     self.assertIs(x.simplify().op, Ops.CONST)
 
 class TestTautologicalCompare(unittest.TestCase):

test/backend/test_custom_kernel.py

@@ -1,5 +1,6 @@
 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
@@ -340,7 +341,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(UOp.const(dtypes.float, Invalid, shape=(2,))))
+    after = x.uop.after(x.uop.shrink(((0, 2),)).store(Invalid))
     self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.])
 
   def test_multi_after_invalid_store_dep_removed(self):
@@ -527,6 +528,77 @@ 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(dtypes.int, i), dtype=y.dtype.scalar())
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), 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].arg == 4)
+    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 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].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
+    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
 
   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(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)
+    c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL)
     c4 = UOp.param(1, dtypes.int, (512,))
-    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)
+    c5 = c4.index(c1.valid(UOp.const(True)))
+    c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE)
     c9 = UOp.param(2, dtypes.uchar, (47040000,))
-    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)
+    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)
     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,6 +76,13 @@ 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,6 +1728,8 @@ 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(dtypes.int, 2)
+    sink = UOp.const(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(dtypes.int, 0)
+  idx = UOp.const(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(dtypes.int, 1)))
+    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int)))
     sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
     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(dtypes.int, 1)))
+    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)))
     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, dtypes.int.min+1))
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int.min+1).cast(dtypes.int))
     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(dtypes.float32, float("inf")))
+    ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(float("inf")).cast(dtypes.float32))
     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(dtypes.int, 1)
+    val = UOp.const(1).cast(dtypes.int)
     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(dtypes.float, 2.0))
+    t = Tensor(UOp.const(2.0).cast(dtypes.float))
     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(dtypes.float, 2.0))
+    t = Tensor(UOp.const(2.0).cast(dtypes.float))
     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(dtypes.float, 2.0)).clone()
+    t = Tensor(UOp.const(2.0).cast(dtypes.float)).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(dtypes.float, 2.0)).numpy(), 2.0)
+    np.testing.assert_equal(Tensor(UOp.const(2.0).cast(dtypes.float)).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(dtypes.float, 2.0)).clone(device="CPU")
+    t = Tensor(UOp.const(2.0).cast(dtypes.float)).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(dtype, arg))
+  if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype))
   elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
   else: uops.append(UOp(op, dtype, tuple(src), arg))
   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(dtypes.int32, 0)].store(alu)
+  out = buf_store[UOp.const(0).cast(dtypes.int32)].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(dtypes.int, 2)
-    c2 = UOp.const(dtypes.int, 3)
+    c1 = UOp.const(2)
+    c2 = UOp.const(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(dtypes.int, 0)).store(a1), out.index(UOp.const(dtypes.int, 1)).store(a2)],
+    uops = to_uops_list([out.index(UOp.const(0)).store(a1), out.index(UOp.const(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(dtypes.int, 0)
-    c2 = UOp.const(dtypes.int, 1)
-    expr = g1.index(c1) * UOp.const(dtypes.int, 4096) + g1.index(c2)
+    c1 = UOp.const(0)
+    c2 = UOp.const(1)
+    expr = g1.index(c1) * UOp.const(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(dtypes.uint, 7)
+    c = UOp.const(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(dtypes.weakint, Invalid))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  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))
   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(dtypes.weakint, Invalid))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  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))
   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(dtypes.weakint, random.randint(-10, 10))])
+  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(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(dtypes.weakint, {n1.as_long()}), UOp.const(dtypes.weakint, {n2.as_long()})," +\
-                f"UOp.const(dtypes.weakint, {n3.as_long()})\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" +\
             "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(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
+  prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
   prog = dev.runtime(prg.to_elf())
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]

test/mockgpu/amd/emu.py

@@ -192,16 +192,16 @@ def _init_sqtt_encoder():
 
   return emit, finish, finalize
 
-def _c(val, dtype=dtypes.uint32): return UOp.const(dtype, val)
+def _c(val, dtype=dtypes.uint32): return UOp.const(val, dtype)
 
 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(dtypes.uint64, 32))
+  return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
 
 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(dtypes.uint64, 32)).cast(dtypes.uint32)
+  return v64.cast(dtypes.uint32), (v64 >> UOp.const(32, dtypes.uint64)).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(ut, mask)).bitcast(ft)
+  if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(mask, ut)).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(dtypes.uint64, 1)).ne(UOp.const(dtypes.uint64, 0))
+  if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(1, dtypes.uint64)).ne(UOp.const(0, dtypes.uint64))
   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(dt, 1) << lane.cast(dt)
+    mask = UOp.const(1, dt) << lane.cast(dt)
     new_bit = _to_u32(val).cast(dt) << lane.cast(dt)
-    cleared = old.cast(dt) & (mask ^ UOp.const(dt, 0xFFFFFFFFFFFFFFFF))
+    cleared = old.cast(dt) & (mask ^ UOp.const(0xFFFFFFFFFFFFFFFF, dt))
     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,7 +365,7 @@ 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(reg_or_addr.dtype, incr) if isinstance(reg_or_addr, UOp) else incr), hi, *args)]
+  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)]
 
 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."""
@@ -374,7 +374,7 @@ def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = Fals
 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."""
   adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32
-  word_addr = addr >> UOp.const(adt, 2)
+  word_addr = addr >> UOp.const(2, adt)
   idx = mem.index(word_addr.valid(active))
   if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))]
   # Sub-word store: read-modify-write with mask
@@ -388,7 +388,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(adt, 1)).valid(active & is_cross))
+  next_idx = mem.index((word_addr + UOp.const(1, adt)).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 +397,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(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)))
+    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)))
   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 +462,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(dtypes.uint64, dword_idx * 4)
-    return self.vmem.index(addr >> UOp.const(dtypes.uint64, 2)).load()
+    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()
 
   def inst_field(self, field) -> UOp:
     """Extract field bits from instruction encoding. Tracks field for canonical key computation."""
@@ -475,15 +475,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(dtypes.uint32, lo_in_dword)
-      return shifted & UOp.const(dtypes.uint32, mask)
+      shifted = word if lo_in_dword == 0 else word >> UOp.const(lo_in_dword, dtypes.uint32)
+      return shifted & UOp.const(mask, dtypes.uint32)
     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(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))
+      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))
 
   def inst_field_signed(self, field) -> UOp:
     """Extract field and sign-extend based on field width."""
@@ -584,7 +584,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(dtypes.uint64, 32) if is_f64 else literal.cast(dtypes.uint64)
+        lit_val = literal.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64) 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 +602,7 @@ class _Ctx:
 
   def inc_pc(self) -> list[UOp]:
     """Increment PC by instruction size in bytes. Returns [store]."""
-    new_pc = self.rpc() + UOp.const(dtypes.uint64, self.inst_size)
+    new_pc = self.rpc() + UOp.const(self.inst_size, dtypes.uint64)
     lo, hi = _split64(new_pc)
     return [self.wsgpr_dyn(_c(PC_LO_IDX), lo), self.wsgpr_dyn(_c(PC_HI_IDX), hi)]
 
@@ -668,7 +668,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(dtypes.float32, 3.4028234663852886e38),
+                  'MAX_FLOAT_F32': UOp.const(3.4028234663852886e38, dtypes.float32),
                   # 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 +709,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].arg), int(val[3].arg)
+        hi_bit, lo_bit = int(val[2].val), int(val[3].val)
         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 +732,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(dtypes.uint32, slice_mask)
+                       val.cast(dtypes.uint32) & UOp.const(slice_mask, dtypes.uint32)
             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(val.dtype, 0.0)).minimum(UOp.const(val.dtype, 1.0))
-          val = _FUNCS['isNAN'](val).where(UOp.const(val.dtype, 0.0), clamped)
+          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)
         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(dtypes.uint32, 0xFFFF)) | (result << UOp.const(dtypes.uint32, 16))
+          hi_result = (old_val & UOp.const(0xFFFF, dtypes.uint32)) | (result << UOp.const(16, dtypes.uint32))
           # GFX9/CDNA zeroes upper 16 bits on lo-half write; RDNA preserves them
-          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))
+          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))
           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 +767,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(dtypes.uint32, ((1 << width) - 1) << lo_bit)
-        result = (result & (mask ^ UOp.const(dtypes.uint32, 0xFFFFFFFF))) | (val_bits << UOp.const(dtypes.uint32, lo_bit))
+        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))
       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 +790,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(dtypes.uint32, 0xFFFFFFFF)),
-                          ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)))
+    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)))
   # 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 +804,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(vcc.dtype, 0)).cast(dtypes.uint32),
-            'EXECZ': exec_val.eq(UOp.const(exec_val.dtype, 0)).cast(dtypes.uint32)}
+            'VCCZ': vcc.eq(UOp.const(0, vcc.dtype)).cast(dtypes.uint32),
+            'EXECZ': exec_val.eq(UOp.const(0, exec_val.dtype)).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 +833,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(dtypes.uint64, 2)
-  vals = [ctx.vmem.index(dword_base + UOp.const(dtypes.uint64, i)) for i in range(ndwords)]
+  dword_base = addr >> UOp.const(2, dtypes.uint64)
+  vals = [ctx.vmem.index(dword_base + UOp.const(i, dtypes.uint64)) for i in range(ndwords)]
   if abs(nval) < 1:
     nbits = int(abs(nval) * 32)
-    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)
+    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)
     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 +941,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(dtypes.bool, True)
+  src_lane, valid = lane_i, UOp.const(True)
 
   if op == 'quad_perm':
     assert isinstance(arg, tuple)
@@ -967,7 +967,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(fallback.dtype, 0) if getattr(inst, 'bc', 0) else fallback
+  invalid = UOp.const(0, fallback.dtype) 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 +1159,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(dtypes.uint64, 0)})[1]:
+    for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(0, dtypes.uint64)})[1]:
       if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32)
     return _c(0)
 
@@ -1189,8 +1189,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(dt, mask)
-  result = UOp.const(dt, 0)
+  def bnot(v): return v ^ UOp.const(mask, dt)
+  result = UOp.const(0, dt)
   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 +1244,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(dtypes.uint32, opsel)}
+  srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(opsel, dtypes.uint32)}
   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 +1398,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(dtypes.int, n_a_elems)
+  b_off = UOp.const(n_a_elems, dtypes.int)
   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 +1407,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(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)
+      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)
     elif is_f32_src:
       return raw  # already uint32 (f32 bit pattern)
     elif is_fp8:
-      return ((raw >> UOp.const(dtypes.uint32, sub_idx * 8)) & UOp.const(dtypes.uint32, 0xFF)).cast(dtypes.uint32)
+      return ((raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)).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(dtypes.uint32, sub_idx * 16)) & UOp.const(dtypes.uint32, 0xFFFF)) << UOp.const(dtypes.uint32, 16)
+      return ((raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)) << UOp.const(16, dtypes.uint32)
     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(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)
+      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)
       # 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(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))
+      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))
       # Return uint32 (f32 bit pattern) — stored directly to uint32 temp array, bitcast to float on read
-      return is_zero.where(UOp.const(dtypes.uint32, 0), f32_bits)
+      return is_zero.where(UOp.const(0, dtypes.uint32), 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(dtypes.int, grp_size)
-  grp_idx = read_lane // UOp.const(dtypes.int, grp_size)
+  lane_in_grp = read_lane % UOp.const(grp_size, dtypes.int)
+  grp_idx = read_lane // UOp.const(grp_size, dtypes.int)
 
   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(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)
+    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)
   else:
     mn_idx = lane_in_grp  # for 16x16 and 4x4
 
@@ -1456,18 +1456,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(dtypes.int, M * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl)
+      a_idx = grp_idx * UOp.const(M * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
     else:
-      a_idx = mn_idx * UOp.const(dtypes.int, K) + grp_idx * UOp.const(dtypes.int, k_per_grp) + UOp.const(dtypes.int, kl)
+      a_idx = mn_idx * UOp.const(K, dtypes.int) + grp_idx * UOp.const(k_per_grp, dtypes.int) + UOp.const(kl, dtypes.int)
     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(dtypes.int, N * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl)
+      b_idx = b_off + grp_idx * UOp.const(N * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
     else:
-      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)
+      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)
     read_stores.append(tmp.index(b_idx).store(b_val))
 
   read_phase = UOp.group(*read_stores).end(read_lane)
@@ -1488,11 +1488,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(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
+    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
 
     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 +1503,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(dtypes.int, 16) + UOp.const(dtypes.int, (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)
 
       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 +1511,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(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)
+        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)
         acc = acc + a_val * b_val
 
       if is_int_out:
@@ -1523,8 +1523,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(dtypes.int, grp_sub)
-    c_grp = compute_lane // UOp.const(dtypes.int, grp_sub)
+    n_idx = compute_lane % UOp.const(grp_sub, dtypes.int)
+    c_grp = compute_lane // UOp.const(grp_sub, dtypes.int)
 
     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 +1534,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(dtypes.int, M * K) + UOp.const(dtypes.int, out_reg * K)
+        m_base = c_grp * UOp.const(M * K, dtypes.int) + UOp.const(out_reg * K, dtypes.int)
         for k in range(K):
-          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)
+          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)
           acc = acc + a_val * b_val
       else:
         # 16x16: K is split across groups. Shared MxK/NxK arrays.
-        m_base = c_grp * UOp.const(dtypes.int, out_per_lane) + UOp.const(dtypes.int, out_reg)
+        m_base = c_grp * UOp.const(out_per_lane, dtypes.int) + UOp.const(out_reg, dtypes.int)
         for k in range(K):
-          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)
+          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)
           acc = acc + a_val * b_val
 
       if is_int_out:
@@ -1570,8 +1570,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(dtypes.int, lane))
-    return cvt((v >> UOp.const(dtypes.uint32, 16)) if half else (v & UOp.const(dtypes.uint32, 0xFFFF)))
+    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)))
 
   # 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 +1593,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(dtypes.uint32, 16)) & UOp.const(dtypes.uint32, 0xFFFF)
+    def f32_to_bf16_bits(v: UOp) -> UOp: return (v.bitcast(dtypes.uint32) >> UOp.const(16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)
     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(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)
+      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)
                 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(dtypes.int, d_map(m, n)[0]), out_cvt(mat_d[m*16+n]), exec_mask)
+      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)
                 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(dtypes.int, d_map(m, n)[0])).bitcast(dtypes.float32)
+    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)
              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(dtypes.int, d_map(m, n)[0]), mat_d[m*16+n].bitcast(dtypes.uint32), exec_mask)
+    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)
               for m in range(16) for n in range(16)]
   return UOp.sink(*stores, *ctx.inc_pc())
 
@@ -1688,8 +1688,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(dtypes.uint32, 0x80000000)
-      if neg_hi_bit: hi = hi ^ UOp.const(dtypes.uint32, 0x80000000)
+      if neg_lo: lo = lo ^ UOp.const(0x80000000, dtypes.uint32)
+      if neg_hi_bit: hi = hi ^ UOp.const(0x80000000, dtypes.uint32)
       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 +1700,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(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)
+      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)
     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(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
+      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
     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(dtypes.uint32, combined_opsel_hi), 'OPSEL': UOp.const(dtypes.uint32, opsel)}
+            'OPSEL_HI': UOp.const(combined_opsel_hi, dtypes.uint32), 'OPSEL': UOp.const(opsel, dtypes.uint32)}
   else:
     def get_half_bits(val: UOp, use_hi: bool, apply_neg: bool = False) -> UOp:
-      bits = ((val >> UOp.const(dtypes.uint32, 16)) if use_hi else val) & UOp.const(dtypes.uint32, 0xFFFF)
+      bits = ((val >> UOp.const(16, dtypes.uint32)) if use_hi else val) & UOp.const(0xFFFF, dtypes.uint32)
       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(dtypes.uint32, 16))
+      return lo | (hi << UOp.const(16, dtypes.uint32))
     # 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(dtypes.uint32, neg)
+    if is_dot_iu: srcs['NEG'] = UOp.const(neg, dtypes.uint32)
   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 +1781,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(dtypes.uint32 if is_lds else dtypes.uint64, 2)
+  addr_shift = UOp.const(2, dtypes.uint32 if is_lds else dtypes.uint64)
 
   # Extract register info - all dynamic for deduplication
   if is_lds:
@@ -1831,20 +1831,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(dtypes.bool, False)
+    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(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(dtypes.uint64, 0)
+      addr_offset = vaddr if sve == 1 else UOp.const(0, dtypes.uint64)
       # 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(dtypes.uint64, 0)) \
-        if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
+      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)
       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(dtypes.uint64, 0)
+    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)
     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 +1874,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(dtypes.uint32, 0)}
+        data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA2': ctx.rvgpr_dyn(data1_reg, lane) if has_data1 else UOp.const(0, dtypes.uint32)}
       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(dtypes.uint64, 0)}
+                '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)}
       # 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(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))
+    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))
     # 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 +1900,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(dtypes.uint32, 0)
+      else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(0, dtypes.uint32)
     if 'STORE' in op_name and data_bits_mem >= 64:
-      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
+      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
     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(dtypes.uint32, 0)
+      srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(0, dtypes.uint32)
     return srcs
 
   def make_stores(dest: str, val: UOp, lane: UOp, active: UOp, writes_return_data: bool) -> list[UOp]:
@@ -1983,7 +1983,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(dtypes.uint64, 0xFFFFFFFFFFFF)
+  base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(0xFFFFFFFFFFFF, dtypes.uint64)
   num_records = ctx.rsgpr_dyn(srsrc + _c(2))
   stride = (ctx.rsgpr_dyn(srsrc + _c(3)) & _c(0x3FFF)).cast(dtypes.uint64)
 
@@ -2000,7 +2000,7 @@ 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(dtypes.uint64, 0))  # safe address when OOB
+  addr = in_bounds.where(addr, UOp.const(0, dtypes.uint64))  # safe address when OOB
   mem = ctx.vmem
 
   stores: list[UOp] = []
@@ -2009,20 +2009,20 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
     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(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
+      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
       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(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
+      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
       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(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
+      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
       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())

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(dt, v)
+def _const(dt, v): return UOp.const(v, dt)
 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.arg, int) and src.arg > 0 and (src.arg & (src.arg + 1)) == 0:
-        widths.append(src.arg.bit_length())
+      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 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.arg - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
+    if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
     if s.op == Ops.MUL and len(s.src) == 2:
-      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]
+      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]
       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].arg == 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].val == 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.arg - right.arg)
+        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val)
         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.arg == 2.0 else left
+      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 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.arg)
+        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val)
       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.arg)
+        w = int(width.val)
         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.arg), int(second.arg)
+        a, b = int(first.val), int(second.val)
         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.arg)
+      idx = int(first.val)
       # 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(dtypes.uint64, 32))
-        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(dtypes.uint64, 8))).cast(dtypes.uint32), val)
+        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)
     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.arg < 0: l = l.cast(dtypes.int)
+      if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 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.arg)
+        return int(expr.val)
       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.arg is v
+      def is_const(c, v): return c.op == Ops.CONST and c.val 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(dtypes.bool, True)
+          remaining = UOp.const(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, Invalid
+from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -36,21 +36,20 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
 
 class TestWeakConstFolding(unittest.TestCase):
   def test_weakint_math(self):
-    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))
+    out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
+    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakint, 2**41))
 
   def test_float_unaries(self):
-    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))
+    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))
 
   def test_weakfloat_math(self):
-    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))
+    out = (UOp.const(1.25) + UOp.const(2.5)).simplify()
+    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakfloat, 3.75))
 
   def test_invalid_poison(self):
-    self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid)
+    self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
 
 class TestBinaryOpsConstFolding(unittest.TestCase):
   def test_add_literal_zero(self):
@@ -121,10 +120,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_dt, from_v).bitcast(to_dt).sink()).src[0]
+          r = full_rewrite(UOp.const(from_v, from_dt).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.arg, to_v, msg)
+          np.testing.assert_equal(r.val, to_v, msg)
 
     t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
     t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})
@@ -145,9 +144,9 @@ class TestBitcastConstFolding(unittest.TestCase):
 
   def test_vec_bitcast(self):
     with Context(SPEC=0):
-      srcs = full_rewrite(UOp.const(dtypes.int32, (-1, -2**31, 75)).bitcast(dtypes.uint32).sink()).src
+      srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src
     self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
-    self.assertEqual(tuple(x.arg for x in srcs), (2**32-1, 2**31, 75))
+    self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75))
 
 # folds advance indexing into basic indexing
 class TestIndexingConstFolding(unittest.TestCase):

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].arg for x in loop_idxs]
+    sizes = [x.src[0].val 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(dtypes.weakint, 0)
+    flat = UOp.const(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(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo())
+    sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
     class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
     specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
     self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)

test/null/test_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(dtypes.float, 1.0), set([x]))[x]
+    gx = compute_gradient(f(x), UOp.const(1.0), set([x]))[x]
 
     for val in [-5., -2.0, 0.0, 2.0, 5.]:
-      tg_out = gx.substitute({x: x.const_like(val)}).ssimplify()
+      tg_out = gx.substitute({x: UOp.const(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(dtypes.float, 1.0), set([x, y]))
+    grads = compute_gradient(f(x, y), UOp.const(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: x.const_like(valx), y: y.const_like(valy)}
+        substitutions = {x: UOp.const(valx), y: UOp.const(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].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)
+    if srcs[0].op is Ops.CONST: return (srcs[0].val,)
+    if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src)
+  return tuple(s.val for s in srcs)
 
 def evaluate_uop(uop, variables):
   if uop.op == Ops.CONST:
-    return uop.arg
+    return uop.val
   elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU:
     return variables[uop.expr]
   elif uop.op in GroupOp.ALU:
@@ -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(dtypes.float32, 10.0) / UOp.const(dtypes.float32, 0.0))
+    optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0))
     self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertTrue(math.isinf(optimized_div_uop.arg) or math.isnan(optimized_div_uop.arg))
+    self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val))
 
   def test_full_graph_rewrite_redundant_operations(self):
-    optimized_uop = apply_rewrite((UOp.const(dtypes.float32, 10.0) + UOp.const(dtypes.float32, 0.0)) * UOp.const(dtypes.float32, 1.0))
+    optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0))
     self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.arg, 10.0)
+    self.assertEqual(optimized_uop.val, 10.0)
 
   def test_full_graph_rewrite_large_graph(self):
-    prev_uop = UOp.const(dtypes.int32, 0)
+    prev_uop = UOp.const(0)
     for i in range(1, 101):
-      prev_uop += UOp.const(dtypes.int32, i)
+      prev_uop += UOp.const(i)
     optimized_uop = apply_rewrite(prev_uop)
     self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.arg, sum(range(1, 101)))
+    self.assertEqual(optimized_uop.val, sum(range(1, 101)))
 
   def test_full_graph_rewrite_division_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(dtypes.float32, 42.0) / UOp.const(dtypes.float32, 1.0))
+    optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0))
     self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.arg, 42.0)
+    self.assertEqual(optimized_uop.val, 42.0)
 
   def test_full_graph_rewrite_modulo_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(dtypes.int32, 42) % UOp.const(dtypes.int32, 1))
+    optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
     self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.arg, 0)
+    self.assertEqual(optimized_uop.val, 0)
 
 
 class TestFoldingAndReduction(unittest.TestCase):
   @unittest.skip("reduce is removed now")
   def test_full_graph_rewrite_constant_reduction_folding(self):
-    const1 = UOp.const(dtypes.int32, 5)
-    const2 = UOp.const(dtypes.int32, 10)
-    const3 = UOp.const(dtypes.int32, 20)
+    const1 = UOp.const(5)
+    const2 = UOp.const(10)
+    const3 = UOp.const(20)
     optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
     expected_sum = 5 + 10 + 20
-    self.assertEqual(optimized_sink.arg, expected_sum)
+    self.assertEqual(optimized_sink.val, expected_sum)
 
   @unittest.skip("reduce is removed now")
   def test_full_graph_rewrite_reduction_with_unused_range(self):
-    const1 = UOp.const(dtypes.int32, 15)
-    const2 = UOp.const(dtypes.int32, 25)
+    const1 = UOp.const(15)
+    const2 = UOp.const(25)
     rng = UOp.range(10, idx=0)
     optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
     expected_sum = 10 * (15 + 25)
-    self.assertEqual(optimized_sink.arg, expected_sum)
+    self.assertEqual(optimized_sink.val, expected_sum)
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_range_reduction(self):
     simple_range = UOp.range(5, idx=0)
     optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
     expected_sum = sum(range(5))
-    self.assertEqual(optimized_sink.arg, expected_sum)
+    self.assertEqual(optimized_sink.val, expected_sum)
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_simple_reduction_folding(self):
     simple_range = UOp.range(4, idx=0)
-    add_uop = simple_range + UOp.const(dtypes.int32, 1)
+    add_uop = simple_range + UOp.const(1)
     optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
     expected_sum = sum(i + 1 for i in range(4))
-    self.assertEqual(optimized_sink.arg, expected_sum)
+    self.assertEqual(optimized_sink.val, 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.arg, sum((i * 10) + j for i in range(8) for j in range(4)))
+    self.assertEqual(optimized_reduce_uop.val, 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.arg, 2)
+    self.assertEqual(optimized_mod_uop.val, 2)
 
   def test_full_graph_rewrite_division_folding_with_define_var(self):
     # index dtype because div-mod rules only work on index
     n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
     optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
     self.assertEqual(optimized_div_uop.op, Ops.MUL)
-    self.assertEqual(optimized_div_uop.src[1].arg, 2)
+    self.assertEqual(optimized_div_uop.src[1].val, 2)
 
   def test_full_graph_rewrite_complex_mod_div_folding(self):
     # index dtype because div-mod rules only work on index
     k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
     optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
     self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertEqual(optimized_div_uop.arg, 1)
+    self.assertEqual(optimized_div_uop.val, 1)
 
   def test_graph_rewrite_div_folding_bug(self):
     lhs = UOp(Ops.ADD, src=(
-      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 32),), arg='lidx0'),)*4),
-      UOp.const(dtypes.int, (0, 256, 512, 768))))
-    rhs = UOp.const(dtypes.int, (2,)*4)
+      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)
     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(dtypes.float32, -1.0).log2().sink(UOp.const(dtypes.float32, 0.0).reciprocal()))
+    optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal()))
     optimized_log2_neg, optimized_recip_zero = optimized_sink.src
-    self.assertTrue(math.isnan(optimized_log2_neg.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}")
+    self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}")
+    self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0,
+                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}")
 
   @unittest.skip("broken")
   def test_full_graph_rewrite_modulo_negative_dividend(self):
@@ -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(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0)
+    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0)
 
   def test_gep_tuple_extraction(self):
     # GEP on a vector dtype to extract multiple elements as a vector
-    base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
+    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
     self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
 
   def test_gep_on_const_stack(self):
     # GEP on a const STACK to extract a single element
-    const_stack = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0)
+    const_stack = UOp.const((1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(const_stack.index(2)).val, 3.0)
 
   def test_gep_tuple_on_const_stack(self):
     # GEP on a const STACK using a tuple to extract multiple elements
-    const_stack = UOp.const(dtypes.float32, (7.0, 8.0, 9.0, 10.0))
+    const_stack = UOp.const((7.0, 8.0, 9.0, 10.0))
     self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
 
   def test_vectorize_multiple_elements(self):
     # Vectorizing multiple elements using GEP
-    base_vector = UOp.const(dtypes.float32, (5.0, 10.0, 15.0, 20.0))
+    base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
     vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4)))
     self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])
 
@@ -213,7 +213,7 @@ from tinygrad.uop.symbolic import symbolic_simple
 
 class TestBottomUpRewrite(unittest.TestCase):
   def test_const_folding(self):
-    a = UOp.const(dtypes.int, 5)
+    a = UOp.const(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(dtypes.int, 3)
+    a = UOp.const(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(dtypes.int, 3)
+    a = UOp.const(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.arg if x.op is Ops.CONST else "+", b))
+def bidir_append(ctx, x, b): ctx.append((x.val if x.op is Ops.CONST else "+", b))
 class TestBidirectional(unittest.TestCase):
   def test_simple(self):
-    a = UOp.const(dtypes.int, 1)
-    b = UOp.const(dtypes.int, 2)
+    a = UOp.const(1)
+    b = UOp.const(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(dtypes.int, 3)
-    b = UOp.const(dtypes.int, 4)
+    a = UOp.const(3)
+    b = UOp.const(4)
     c = a+b
-    cn = UOp.const(dtypes.int, 7)
-    d = UOp.const(dtypes.int, 2)
+    cn = UOp.const(7)
+    d = UOp.const(2)
     def visit_const(c:UOp):
-      print(f"visit {c.arg}")
-      assert c.arg not in (3,4)
+      print(f"visit {c.val}")
+      assert c.val 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(dtypes.int, 3)
+    a = UOp.const(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(dtypes.int, 4))
+    self.assertIs(ret, UOp.const(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.arg if x.op is Ops.CONST else x.op)
+      ctx.append(x.val if x.op is Ops.CONST else x.op)
       return None
     pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(dtypes.int, 1)
-    b = UOp.const(dtypes.int, 2)
+    a = UOp.const(1)
+    b = UOp.const(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(dtypes.int, 3)
+    a = UOp.const(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(dtypes.int, 4))
+    self.assertIs(ret, UOp.const(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.arg if x.op is Ops.CONST else x.op)
+      ctx.append(x.val if x.op is Ops.CONST else x.op)
       return None
     pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(dtypes.int, 1)
-    b = UOp.const(dtypes.int, 2)
+    a = UOp.const(1)
+    b = UOp.const(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.arg if x.op is Ops.CONST else x.op, "bpm"))
+      ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
       return None
     def pm_visit(ctx, x):
-      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
+      ctx.append((x.val 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(dtypes.int, 1)
-    b = UOp.const(dtypes.int, 2)
+    a = UOp.const(1)
+    b = UOp.const(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.arg if x.op is Ops.CONST else x.op, "bpm"))
+      ctx.append((x.val 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.arg == 1: return x.replace(arg=10)
+      if x.op is Ops.CONST and x.val == 1: return x.replace(arg=10)
       return None
     def pm_match(ctx, x):
-      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
+      ctx.append((x.val 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(dtypes.int, 1)
-    b = UOp.const(dtypes.int, 2)
+    a = UOp.const(1)
+    b = UOp.const(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(dtypes.int, 10) + b)
+    self.assertIs(ret, UOp.const(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(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)
+    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)
 
 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(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)
+    c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK)
+    c3 = ((c1*UOp.const(32))+c2)
     c4 = UOp.param(1, dtypes.float, (163840,))
-    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))))
+    c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE)
+    c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920))))
     c7 = UOp.param(2, dtypes.float, (64,))
     c8 = c7.index(c3)
-    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()
+    c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)
     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].arg * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
+        replace_map[id(orig)] = (new.src[0], new.src[1].val * 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 dtypes, Tensor, UOp, getenv
+from tinygrad import 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(dtypes.int, 100+i)
+    for i in range(self.N): UOp.const(100+i)
 
   def test_uop_list_creation(self):
-    [UOp.const(dtypes.int, 100+i) for i in range(self.N)]
+    [UOp.const(100+i) for i in range(self.N)]
 
   def test_uop_add_2n(self):
-    a = UOp.const(dtypes.int, 2)
+    a = UOp.const(2)
     for _ in range(self.N): a = a + a
 
   def test_uop_toposort(self):
-    a = UOp.const(dtypes.int, 0)
-    for i in range(self.N): a = a + UOp.const(dtypes.int, 100+i)
+    a = UOp.const(0)
+    for i in range(self.N): a = a + UOp.const(100+i)
     self.start_time()
     self.assertEqual(len(a.toposort()), 2*self.N+1)
 
   def test_uop_toposort_2n(self):
-    a = UOp.const(dtypes.int, 0)
+    a = UOp.const(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(dtypes.int, 2)
+    a = UOp.const(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(dtypes.int, 2)
+    a = UOp.const(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.float), lambda x: x.rtag())])
-    c1 = UOp.const(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.int, 1)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.weakfloat), lambda x: x.rtag())])
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(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(dtypes.float, 1.0)
+    c1 = UOp.const(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(dtypes.bool, False)
+    c1 = UOp.const(False)
     c2 = UOp(Ops.CAST, arg=dtypes.int, src=(c1,))
-    c3 = UOp.const(dtypes.float, 1.0)
+    c3 = UOp.const(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(dtypes.float, 0.0)
-    c2 = UOp.const(dtypes.bool, False)
+    c1 = UOp.const(0.0)
+    c2 = UOp.const(False)
     c3 = UOp(Ops.MAX, src=(c1, c1))
     c4 = UOp(Ops.MUL, src=(c1, c1))
-    c5 = UOp.const(dtypes.int, -1)
+    c5 = UOp.const(-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.arg in {1, -1} else None)
+       lambda x,c: x.rtag() if c.val in {1, -1} else None)
     ])
-    y1 = UOp.const(dtypes.int, 1)
-    y2 = UOp.const(dtypes.int, 2)
-    y3 = UOp.const(dtypes.int, -1)
+    y1 = UOp.const(1)
+    y2 = UOp.const(2)
+    y3 = UOp.const(-1)
     c1 = UOp(Ops.MUL, src=(y1, y2))
     c2 = UOp(Ops.MUL, src=(y2, y2))
     c3 = UOp(Ops.MUL, src=(y3, y2))
@@ -114,26 +114,27 @@ 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(dtypes.float, 1.0)
-    y2 = UOp.const(dtypes.float, 1.0)
+    y1 = UOp.const(1.0)
+    y2 = UOp.const(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):
-    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)
+    # 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)
     self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
 
   def test_dtype_set(self):
-    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)
+    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)
     self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), c2.rtag())
     self.assertEqual(matcher.rewrite(c3), None)
@@ -141,8 +142,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(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(2.0)
     c3 = UOp(Ops.ADD, src=(c1,c2))
     self.assertEqual(matcher.rewrite(c3), c3.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
@@ -158,8 +159,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(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(2.0)
     c3 = UOp(Ops.ADD, src=(c1,c2))
     c4 = UOp(Ops.ADD, src=(c3,c2))
     c5 = UOp(Ops.ADD, src=(c2,c3))
@@ -171,8 +172,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(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(2.0)
     c3 = UOp(Ops.ADD, src=(c1,c2))
     c4 = UOp(Ops.ADD, src=(c2,c3))
     self.assertEqual(matcher.rewrite(c3), c3.rtag())
@@ -180,9 +181,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(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
-    c3 = UOp.const(dtypes.float, 3.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(2.0)
+    c3 = UOp.const(3.0)
     c4 = UOp(Ops.EXP2, src=(c1,))
     c5 = UOp(Ops.ADD, src=(c1,c2))
     c6 = UOp(Ops.MULACC, src=(c1,c2,c3))
@@ -191,8 +192,8 @@ class TestPatternMatcher(unittest.TestCase):
     self.assertEqual(matcher.rewrite(c6), c6.rtag())
 
   def test_deep_src_permutations(self):
-    c1 = UOp.const(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(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).uop.const_like(10)
+    constv = Tensor.empty(2, 2).const_like(10).uop
     check_schedule(constv, 0)
 
   def test_const_schedule_contig(self):
-    constv = Tensor.empty(2, 2).uop.const_like(10).contiguous()
+    constv = Tensor.empty(2, 2).const_like(10).uop.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
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite, pm_lower_index_dtype
 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(dtypes.weakint, nmax),), arg=expr)
+def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(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(dtypes.bool, True), "svalid is None but valid is not True")
+      self.assertEqual(off.src[1].get_valid(), UOp.const(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(dtypes.weakint, 0)))
+    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(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(dtypes.weakint, 0)))
+    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(0)))
     try:
       self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
     except AssertionError:
@@ -488,10 +488,20 @@ 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(dtypes.weakint, 0))
+    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(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
@@ -499,8 +509,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(dtypes.weakint, 0)
-    true_gate = UOp.const(dtypes.bool, True)
+    idx = UOp.const(0)
+    true_gate = UOp.const(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)
@@ -516,50 +526,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(dtypes.weakint, 4), r)
+    load = get_gated_load_uop(r < UOp.const(4), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].arg, 4)
+    self.assertEqual(ranges[0].src[0].val, 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(dtypes.weakint, 4), r)
-    load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r)
+    load1 = get_gated_load_uop(r < UOp.const(4), r)
+    load2 = get_gated_load_uop(r < UOp.const(8), r)
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].arg, 8)
+    self.assertEqual(ranges[0].src[0].val, 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(dtypes.weakint, 300), r)
+    load = get_gated_load_uop(r < UOp.const(300), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].arg, 204)
+    self.assertEqual(ranges[0].src[0].val, 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(dtypes.weakint, 4), r)
+    load1 = get_gated_load_uop(r < UOp.const(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].arg, 204)
+    self.assertEqual(ranges[0].src[0].val, 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(dtypes.weakint, 4), r)
+    gated_load = get_gated_load_uop(r < UOp.const(4), r)
     red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
     ranges = self.get_ranges(red.sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].arg, 204)
+    self.assertEqual(ranges[0].src[0].val, 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(dtypes.weakint, 1), r)
+    load = get_gated_load_uop(r < UOp.const(1), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 0)
 
@@ -567,19 +577,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(dtypes.float, 1), Invalid)
+    x = (r < 4).where(UOp.const(1.0), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].arg, 4)
+    self.assertEqual(ranges[0].src[0].val, 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(dtypes.float, 1), Invalid)
+    x = (r < 4).where(UOp.const(1.0), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
     self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].arg, 4)
+    self.assertEqual(ranges[0].src[0].val, 4)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_symbolic_failures.py

@@ -1,5 +1,6 @@
 import unittest
 from tinygrad import Variable
+from tinygrad.uop.ops import UOp
 
 
 class TestFuzzFailure(unittest.TestCase):
@@ -8,7 +9,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 = v1.const_like(8), v2.const_like(0), v3.const_like(0)
+    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(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)
@@ -19,7 +20,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 = v1.const_like(6), v2.const_like(0), v3.const_like(0)
+    v1_val, v2_val, v3_val = UOp.const(6), UOp.const(0), UOp.const(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)
@@ -29,7 +30,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 = v1.const_like(0), v2.const_like(0), v3.const_like(0)
+    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(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)
@@ -39,7 +40,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 = v1.const_like(2), v2.const_like(0), v3.const_like(2)
+    v1_val, v2_val, v3_val = UOp.const(2), UOp.const(0), UOp.const(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)
@@ -49,7 +50,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 = v1.const_like(0), v2.const_like(0), v3.const_like(0)
+    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(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)
@@ -60,7 +61,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 = v1.const_like(8), v2.const_like(3), v3.const_like(2)
+    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(3), UOp.const(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)
@@ -70,7 +71,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 = v1.const_like(0), v2.const_like(2), v3.const_like(65)
+    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(2), UOp.const(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)
@@ -80,7 +81,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 = v1.const_like(0), v2.const_like(0), v3.const_like(0)
+    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(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)
@@ -90,7 +91,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 = v1.const_like(0), v2.const_like(1), v3.const_like(0)
+    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(1), UOp.const(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)
@@ -101,7 +102,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 = v1.const_like(9), v2.const_like(0), v3.const_like(0)
+    v1_val, v2_val, v3_val = UOp.const(9), UOp.const(0), UOp.const(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)
@@ -111,7 +112,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 = v1.const_like(0), v2.const_like(7), v3.const_like(0)
+    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(7), UOp.const(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(dtypes.float, 2.0)
+    u = UOp.const(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(dtypes.weakint, 1).stack(UOp.const(dtypes.weakint, 2)).shape, (2,))
+    self.assertEqual(UOp.const(1).stack(UOp.const(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(dtypes.int, 0)).load()
+    loaded_value = input_buf.index(UOp.const(0)).load()
     def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
       return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))
 
@@ -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(dtypes.float64, 12 * math.pi + 0.1)))
+    r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(12 * math.pi + 0.1).cast(dtypes.float64)))
     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(dtypes.float64, x)))
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(float(x)).cast(dtypes.float64)))
       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(dtypes.float64, 2.0)))
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(2.0).cast(dtypes.float64)))
       np.testing.assert_equal(mantissa, 0.5)
       np.testing.assert_equal(exponent, 2)
 
-    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 5.0)))
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(5.0).cast(dtypes.float64)))
     np.testing.assert_equal(mantissa, 0.625)
     np.testing.assert_equal(exponent, 3)
 
-    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 1000.0)))
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(1000.0).cast(dtypes.float64)))
     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(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)
+    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)
 
   def test_pow2if(self):
-    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)
+    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)
 
 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.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)),
+  (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)),
 ])
 
 def const_values(u:UOp):
-  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)
+  if u.op is Ops.CONST: return (u.val,)
+  if u.op is Ops.STACK: return tuple(x.val for x in u.src)
   raise AssertionError(f"expected const-like UOp, got {u.op}")
 
 class TestGraphRewriteConst(unittest.TestCase):
   def test_gep_const(self):
-    v1 = UOp.const(dtypes.int, (0,1,2))
+    v1 = UOp.const((0,1,2), dtypes.int)
     v2 = v1.index(1)
     ret = graph_rewrite(v2, sym)
     self.assertEqual(ret.dtype, dtypes.int)
-    self.assertEqual(ret.arg, 1)
+    self.assertEqual(ret.val, 1)
 
   def test_add_const(self):
-    v1 = UOp.const(dtypes.int, (0,1,2))
-    v2 = UOp.const(dtypes.int, (5,6,7))
+    v1 = UOp.const((0,1,2))
+    v2 = UOp.const((5,6,7))
     ret = graph_rewrite(v1+v2, sym)
     self.assertEqual(ret.op, Ops.STACK)
     self.assertEqual(const_values(ret), (5,7,9))
 
   def test_add_const_lose_v(self):
-    v1 = UOp.const(dtypes.int, (0,1,2))
-    v2 = UOp.const(dtypes.int, (2,1,0))
+    v1 = UOp.const((0,1,2))
+    v2 = UOp.const((2,1,0))
     ret = graph_rewrite(v1+v2, sym)
     self.assertEqual(ret.op, Ops.STACK)
     self.assertEqual(const_values(ret), (2,2,2))
@@ -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].arg, expected)
+    self.assertEqual(results[0].val, expected)
 
   @xfail_broken_const_wraparound
   def test_cast(self):
     t = self._test
-    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8), 0xD6)
-    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)
+    t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8), 0xD6)
+    t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)
 
   @xfail_broken_const_wraparound
   def test_mul(self):
     t = self._test
-    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD, 1147018174)
-    t(UOp.const(dtypes.int, 0xABCD17D6) * 10, -1241321892)
+    t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD, 1147018174)
+    t(UOp.const(0xABCD17D6, dtypes.int) * 10, -1241321892)
 
   @xfail_broken_const_wraparound
   def test_div(self):
     t = self._test
-    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD // 11, 104274379)
-    t(UOp.const(dtypes.int, 0xABCD17D6) * 10 // 11, -112847444)
+    t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD // 11, 104274379)
+    t(UOp.const(0xABCD17D6, dtypes.int) * 10 // 11, -112847444)
 
   @xfail_broken_const_wraparound
   def test_neg(self):
     t = self._test
-    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)
+    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)
 
   @xfail_broken_const_wraparound
   def test_neg_min_int(self):
     t = self._test
-    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)
+    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)
 
   @xfail_broken_const_wraparound
   def test_payne_hanek_reduction_bug(self):
     t = self._test
-    a = (UOp.const(dtypes.uint, 43748177600).cast(dtypes.uint) | 36).cast(dtypes.ulong)
-    b = 2536655455 * a + 4294967296 * UOp.const(dtypes.ulong, 25366554550)
+    a = (UOp.const(43748177600, dtypes.uint).cast(dtypes.uint) | 36).cast(dtypes.ulong)
+    b = 2536655455 * a + 4294967296 * UOp.const(25366554550, dtypes.ulong)
     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(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)
+    a1 = UOp.variable("a1", UOp.const(0), UOp.const(11), dtypes.int)
+    a2 = UOp.variable("a2", UOp.const(0), UOp.const(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(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(2.0)
     nout = graph_rewrite(c1+c2, simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.arg, 3.0)
+    self.assertEqual(nout.val, 3.0)
 
   def test_depth_2_late(self):
-    c1 = UOp.const(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
-    c3 = UOp.const(dtypes.float, 3.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(2.0)
+    c3 = UOp.const(3.0)
     nout = graph_rewrite(c1*c2*(c3+c3), simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.arg, 12.0)
+    self.assertEqual(nout.val, 12.0)
 
   def test_double(self):
-    c1 = UOp.const(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
-    c3 = UOp.const(dtypes.float, 3.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(2.0)
+    c3 = UOp.const(3.0)
     nout = graph_rewrite(c1+c2+c3, simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.arg, 6.0)
+    self.assertEqual(nout.val, 6.0)
 
   def test_triple(self):
-    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)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(2.0)
+    c3 = UOp.const(3.0)
+    c4 = UOp.const(4.0)
     nout = graph_rewrite(c1+c2+c3+c4, simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.arg, 10.0)
+    self.assertEqual(nout.val, 10.0)
 
   def test_diamond(self):
-    c1 = UOp.const(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
-    c3 = UOp.const(dtypes.float, 3.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(2.0)
+    c3 = UOp.const(3.0)
     nout = graph_rewrite((c1+c2)+(c1+c3), simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.arg, 7.0)
+    self.assertEqual(nout.val, 7.0)
 
   def test_magic_4(self):
-    c1 = UOp.const(dtypes.int, 4.0)
+    c1 = UOp.const(4)
     nout = graph_rewrite(c1, simple_pm)
     self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.arg, 3.0)
+    self.assertEqual(nout.val, 3.0)
 
   def test_depth_2_fold(self):
     v = UOp.variable("v", 0, 1, dtypes.float)
-    c1 = UOp.const(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
+    c1 = UOp.const(1.0)
+    c2 = UOp.const(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].arg, 3.0)
+    self.assertEqual(nout.src[1].val, 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], a.const_like(2))
+    self.assertIs(tst.src[1], UOp.const(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, a.const_like(2)+a, (4+d)+c+(2+a)+b]
+    outs = [2+a, 2+a+d+3+b+c+4, UOp.const(2)+a, (4+d)+c+(2+a)+b]
     for out in outs:
       sink = graph_rewrite(out, sym)
       print(sink.render())
@@ -191,59 +191,59 @@ class TestGraphRewrite(unittest.TestCase):
 
 class TestUOpGraph(unittest.TestCase):
   def test_add_constant_fold(self):
-    c1 = UOp.const(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
+    c1 = UOp.const(1.0, dtypes.float)
+    c2 = UOp.const(2.0, dtypes.float)
     out = c1+c2
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.arg, 3.0)
+    self.assertEqual(out.val, 3.0)
 
   def test_where_same_fold(self):
     v = UOp.variable('tmp', 0, 1)
-    c0 = UOp.const(dtypes.weakint, 0)
+    c0 = UOp.const(0)
     vc = v != c0
-    c1 = UOp.const(dtypes.float, 1.0)
+    c1 = UOp.const(1.0, dtypes.float)
     out = vc.where(c1, c1)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.arg, 1.0)
+    self.assertEqual(out.val, 1.0)
 
   def test_where_const_fold(self):
-    bf = UOp.const(dtypes.bool, False)
-    c1 = UOp.const(dtypes.float, 1.0)
-    c2 = UOp.const(dtypes.float, 2.0)
+    bf = UOp.const(False)
+    c1 = UOp.const(1.0, dtypes.float)
+    c2 = UOp.const(2.0, dtypes.float)
     out = bf.where(c1, c2)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.arg, 2.0)
+    self.assertEqual(out.val, 2.0)
 
   def test_const_cast(self):
-    bf = UOp.const(dtypes.bool, False)
+    bf = UOp.const(False)
     out = bf.cast(dtypes.int)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.arg, 0)
+    self.assertEqual(out.val, 0)
 
   def test_const_bitcast(self):
-    bf = UOp.const(dtypes.float, 1.0)
+    bf = UOp.const(1.0, dtypes.float)
     out = bf.bitcast(dtypes.uint32)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
     out = uops[-2]
     self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.arg, 0x3F800000)
+    self.assertEqual(out.val, 0x3F800000)
 
   @unittest.expectedFailure
   def test_const_shape_change_bitcast(self):
-    bf = UOp.const(dtypes.uint8, 0x3F)
+    bf = UOp.const(0x3F).cast(dtypes.uint8)
     out = bf.bitcast(dtypes.half)
     uops = to_uops_list([out])
     self.assertEqual(len(uops), 2)  # +1 for SINK
@@ -251,7 +251,7 @@ class TestUOpGraph(unittest.TestCase):
   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(None, 1.0).cast(dtypes.float))
+    lhs = UOp.stack(UOp.invalid(), UOp.const(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 +259,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(dtypes.int, 0)
+    idx = UOp.const(0)
     ld = d0.load(idx, dtype=dtypes.float)
     vec = UOp(Ops.STACK, dtypes.float, (ld,))
     x = vec.index(0)
@@ -273,7 +273,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(dtypes.int, 0)
+    idx = UOp.const(0)
     def _test_vec(geps, count=4):
       vec = UOp(Ops.STACK, dtypes.float, geps)
       out = d0.index(idx).store(vec)
@@ -310,7 +310,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_gep_vec_const_fold(self):
     for vec_size in [2, 4, 8]:
-      consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
+      consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
       vec = UOp(Ops.STACK, src=tuple(consts))
       with Context(SPEC=0):
         uops = to_uops_list([vec.index(i) for i in range(vec_size)])
@@ -320,7 +320,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(dtypes.int, 0)
+    idx = UOp.const(0)
     ld = d1.index(idx)
     alu = (ld<1).cast(dtypes.bool)
     out = d0.index(idx).store(alu)
@@ -330,7 +330,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(dtypes.int, 0)
+    idx = UOp.const(0, dtypes.int)
     ld = d1.index(idx)
     alu = ld.cast(dtypes.float).cast(dtypes.float)
     out = d0.index(idx).store(alu)
@@ -339,8 +339,8 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_depth_2_const_fold(self):
     v = UOp.variable("tmp", 0, 1, dtypes.int)
-    c2 = UOp.const(dtypes.int, 2)
-    c4 = UOp.const(dtypes.int, 4)
+    c2 = UOp.const(2, dtypes.int)
+    c4 = UOp.const(4, dtypes.int)
     vc = v+c2
     out = vc+c4
     uops = to_uops_list([out])
@@ -348,19 +348,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].arg, 6)
+    self.assertEqual(out.src[1].val, 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(dtypes.int, 0))
+      v = d0.index(UOp.const(0))
       uops = to_uops_list([v.bitcast(dt)])
       self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
 
   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(dtypes.int, 0)
+    assert uops[0] == UOp.const(0, dtypes.int)
     assert uops[-1].op == Ops.SINK
 
   def test_where_on_gated_load_fold(self):
@@ -372,7 +372,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].arg==5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
@@ -382,7 +382,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].arg==5
+      if u.op is Ops.LOAD: assert u.src[1].val==5
 
   def test_where_on_gated_load_with_cast(self):
     ridx0 = UOp.range(100, 0)
@@ -394,13 +394,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].arg == 5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 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(dtypes.float, 0)).cast(dtypes.half)
+    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.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 +410,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(dtypes.float, 0), ld).cast(dtypes.half)
+    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.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 +426,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].arg==5
+      if u.op is Ops.STORE: assert u.src[1].val==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(dtypes.weakint, 512), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
+    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)
-    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
+    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(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(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -449,19 +449,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(dtypes.weakint, 512), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
+    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)  # c4 is a load
-    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
+    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(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(dtypes.weakint, 60000)
+    comparison = (loaded_idx + c5) < UOp.const(60000)
     c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -470,9 +470,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(dtypes.int, 0)
+    idx = UOp.const(0)
     ld0 = glbl1.index(UOp.invalid())
-    ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True)))
+    ld1 = glbl2.index(idx.valid(UOp.const(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 +484,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(dtypes.bool, True)))
+    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(True)))
     uops = to_uops_list([glbl0.index(lidx).store(ld1+ld0)])
 
     # the gate and invalid value are deleted from ld1
@@ -492,10 +492,10 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_fold_gated_store(self):
     glbl = UOp.param(0, dtypes.int, (1,))
-    idx0 = UOp.const(dtypes.int, 0)
-    val = UOp.const(dtypes.int, 42)
+    idx0 = UOp.const(0)
+    val = UOp.const(42)
     st0 = glbl.index(UOp.invalid()).store(val)
-    st1 = glbl.index(idx0.valid(UOp.const(dtypes.bool, True))).store(val)
+    st1 = glbl.index(idx0.valid(UOp.const(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 +503,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(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))])
+    idx = UOp.const(0)
+    bad_gate = UOp.const(1)
+    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))])
 
   def test_after_end(self):
     r = UOp.range(10, 0)
@@ -513,7 +513,7 @@ class TestUOpGraph(unittest.TestCase):
     c = r + 1
     self.assertIn(r, c.ranges)
 
-    e = UOp.const(dtypes.int, 1).end(r)
+    e = UOp.const(1).end(r)
     self.assertNotIn(r, e.ranges)
 
     a = c.after(e)
@@ -563,7 +563,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(dtypes.float, 42.0)
+    c = UOp.const(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 +571,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.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
+    const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
     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(dtypes.float, 3.14)
+    c = UOp.const(3.14)
     r1 = UOp.range(3, 0)
     r2 = UOp.range(4, 1)
     bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU"))
@@ -586,14 +586,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.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
+    const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
     self.assertIn(3.14, const_vals)
 
 class TestUOpTags(unittest.TestCase):
   def test_inc_by_one(self):
-    g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)
+    g = UOp.const(1) + UOp.const(1)
     assert g.ssimplify() == 2
-    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_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_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 +682,36 @@ class TestUOpGetItem(unittest.TestCase):
 
 class TestUOpBroadcast(unittest.TestCase):
   def test_broadcast_row(self):
-    a = UOp.const(dtypes.float, 1, shape=(4, 8))
-    b = UOp.const(dtypes.float, 2, shape=(4, 1))
+    a = UOp.const(1.0).expand((4, 8))
+    b = UOp.const(2.0).expand((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(dtypes.float, 1, shape=(4, 8))
-    b = UOp.const(dtypes.float, 2, shape=(1, 8))
+    a = UOp.const(1.0).expand((4, 8))
+    b = UOp.const(2.0).expand((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(dtypes.float, 1, shape=(4, 8))
-    b = UOp.const(dtypes.float, 2, shape=(8,))
+    a = UOp.const(1.0).expand((4, 8))
+    b = UOp.const(2.0).expand((8,))
     c = a * b
     self.assertEqual(c.shape, (4, 8))
     self.assertEqual(c.op, Ops.MUL)
 
   def test_broadcast_scalar(self):
-    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    a = UOp.const(1.0).expand((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(dtypes.float, 1, shape=(1, 1, t))
-    b = UOp.const(dtypes.float, 2, shape=(1, 1, t))
+    a = UOp.const(1.0).expand((1, 1, t))
+    b = UOp.const(2.0).expand((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, dtypes
+from tinygrad import UOp
 
 class TestUOpRepr(unittest.TestCase):
   def test_simple_const(self):
-    a = UOp.const(dtypes.int, 42)
-    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.int, arg=42, src=())")
+    a = UOp.const(42)
+    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.weakint, arg=42, src=())")
   def test_different_consts(self):
-    a, b = UOp.const(dtypes.int, 42), UOp.const(dtypes.int, 3)
+    a, b = UOp.const(42), UOp.const(3)
     expected = (
-      "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=()),))"
+      "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=()),))"
     )
     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(dtypes.int, 42)
+    a = UOp.const(42)
     expected = (
-      "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
+      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.CONST, dtypes.weakint, 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(dtypes.int, 1)) + a
+    b = (a:=UOp.const(1)) + a
     expected = (
-      "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" +
+      "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" +
       "    x1,)),\n" +
       "  x0,))"
     )

test/null/test_uop_resolve.py

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

test/null/test_uop_symbolic.py

@@ -12,13 +12,12 @@ 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, 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)
+  if isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(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(dtypes.weakint, val)
+def uconst(val): return UOp.const(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)
 
@@ -104,16 +103,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(), a.const_like(5))
-    self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), a.const_like(1))
+    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))
 
   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(), 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*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*(-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)
@@ -883,7 +882,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, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)")
+    self.helper_test_variable((idx<4).where(idx//4, uconst(-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
@@ -942,7 +941,7 @@ class TestSymbolic(unittest.TestCase):
 
   def test_where_removal(self):
     cond = Variable("a", 0, 3) < 2
-    u1, u0 = cond.const_like(True), cond.const_like(False)
+    u1, u0 = UOp.const(True), UOp.const(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)")
@@ -1017,12 +1016,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(dtypes.int, 0)), expr)).sink())
+    uops = get_uops(UOp(Ops.STORE, src=(glbl.index(UOp.const(0, dtypes.int)), 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']<2).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))
+    self.assertEqual(rewritten_uop, (pvar['s']<UOp.const(2, dtypes.int)).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))
 
   def test_where_merge_branches(self):
     cond1 = Variable("s", 0, 10) < 6
@@ -1179,7 +1178,7 @@ class TestSymbolicVariables(unittest.TestCase):
     b = Variable("x", 1, 1).bind(1)
     s = b.simplify()
     self.assertEqual(s.op, Ops.CONST)
-    self.assertEqual(s.arg, 1)
+    self.assertEqual(s.val, 1)
 
 class TestSymInfer(unittest.TestCase):
   def test_sym_infer(self):
@@ -1203,9 +1202,13 @@ 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(dtypes.uint, 1)).bitcast(dtypes.int) + 2)
+    expr = ((a.bitcast(dtypes.uint) << UOp.const(1)).bitcast(dtypes.int) + 2)
     ret = sym_infer(expr, {a.expr: 2})
     assert isinstance(ret, int)
     assert ret == 6
@@ -1216,7 +1219,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(dtypes.float, 1.5).bitcast(dtypes.uint), {}) == 1069547520
+    assert sym_infer(UOp.const(1.5).cast(dtypes.float).bitcast(dtypes.uint), {}) == 1069547520
 
   def test_sym_infer_deeply_nested(self):
     # build an expression that exceeds Python's nested parentheses limit for eval
@@ -1337,7 +1340,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(UOp.const(dtypes.weakint, 0), UOp.invalid()),
+    self.assertIs(idx.simplify(), (ridx<5).where(uconst(0), UOp.invalid()),
                   "multiplying an index by 0 should preserve the invalid")
 
   def test_alu_moves_inside_invalid(self):
@@ -1362,24 +1365,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(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)))
+    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)))
 
 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(dtypes.weakint, 0))
+    index = buf.index(UOp.const(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(dtypes.int, 0)), sym).op, Ops.NOOP)
+    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(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(dtypes.int, 1)), sym).op, Ops.NOOP)
+    self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(1)), sym).op, Ops.NOOP)
     # Negative: store(idx, load(idx) + 1) should NOT fold
-    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(dtypes.int, 1)), sym).op, Ops.STORE)
+    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(1)), sym).op, Ops.STORE)
 
 class TestMoveWhereOnLoad(unittest.TestCase):
   def test_bool_index_preserves_dtype(self):
@@ -1390,7 +1393,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, idx.const_like(0))
+    expr = cond.where(idx, UOp.const(0))
     out = graph_rewrite(expr, pm_move_where_on_load)
     type_verify(out, spec_shared)  # Invalid matches any dtype
 
@@ -1475,7 +1478,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 = v1.const_like(8), v2.const_like(0), v3.const_like(0)
+    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(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(dtypes.int32, 42)
+    uop = UOp.const(42)
     self.assertEqual(uop.vmin, 42)
     self.assertEqual(uop.vmax, 42)
 
   def test_vmin_vmax_cmpne(self):
-    uop = UOp.const(dtypes.int32, 42)
+    uop = UOp.const(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(dtypes.float, 0.0)
-    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(dtypes.int, 0), dtype=dtypes.float)
+    x = UOp.const(0.0)
+    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(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(dtypes.weakint, (0, 4, Invalid, Invalid))
+    x = UOp.const((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, x.const_like(3))
+    floordiv = x.alu(Ops.FLOORDIV, UOp.const(3))
     self.assertEqual(floordiv.vmin, -3)
     self.assertEqual(floordiv.vmax, 2)
-    floormod = x.alu(Ops.FLOORMOD, x.const_like(3))
+    floormod = x.alu(Ops.FLOORMOD, UOp.const(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, x.const_like(-3))
+    floormod_neg = x.alu(Ops.FLOORMOD, UOp.const(-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(dtypes.int32, (42,))
+    uop = UOp.const((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(dtypes.int32, (10, 20, -5, 7))
+    uop = UOp.const((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(dtypes.int32, (7, 7, 7))
+    uop = UOp.const((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(dtypes.int32, (-10, -20, -5, -15))
+    uop = UOp.const((-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(dtypes.float32, (1.5, -3.2, 0.0))
+    uop = UOp.const((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(dtypes.bool, (True, False, False))
+    uop = UOp.const((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(dtypes.int, 0)
+    idx = UOp.const(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(dtypes.int32, 42)
+    uop = UOp.const(42)
     self.assertEqual(uop.const_factor(), 42)
 
   def test_const_factor_addition(self):
     # const_factor for an addition of constants
-    uop = UOp.const(dtypes.int32, 30) + UOp.const(dtypes.int32, 12)
+    uop = UOp.const(30) + UOp.const(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(dtypes.int32, 5) * UOp.const(dtypes.int32, 7)
+    uop = UOp.const(5) * UOp.const(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(dtypes.int32, 42)
+    uop = UOp.const(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(dtypes.int32, 42)
+    uop = UOp.const(42)
     result = uop.divides(5)
     self.assertIsNone(result)  # 42 is not divisible by 5
 

test/null/test_uops.py

@@ -5,23 +5,24 @@ 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, 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, 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.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(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)
+    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)
 
   def test_same_dtype_fast_path(self):
-    src = (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.weakint, 2))
+    src = (UOp.const(1), UOp.const(2))
     self.assertEqual(dtype_from_uop(Ops.ADD, src, None), dtypes.weakint)
 
   def test_where_promotion(self):
-    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)
+    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)
     idx = UOp.range(4, 0)
     self.assertEqual(idx.valid(idx < 4).dtype, dtypes.weakint)
 
@@ -37,33 +38,39 @@ 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(dtypes.int32, 3).dtype, dtypes.int32)
+    self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)
 
   def test_weak_dtype_rejected_by_program_spec(self):
     for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
-      with self.assertRaises(RuntimeError): type_verify(UOp.const(weak, value).sink(), spec_program)
-      type_verify(UOp.const(concrete, value).sink(), spec_program)
+      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)
 
   def test_invalid_dtype_and_consumers(self):
     invalid = UOp.invalid()
     self.assertIs(invalid.dtype, dtypes.bool)
-    self.assertIs(UOp.const(dtypes.float32, Invalid), invalid)
+    self.assertIs(UOp.const(Invalid, dtypes.float32), 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.arg is Invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
-    invalid, value = UOp.invalid(), UOp.const(dtypes.float32, 1)
+      if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
+    invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32)
     for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)),
-              UOp.const(dtypes.bool, True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
+              UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
               UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
     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(dtypes.float, 0))
+    self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float))
     type_verify(out.sink(), spec_program)
 
   def test_remove_invalid_stack_lanes(self):
-    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(dtypes.half, 1), UOp.invalid()))
+    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(1, dtypes.half), UOp.invalid()))
     out = graph_rewrite(stack, pm_remove_invalid)
-    self.assertEqual(out.src, (UOp.const(dtypes.half, 1), UOp.const(dtypes.half, 0)))
+    self.assertEqual(out.src, (UOp.const(1, dtypes.half), UOp.const(0, dtypes.half)))
     type_verify(out.sink(), spec_program)
 
 class TestLowerIndexDtype(unittest.TestCase):
@@ -72,7 +79,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(dtypes.weakint, 4)))
+    shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(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)
@@ -111,6 +118,20 @@ 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)
@@ -198,9 +219,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(dtypes.weakint, 1)
-    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(dtypes.weakint, 2)).valid(gate)))
-    val = UOp.const(dtypes.float, 42.0)
+    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)
     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)
@@ -215,10 +236,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(dtypes.weakint, 2)
-    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(dtypes.weakint, 1))))
+    idx = gidx0 * UOp.const(2)
+    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(1))))
     idx1 = UOp(Ops.INDEX, src=(gmem1, idx))
-    val = UOp.const(dtypes.float, 42.0)
+    val = UOp.const(42.0).cast(dtypes.float)
     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)
@@ -235,11 +256,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(dtypes.weakint, 2)
-    gate = gidx0<UOp.const(dtypes.weakint, 1)
+    idx = gidx0*UOp.const(2)
+    gate = gidx0<UOp.const(1)
     idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gate)))
     idx1 = UOp(Ops.INDEX, src=(gmem1, idx.valid(gate)))
-    val = UOp.const(dtypes.float, 42.0)
+    val = UOp.const(42.0).cast(dtypes.float)
     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]
@@ -257,7 +278,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(dt, 2)
+      c = UOp.const(2).cast(dt)
       l = g.index(c)
       a = UOp(Ops.CDIV, dt, (l, c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -270,7 +291,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(dt, 8)
+      c = UOp.const(8).cast(dt)
       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]
@@ -282,7 +303,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(dt, 2)
+      c = UOp.const(2).cast(dt)
       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]
@@ -294,7 +315,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(dtypes.uint, 3)
+    c = UOp.const(3).cast(dtypes.uint)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
     uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -314,7 +335,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, x.const_like(3)), vals=(val,)), cdiv(val, 3))
+      self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3).cast(x.dtype)), vals=(val,)), cdiv(val, 3))
 
   @Context(DISABLE_FAST_IDIV=0)
   def test_fast_idiv_remove_powers_of_two(self):
@@ -329,7 +350,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(dtypes.uint, 7)
+    c = UOp.const(7).cast(dtypes.uint)
     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)
@@ -340,7 +361,7 @@ class TestFastIdiv(unittest.TestCase):
 
   def test_disable_fast_idiv(self):
     g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(dtypes.uint, 3)
+    c = UOp.const(3).cast(dtypes.uint)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
     with Context(DISABLE_FAST_IDIV=1):
@@ -352,8 +373,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(dtypes.float, 2.0)
-    b = UOp.const(dtypes.float, 3.0)
+    a = UOp.const(2.0)
+    b = UOp.const(3.0)
 
     add = UOp(Ops.ADD, src=(a, b))
     mul = UOp(Ops.MUL, src=(a, b))
@@ -368,8 +389,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(dtypes.int, 8),), arg='gidx0')
-    self.assertEqual(UOp.const(dtypes.int, 17).const_factor(), 17)
+    gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(8),), arg='gidx0')
+    self.assertEqual(UOp.const(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)
@@ -388,26 +409,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(dtypes.float, 0.0)
-    neg_zero = UOp.const(dtypes.float, -0.0)
+    pos_zero = UOp.const(0.0)
+    neg_zero = UOp.const(-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(dtypes.float, float('nan'))
-    nan2 = UOp.const(dtypes.float, float('nan'))
+    nan1 = UOp.const(float('nan'))
+    nan2 = UOp.const(float('nan'))
     self.assertIs(nan1, nan2)
 
 class TestUOpStr(unittest.TestCase):
   def test_uop_str(self):
-    a = UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 3.0)
+    a = UOp.const(2.0) + UOp.const(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(dtypes.int, x) for x in range(4)))
+    vec = UOp(Ops.STACK, src=tuple(UOp.const(x) for x in range(4)))
     assert str(eval(str(vec))) == str(vec)
 
   def test_reduceop_arg(self):
@@ -442,16 +463,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, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
+    u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
     self.assertEqual(u.render(simplify=False), "{0,0,0}")
   def test_render_vectorize_different(self):
-    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
+    u = UOp(Ops.STACK, src=tuple(UOp.const(i) for i in range(3)))
     self.assertEqual(u.render(simplify=False), "{0,1,2}")
   def test_render_vectorize_same_simplified(self):
-    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
+    u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
     self.assertEqual(u.render(), "{0,0,0}")
   def test_render_vectorize_different_simplified(self):
-    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
+    u = UOp(Ops.STACK, src=tuple(UOp.const(i) for i in range(3)))
     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(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)
+      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)
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(dtypes.int, 16)).load(dtype=dtypes.int)])  # off by one
+        to_uops_list([buf.index(UOp.const(16)).load(dtype=dtypes.int)])  # off by one
       with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(dtypes.int, 42)).load(dtype=dtypes.int)])  # way out
+        to_uops_list([buf.index(UOp.const(42)).load(dtype=dtypes.int)])  # way out
 
   def test_variable_index(self):
     with Context(CHECK_OOB=1, SPEC=2):
@@ -138,11 +138,19 @@ 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(dtypes.weakint, 16),), arg="gidx0")
+      gidx = UOp(Ops.SPECIAL, src=(UOp.const(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
@@ -156,12 +164,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(dtypes.int, 416),), arg="gidx0")
-      lidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 10),), arg="lidx0")
+      gidx = UOp(Ops.SPECIAL, src=(UOp.const(416),), arg="gidx0")
+      lidx = UOp(Ops.SPECIAL, src=(UOp.const(10),), arg="lidx0")
 
       gate = (gidx<400) & (lidx<8)
 
-      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(dtypes.uint, 1))
+      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(1))
 
       barrier = UOp(Ops.BARRIER, src=(local_store,))
       if_barrier = UOp(Ops.IF, src=(gate, barrier))
@@ -179,7 +187,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, ridx<16&mask))))
+      ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
       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.arg <= 3
-      return x.replace(arg=x.arg+1)
+      assert x.val <= 3
+      return x.replace(arg=x.val+1)
     err_pm = PatternMatcher([(UPat.cvar("x"), count_3),])
-    a = UOp.const(dtypes.int, 1)
+    a = UOp.const(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(dtypes.int, 3)
-    b = UOp.const(dtypes.int, 4)
+    a = UOp.const(3)
+    b = UOp.const(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(dtypes.int, 3)
+      inner = UOp.const(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(a.dtype, 0)
-      y = UOp.const(dtypes.float, math.pi)
+      z = UOp.const(0, a.dtype)
+      y = UOp.const(math.pi, dtypes.float)
       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(dtypes.float, 1.0, shape=(3,4))  # creates CONST->EXPAND chain
+    c = UOp.const(1.0).expand((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(dtypes.int, 1)
+    c = UOp.const(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(dtypes.int, 2))
+    const_stack = c.stack(UOp.const(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(dtypes.int, 1))
+      test(c:=UOp.const(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(dtypes.int, 1))
+      test(c:=UOp.const(1))
       test(c+1)
     ls = viz.list_items()
     self.assertEqual(len(ls), 2)
@@ -425,21 +425,22 @@ 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(dtypes.int, 1)
+      c = UOp.const(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.const_like(2)), id(c+2)])
+    self.assertEqual(list(graph), [id(c), id((c+2).src[1]), id(c+2)])
     self.assertTrue(graph[id(c)]["exclude"])
-    self.assertTrue(graph[id(c.const_like(2))]["exclude"])
+    self.assertTrue(graph[id((c+2).src[1])]["exclude"])
     self.assertFalse(graph[id(c+2)]["exclude"])
 
   def test_recurse(self):

test/opt/test_tensor_cores.py

@@ -81,6 +81,13 @@ 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(dtypes.float, 2.0)))
+    s.assign(Tensor(UOp.const(2.0).cast(dtypes.float)))
     np.testing.assert_equal(s.numpy(), [2, 2, 2, 2])
 
   def test_nested_after_contiguous_store(self):

test/unit/test_dtype_weak.py

@@ -3,7 +3,8 @@ 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
+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.spec import spec_shared, type_verify
 from tinygrad.engine.jit import JitError
 
@@ -11,18 +12,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(dtypes.weakfloat, 1.0).rand_like()
-    with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).randn_like()
+    with self.assertRaises(ValueError): Tensor.const(1.0).rand_like()
+    with self.assertRaises(ValueError): Tensor.const(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(weak, value).expand(3)
+      t = Tensor.const(value, weak).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(dtypes.weakfloat, 1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float32)
+    self.assertEqual((Tensor.const(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(weak, value)
+      t = Tensor.const(value, weak)
       self.assertEqual(t.dtype, weak)
       self.assertEqual(t.data().itemsize, strong.itemsize)
       self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize)
@@ -32,7 +33,7 @@ class TestWeakPromotion(unittest.TestCase):
       self.assertEqual(t.contiguous().dtype, weak)
 
   def test_assign_into_weak_commits(self):
-    t = Tensor.const(dtypes.weakfloat, 0.5)
+    t = Tensor.const(0.5)
     t.assign(Tensor(1.0, dtype=dtypes.default_float))
     self.assertEqual((t.dtype, t.item()), (dtypes.default_float, 1.0))
 
@@ -43,10 +44,59 @@ 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_uop_scalar_const_unchanged(self):
-    for dtype, value in ((dtypes.weakint, 1), (dtypes.int32, 1), (dtypes.float32, 0.5)):
+  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)):
       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].dtype), (dtype, dtype))
+      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)))
 
   @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):
@@ -79,17 +129,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(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(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((v & 3).dtype, dtypes.weakint)
-    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))
+    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))
     for op in (Ops.SHL, Ops.SHR):
       with self.assertRaises(RuntimeError):
-        UOp.const(dtypes.float32, 1).alu(op, UOp.const(dtypes.int32, 1))
+        UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32))
     # float bitwise builds, the spec rejects it
     with Context(SPEC=1):
-      f32, wf = UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.weakfloat, 1.0)
+      f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0)
       for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, dtypes.float32, (f32, f32)), UOp(Ops.AND, dtypes.int32, (wf, wf))):
         with self.assertRaises(RuntimeError): type_verify([bad], spec_shared)
 
@@ -128,7 +178,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(dtypes.weakfloat, 0.5).reshape(1)
+    w05 = Tensor.const(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
@@ -149,7 +199,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(weak, value), Tensor.const(weak, value))
+        return Tensor([True], device="CPU").where(Tensor.const(value, weak), Tensor.const(value, weak))
       self.assertEqual(weak_val().dtype, weak)
       self.assertEqual(weak_val().to("CPU").dtype, weak)
       self.assertEqual(weak_val().data().format, strong.fmt)
@@ -165,7 +215,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(dtypes.weakfloat, 0.5), devful):
+    for t in (Tensor.const(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)
@@ -176,7 +226,7 @@ class TestWeakMaterializationEntries(unittest.TestCase):
 
   def test_empty_reads_commit(self):
     for weak, strong in ((dtypes.weakfloat, dtypes.default_float),):
-      empty = Tensor.const(weak, 0).reshape(1).shrink(((0, 0),))
+      empty = Tensor.const(0, weak).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 dtypes, Invalid
+from tinygrad.dtype import 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(UOp.const(dtypes.float32, Invalid, shape=(2,))))
+      after = state.uop.after(state.uop.shrink(((0, 2),)).store(Invalid))
       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(dtypes.weakint, 0) < UOp.const(dtypes.weakint, 1)
-    idx = UOp(Ops.STACK, src=tuple(UOp.const(dtypes.weakint, x) for x in range(3)))
+    cond = UOp.const(0) < UOp.const(1)
+    idx = UOp(Ops.STACK, src=tuple(UOp.const(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.assertIs(out.src[2].arg, Invalid)
+    self.assertTrue(out.src[2].is_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(dtypes.float, 2.0))).item()
+      f(Tensor(UOp.const(2.0).cast(dtypes.float))).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(dtypes.float, 2.0) + UOp.const(dtypes.float, 1.0))).item()
+      f(Tensor(UOp.const(2.0).cast(dtypes.float) + UOp.const(1.0).cast(dtypes.float))).item()
 
   def test_jit_init_empty_alt(self):
     @TinyJit

test/unit/test_llm_server.py

@@ -6,6 +6,8 @@ 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):
@@ -14,7 +16,7 @@ class TestTransformerGenerate(unittest.TestCase):
 
     captured_inputs = []
     def mock_call(self, tokens, start_pos, temperature):
-      captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val))
+      captured_inputs.append((tokens.shape, start_pos))
       return Tensor([[42]])
 
     with patch.object(Transformer, '__call__', mock_call):
@@ -31,9 +33,7 @@ class TestTransformerGenerate(unittest.TestCase):
       next(gen)
 
     # should process tokens[6:] = [42, 10, 11, 12] since first 6 have cached k/v
-    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)
+    self.assertEqual(captured_inputs, [((1, V_TOKS.bind(4)), V_START_POS.bind(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 if isinstance(start_pos, int) else start_pos.val))
+      captured_inputs.append((tokens.shape, start_pos))
       return Tensor([[42]])
 
     with patch.object(Transformer, '__call__', mock_call):
@@ -55,9 +55,7 @@ class TestTransformerGenerate(unittest.TestCase):
       next(gen)
 
     # should process all 3 tokens from start
-    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)
+    self.assertEqual(captured_inputs, [((1, V_TOKS.bind(3)), V_START_POS.bind(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(dtypes.weakint, offset))
+      buf.src = (src, UOp.const(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(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])
+    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])
 
   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(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)
+    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)
 
   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].arg * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize
+    byte_offset = buf.src[1].val * 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(None, offset)), src.numel())
+  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(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
+from tinygrad.uop.ops import AxisType, pm_commit_weak, pm_cast_weak
 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, Invalid
+from tinygrad.dtype import dtypes, AddrSpace
 
 # 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(r.dtype, tuple(range(r.vmax+1))) \
+   lambda ctx, r: UOp.const(tuple(range(r.vmax+1)), r.dtype) \
     .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.arg is Invalid for x in b.src): return None
+  if not all(x.shape == b.shape or x.base.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(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)))
+    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)))
   return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
 
 def do_stack_wmma(u:UOp):
@@ -330,10 +330,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads")
 
   # devectorize
-  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
-
-  # simplify indexing
-  sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
+  sink = graph_rewrite(sink, symbolic_simple+devectorizer2+indexing_simplify, ctx=ren, name="devectorize2")
 
   # some coalescing misses without this
   sink = graph_rewrite(sink, sym, name="early symbolic")
@@ -362,7 +359,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, ctx=(set(), ren), name="decomp dtypes")
+  sink = graph_rewrite(sink, pm_dtype_decomps+pm_commit_weak, 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 +368,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_decomp+extra_matcher+pm_split_ends
+  pm_final_rewrite = pm_commit_weak+pm_cast_weak+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(dt, 0)
+  zero = UOp.const(0, dt)
   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(dtypes.uint, 0), z), (z, z)
+      q, r = (z:=UOp.const(0, dtypes.uint), z), (z, z)
       for i in range(63, -1, -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]
+        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]
         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(x.tag[1], truncate[x.tag[1]]((x.arg >> 32) if x.tag[0] == 1 else (x.arg & 0xFFFFFFFF))))
+   UOp.const(truncate[x.tag[1]]((x.val >> 32) if x.tag[0] == 1 else (x.val & 0xFFFFFFFF)), x.tag[1]))
 ])
 
 # 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.arg-1) if c.arg 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.val-1) if c.val 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.arg, 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.val, 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.arg, 0)) else None)]
+      lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 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.arg, 0)) else None)]
+        if (v:=powers_of_two.get(c.val, 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.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
+        lambda ctx, x, d: fast_idiv(ctx, x, d.val) 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.arg+1==c2.arg-1 else None),  # (c-1)<x & x<(c+1) -> x==c
+        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
     ]
   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.arg), 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.val), 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(None, 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(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().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 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 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(dtypes.uint32, 0), i) for i in range(4)]
+  a = [_take(UOp.const(0, dtypes.uint32), 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.arg is not True:
+  if valid.op is not Ops.CONST or valid.val is not True:
     return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid), dtype=dtypes.float)
   else:
     return buf.index(cidx.src[1], cidx.src[0], dtype=dtypes.float)
@@ -111,10 +111,10 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
       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].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
+      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
       else: root_src, arg = idx, 0
       memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)
 
@@ -127,7 +127,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.fp8s) and not is_image_shape(buf._shape):
+    elif buf.dtype not in (dtypes.float, dtypes.half, dtypes.int, dtypes.uint, *dtypes.fp8s) and not is_image_shape(buf._shape):
       pass
     elif buf.addrspace == AddrSpace.REG:
       pass
@@ -141,12 +141,12 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
     grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.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(None, full_grp[0])
+        offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(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(None, len(grp)))) if len(grp) > 1 else buf.index(offset)
+        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset)
         if op == Ops.STORE:
           datas = []
           for i,g in enumerate(grp):

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.arg is Invalid else
+  return l.replace(src=(l.src[0], l.vconst_like(0) if a.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(dtypes.int32, offset)
+        self.spills[v] = UOp.const(offset, dtypes.int32)
         self.stack_size = offset + sz
       r = alloc(cons if cons is not None else v.cons, i)
       self.insert_before.setdefault(i, []).append((v, r))
@@ -84,7 +84,7 @@ class LinearScanRegallocContext:
 
       # allocate stack array
       if u.op is Ops.BUFFER:
-        self.locals[u] = UOp.const(dtypes.int32, self.stack_size)
+        self.locals[u] = UOp.const(self.stack_size, dtypes.int32)
         self.stack_size += u.max_numel() * u.dtype.itemsize
 
       # loop prologue, avoid loading inside the loop

tinygrad/codegen/opt/heuristic.py

@@ -26,22 +26,18 @@ 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)):
-    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]))
+    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]))
       return tk
 
   # make a copy so it does not mutate the input
@@ -130,8 +126,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].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
+            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
         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(b.dtype, Invalid))
+          replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(Invalid, b.dtype))
       self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
     elif opt.op is OptOps.SWAP:
       try:
@@ -245,6 +245,8 @@ 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")})
 
@@ -302,7 +304,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(tc.dtype_out, (0.0,)*tc.elements_per_thread[2]),
+            tc_uop = UOp.wmma(srcs[0], srcs[1], UOp.const((0.0,)*tc.elements_per_thread[2], tc.dtype_out),
                               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].arg < c.arg)}
+  ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].val < c.val)}
   # 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}
 
@@ -62,7 +62,7 @@ pm_simplify_ranges = PatternMatcher([
 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.arg) is not None: ctx[r] = c
+    and r.src[0].op is Ops.CONST and r.src[0].divides(c.val) 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,6 +17,7 @@ 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)
@@ -74,8 +75,7 @@ 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: tuple[ConstType, ...]|ConstType):
-    if isinstance(val, tuple): return tuple(map(self.const, val))
+  def const(self, val: ConstType):
     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,6 +165,8 @@ 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].arg * arg_uops[1].dtype.itemsize
+    offset = ast.src[1].val * 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].arg*bufs[1].dtype.itemsize)
+  bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].val*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.arg for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
+  shape, pos_var = tuple(s.val for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
   with track_stats(ctx, call, bufs[0].device, bufs, ctx.var_vals):
     bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
   return None

tinygrad/function.py

@@ -1,7 +1,6 @@
 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
@@ -22,8 +21,7 @@ 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.op is Ops.CONST and u.src[1].base.arg is Invalid
-          and not u.src[0].buf_uop.is_realized}
+          if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}
 
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):

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(t.dtype, tuple(2**(i*b) for i in range(8//b))), 0xff >> (8 - b)
+    shift_tensor, bitmask = Tensor.const(tuple(2**(i*b) for i in range(8//b)), t.dtype), 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(dtypes.uint32, (0, 7, 14, 21))).bitwise_and(0x7F).reshape((-1, 32)).cast(dtypes.int32)
+      sign_idx = scale_words.unsqueeze(-1).rshift(Tensor.const((0, 7, 14, 21), dtypes.uint32)).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(dtypes.uint16, (0, 2, 4, 6, 8, 10, 12, 14))
+      scale_shifts = Tensor.const((0, 2, 4, 6, 8, 10, 12, 14), dtypes.uint16)
       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,9 +432,10 @@ 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:
-      sp, nt = v_start_pos.bind(start_pos), v_toks.bind(min(chunk_size, len(tokens) - start_pos))
+      n_toks = min(chunk_size, len(tokens) - start_pos)
+      sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks)
       out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
-      start_pos += nt.val
+      start_pos += n_toks
       # chunked prefill: keep processing until all prompt tokens are consumed
       if start_pos < len(tokens): continue
       tokens.append(int(out.item()))

tinygrad/llm/serve.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import json, pathlib, re, time, typing, uuid
+import json, pathlib, queue, re, threading, time, traceback, typing, uuid
 from typing import TYPE_CHECKING
 from tinygrad.helpers import DEBUG, colored, stderr_log
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
@@ -60,6 +60,11 @@ class StreamRouter:
     if emit: yield "content", emit
     if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
 
+class Request:
+  def __init__(self, ids:list[int], model_name:str, include_usage:bool, max_tokens:int|None, temperature:float):
+    self.ids, self.model_name, self.include_usage, self.max_tokens, self.temperature = ids, model_name, include_usage, max_tokens, temperature
+    self.out: queue.Queue[dict|None] = queue.Queue()
+
 class Handler(HTTPRequestHandler):
   server: LLMServer
   def log_request(self, code='-', size='-'): pass
@@ -67,46 +72,10 @@ class Handler(HTTPRequestHandler):
     if self.path == "/v1/models": self.send_data(json.dumps({"object":"list","data":[{"id":self.server.model_name,"object":"model"}]}).encode())
     else: self.send_data((pathlib.Path(__file__).parent / "chat.html").read_bytes(), content_type="text/html")
   def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0):
-    model, tok = self.server.model, self.server.tok
-    prompt_tokens = len(ids)
-    cache_start_pos = model.get_start_pos(ids)
-    stderr_log(f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
-    tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
-    def chunk(d:dict): return {"choices": [{"index":0, "delta":d, "finish_reason":None}], **tmpl}
-    yield chunk({"role":"assistant", "content":""})
-    out: list[int] = []
-    finish_reason = "stop"
-    st = time.perf_counter()
-    dec = tok.stream_decoder()
-    router = StreamRouter()
-    for next_id in model.generate(ids, temperature=temperature):
-      if len(out) == 0: stderr_log(f"prefill:{(prompt_tokens-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
-      if tok.is_end(next_id): break
-      out.append(next_id)
-      for field, delta in router.route(dec(next_id)): yield chunk({field:delta})
-      if max_tokens is not None and len(out) >= max_tokens:
-        finish_reason = "length"
-        break
-    for field, delta in router.route(dec(), final=True): yield chunk({field:delta})
-    tool_calls: list[dict] = []
-    for m in re.finditer(r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", router.buf, re.DOTALL):
-      if (parsed := parse_tool_call(m.group(1))) is None:
-        stderr_log(f"failed to parse tool call: {m.group(1)[:200]}")
-        yield chunk({"content":m.group(0)})  # don't silently drop output the client can't use
-      else:
-        name, args = parsed
-        tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function",
-                           "function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}})
-    if tool_calls:
-      yield chunk({"tool_calls":tool_calls})
-      if finish_reason == "stop": finish_reason = "tool_calls"
-    yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
-    if include_usage:
-      yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out),
-                                      "total_tokens": prompt_tokens + len(out)}, **tmpl}
-    et = time.perf_counter()
-    stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s  {colored('--', 'BLACK')}  "
-               f"out:{len(out):5d}  {colored('--', 'BLACK')}  total:{et-st:6.2f}s\n")
+    # hand off to the scheduler thread (the only thread that touches the model), drain our private stream
+    req = Request(ids, model_name, include_usage, max_tokens, temperature)
+    self.server.waiting.put(req)
+    while (c := req.out.get()) is not None: yield c
 
   def do_POST(self):
     request_st = time.perf_counter()
@@ -152,4 +121,58 @@ class Handler(HTTPRequestHandler):
 class LLMServer(TCPServerWithReuse):
   def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer, template:typing.Any):
     self.model, self.model_name, self.tok, self.template = model, model_name, tok, template
+    self.waiting: queue.Queue[Request] = queue.Queue()
+    threading.Thread(target=self.scheduler, daemon=True, name="LLMScheduler").start()
     super().__init__(server_address, Handler)
+
+  def scheduler(self):
+    # single thread owns the model; handler threads enqueue requests and drain req.out
+    while True:
+      req = self.waiting.get()                                   # admit FCFS
+      try:
+        for c in self._generate(req): req.out.put(c)             # TODO: batch decodes across requests (batched_step)
+      except Exception: traceback.print_exc()
+      req.out.put(None)                                          # retire: None ends the client's stream
+
+  def _generate(self, req:Request):
+    ids, model_name, include_usage, max_tokens, temperature = req.ids, req.model_name, req.include_usage, req.max_tokens, req.temperature
+    model, tok = self.model, self.tok
+    prompt_tokens = len(ids)
+    cache_start_pos = model.get_start_pos(ids)
+    stderr_log(f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
+    tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
+    def chunk(d:dict): return {"choices": [{"index":0, "delta":d, "finish_reason":None}], **tmpl}
+    yield chunk({"role":"assistant", "content":""})
+    out: list[int] = []
+    finish_reason = "stop"
+    st = time.perf_counter()
+    dec = tok.stream_decoder()
+    router = StreamRouter()
+    for next_id in model.generate(ids, temperature=temperature):
+      if len(out) == 0: stderr_log(f"prefill:{(prompt_tokens-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
+      if tok.is_end(next_id): break
+      out.append(next_id)
+      for field, delta in router.route(dec(next_id)): yield chunk({field:delta})
+      if max_tokens is not None and len(out) >= max_tokens:
+        finish_reason = "length"
+        break
+    for field, delta in router.route(dec(), final=True): yield chunk({field:delta})
+    tool_calls: list[dict] = []
+    for m in re.finditer(r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", router.buf, re.DOTALL):
+      if (parsed := parse_tool_call(m.group(1))) is None:
+        stderr_log(f"failed to parse tool call: {m.group(1)[:200]}")
+        yield chunk({"content":m.group(0)})  # don't silently drop output the client can't use
+      else:
+        name, args = parsed
+        tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function",
+                           "function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}})
+    if tool_calls:
+      yield chunk({"tool_calls":tool_calls})
+      if finish_reason == "stop": finish_reason = "tool_calls"
+    yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
+    if include_usage:
+      yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out),
+                                      "total_tokens": prompt_tokens + len(out)}, **tmpl}
+    et = time.perf_counter()
+    stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s  {colored('--', 'BLACK')}  "
+               f"out:{len(out):5d}  {colored('--', 'BLACK')}  total:{et-st:6.2f}s\n")

tinygrad/mixin/creation.py

@@ -9,7 +9,7 @@ if TYPE_CHECKING:
 
 class CreationMixin(DTypeMixin, MovementMixin):
   @staticmethod
-  def const(dtype, b): raise NotImplementedError
+  def const(b, dtype=None): 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(dt, fill_value)
+    val = cls.const(fill_value, dt)
     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, Invalid
+from tinygrad.dtype import DType, DTypeLike, dtypes, to_dtype
 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.arg 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.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
+from tinygrad.dtype import dtypes, ConstType, PyConst, least_upper_dtype, least_upper_float, weak_dtype
 from tinygrad.helpers import argfix, polyN
 from tinygrad.mixin.creation import CreationMixin
 
@@ -21,7 +21,12 @@ 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)
-    return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype)
+    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)
 
   def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self:
     lhs, rhs = self._broadcasted(x, reverse)
@@ -390,6 +395,8 @@ 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(dtypes.bool, True)).where(vb, self)
+    return (type(self).uprod(*per_dim) if per_dim else type(self).const(True)).where(vb, self)
 
   @classmethod
   def arange(cls, start, stop=None, step=1, dtype:DTypeLike|None=None) -> Self:
@@ -703,6 +703,28 @@ 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`.
@@ -1389,7 +1411,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     ret = self
     for dim in range(dims):
       ret = ret.transpose(0, dim)
-      ret = sum(type(self).const(ret.dtype, tuple(float(m[k]) for m in mat)).reshape((len(mat),)+(1,)*(ret.ndim-1)) * ret[k]
+      ret = sum(type(self).const(tuple(float(m[k]) for m in mat), ret.dtype).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)
@@ -1876,7 +1898,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(dtype, tuple(l))
+      return type(self).const(tuple(l), dtype)
     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])
 
@@ -1896,7 +1918,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(dtypes.uint8, v).expand(l) for l, v in mb if l > 0)).unsqueeze(0)
+    pad_mask = type(self).cat(*(type(self).const(v, dtypes.uint8).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.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"),
+  (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"),
   # 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.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))})"),
+  (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))})"),
   # default const render
-  (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
+  (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.val)),
 
   # 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.arg}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.arg]}")
+      return self[buf]+(f"[{idx.val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.val]}")
     return f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})"
 
   def render_buffer(self, x:UOp):
@@ -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.arg)))[0] >> 16)}u"),
+      lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.val)))[0] >> 16)}u"),
     # load/store image (OpenCL)
     (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"),
     (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))),
@@ -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.bfloat16: "bfloat"}
+  type_map = {dtypes.uint32: "uint", 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.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"}
+  type_map = {dtypes.uint32: "uint", 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 = ["#define INFINITY (__int_as_float(0x7f800000))", "#define NAN (__int_as_float(0x7fffffff))",
+    prefix = ["typedef unsigned int uint;", "#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>")
@@ -494,10 +494,10 @@ class HIPRenderer(CStyleLanguage):
         (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]},"
           f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[0][2] == 128 else None),
         (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"),
-        (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.arg) 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.val) 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.arg}f, {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.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 +522,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(dtypes.float, x.arg))),
+    (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.val, dtypes.float))),
   ])
 
   def asm(self, prg:UOp, lin:UOp) -> bytes:
@@ -538,7 +538,7 @@ class HIPRenderer(CStyleLanguage):
     prefix, ockl = [], []
     type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" }
     used_dtypes = uops_to_dtypes(uops)
-    if any(u.op is Ops.CONST and not math.isfinite(u.arg) for u in uops):
+    if any(u.op is Ops.CONST and not math.isfinite(u.val) for u in uops):
       prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"]
     if any(u.op is Ops.SPECIAL for u in uops):
       prefix.append("typedef long unsigned int size_t;")

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(dtypes.int32, 0), dtype=dtypes.uint64)
+  local_idx = local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64)
   # the selected address is a 64bit value, the AFTER orders the load after the scratch store and carries the element dtype for the encoder
   sel = gate.where(addr.replace(dtype=dtypes.uint64), local_idx)
   ptr = UOp(Ops.AFTER, addr.dtype, (sel, (local_idx if x.max_numel() == 1 else local).store(alt)))
@@ -174,7 +174,7 @@ def gated_load(ctx, addr:UOp, alt:UOp, gate:UOp, x:UOp):
 
 def gated_store(addr:UOp, gate:UOp, val:UOp):
   local = scratch_buffer(addr.src[0].dtype.scalar(), val.max_numel(), -1)
-  sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.const(dtypes.int32, 0), dtype=dtypes.uint64))
+  sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64))
   return UOp(Ops.AFTER, addr.dtype, (sel,)).store(val)
 
 # legalize the new style graph for isel. NOTE: this runs after the spec is verified, some of these rewrites violate it
@@ -195,7 +195,7 @@ pre_isel_matcher = PatternMatcher([
   # if gate in scalar int cmove is not a comparison need to add one to set the flag
   # NOTE: the 0 is int so the bool gate zero-extends and compares as int (a byte compare renders different kernels)
   (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")),
-   lambda m,a,b: m.ne(UOp.const(dtypes.int, 0)).where(a,b) if m.op not in GroupOp.Comparison else None),
+   lambda m,a,b: m.ne(UOp.const(0, dtypes.int)).where(a,b) if m.op not in GroupOp.Comparison else None),
 ])
 
 # ***** X86 registers *****
@@ -221,15 +221,15 @@ reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"},
 
 # ***** X86 instruction selection *****
 def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX else s
-def lane(x:UOp, i:int) -> int: return s.src[1].arg if (s:=x.src[i]).op is Ops.INDEX else 0
+def lane(x:UOp, i:int) -> int: return s.src[1].val if (s:=x.src[i]).op is Ops.INDEX else 0
 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt]
 def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, arg=X86Ops.DEFINE, tag=None if reg is None else (reg,))
-def imm(dt:DType, v:int) -> UOp: return UOp.const(dt, truncate[dt](v)).rtag()
+def imm(dt:DType, v:int) -> UOp: return UOp.const(truncate[dt](v), dt).rtag()
 def to_imm(c:UOp) -> UOp|None:
   if c.op is not Ops.CONST: return None
-  if c.dtype is dtypes.int64: return imm(dtypes.int32, c.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)
+  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)
   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].arg * scale), sz)
-  if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.arg * 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].val * scale), sz)
+  if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.val * 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.arg),) + 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.val),) + x.src[1:])),
   (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None),
   # really all a backedge END is is an IF with a tag referencing the RANGE start label
   (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"),
@@ -367,10 +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.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.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.val),)) if not x.tag else None),
+  (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.val),)) if not x.tag else None),
   (UPat.cvar("x", dtypes.floats), lambda x:
-   UOp.const(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),
+   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),
   # 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(dt, 1))).f(Ops.NOOP, dtype=dtypes.bool)),
+   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)),
   # 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.arg))) if _is_vec_xmm(y) else None),
+   lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
   (UPat.var("y", dtypes.int16s).index(UPat.cvar("c"), name="x"),
-   lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None),
+   lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
   (UPat.var("y", dtypes.int32s).index(UPat.cvar("c"), name="x"),
-   lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None),
+   lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
   (UPat.var("y", dtypes.int64s).index(UPat.cvar("c"), name="x"),
-   lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None),
+   lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
   (UPat.var("y", dtypes.floats).index(UPat.cvar("c"), name="x"),
-   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),
+   lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.val * x.dtype.itemsize))) if _is_vec_xmm(y) else None),
   # packed bitwise
   ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.max_numel() > 1 else None),
   ((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.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.SHLi, src=(a, imm(dtypes.uint8, c.val)))),
+  (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))),
+  (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))),
   (UPat.var("a", dtypes.ints) + UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   (UPat.var("a", dtypes.ints) * UPat.cvar("c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
@@ -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.arg if sz_uop is not None else rm_uop.dtype.itemsize
+    rm_sz = sz_uop.val if sz_uop is not None else rm_uop.dtype.itemsize
     reg_sz = reg_uop.dtype.itemsize if reg_uop is not None else 0
     sz = reg_sz or rm_sz
 
@@ -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.arg != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
+      if disp_uop.val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
       else: mod = 0b00
     else: mod = 0b11
     # x 0b0 and idx 0b100 means rsp which means no index exists
@@ -664,10 +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.arg)
+      inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.val)
     # IMM byte
     if imm_uop is not None:
-      if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.arg)
+      if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.val)
       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].arg),
+   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),
   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].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.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.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.arg) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize, o) if \
+        return [str(s.val) 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.arg}" if greg(idx) else "") + (f" + {disp.arg}" if disp.arg else "") + "]"]
+        return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.val}" if greg(idx) else "") + (f" + {disp.val}" if disp.val else "") + "]"]
 
       if len(x.src) > 4 and x.arg in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:])
       elif len(x.src) > 3 and x.arg in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:])

tinygrad/renderer/llvmir.py

@@ -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.arg}" if buf.addrspace == AddrSpace.ALU else None),
+   f"  {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.val}" if buf.addrspace == AddrSpace.ALU else None),
 
   # load/store
   (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"),
@@ -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.arg, u.dtype)
+      elif u.op is Ops.CONST: r[u] = lconst(u.val, u.dtype)
       elif u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype):
         r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast
       else:
@@ -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(x.src[2].dtype, 0.0)
+          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)
             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, x.dtype.max+x.arg+1) if x.arg < 0 else None),
+    (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype.max+x.val+1, x.dtype) if x.val < 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.arg, x.dtype)),
+    (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.val, x.dtype)),
     (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, x.dtype.itemsize if x.addrspace is AddrSpace.ALU else 8)),
     (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
     (UPat(Ops.STORE, src=(UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val"))),
@@ -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].arg
+    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
     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].arg)
+        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)
       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.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.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(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].arg]
+        r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].val]
         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(dtypes.uint32, 0xffffffff) << sext_am, UOp.const(dtypes.uint32, 0)) \
+  return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(0xffffffff, dtypes.uint32) << sext_am, UOp.const(0, dtypes.uint32)) \
         | 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(dtypes.uint32, 0), gate) if gate is not None else ()), dtype=dtypes.uint32)
+  buf = UOp.load(idx, *((UOp.const(0, dtypes.uint32), 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.arg else "false"),
+    (UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.val else "false"),
     (UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"),
-     lambda x: f"bitcast<u32>({x.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)}"),
+     lambda x: f"bitcast<u32>({x.val})" if x.val < 0 else f"{x.val&0xFFFFFFFF}u"),
+    (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}"),
     (UPat(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].arg * 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].val * 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), "AMD")):
+    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 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)

tinygrad/runtime/ops_cpu.py

@@ -28,20 +28,20 @@ def wait_prog():
   return (v:=UOp.param(0, dtypes.uint32, (1,), volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32))
 
 def timestamp_prog():
-  if WIN: val = UOp.const(dtypes.uint64, 0)
+  if WIN: val = UOp.const(0, dtypes.uint64)
   else:
     fn, ts = UOp.param(1, dtypes.uint64, (1,)), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG)
-    call = fn[0].load().call(UOp.const(dtypes.int, 6 if OSX else 1), ts[0], ret_dtype=dtypes.void) # clock_gettime(CLOCK_MONOTONIC, &ts)
+    call = fn[0].load().call(UOp.const(6 if OSX else 1, dtypes.int), ts[0], ret_dtype=dtypes.void) # clock_gettime(CLOCK_MONOTONIC, &ts)
     val = ts.after(call)[0].load() * 1_000_000_000 + ts.after(call)[1].load()
   return UOp.param(0, dtypes.uint64, (1,))[0].store(val)
 
 def quit_prog():
   fn = UOp.param(0, dtypes.uint64, (1 if WIN else 3,))
-  if WIN: return fn[0].load().call(UOp.const(dtypes.uint64, 0), ret_dtype=dtypes.void) # ExitThread(0)
+  if WIN: return fn[0].load().call(UOp.const(0, dtypes.uint64), ret_dtype=dtypes.void) # ExitThread(0)
   sem = UOp.param(1, dtypes.uint64, (1,))
 
   close = fn[2].load().call(sem[0], ret_dtype=dtypes.void) # sem_close(sem)
-  return fn.after(close)[0].load().call(UOp.const(dtypes.uint64, 0), ret_dtype=dtypes.void) # pthread_exit(0)
+  return fn.after(close)[0].load().call(UOp.const(0, dtypes.uint64), ret_dtype=dtypes.void) # pthread_exit(0)
 
 def worker_prog():
   ring = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True)
@@ -56,7 +56,7 @@ def worker_prog():
   return entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(cur)
 
 def host_wait(ctx, dst:UOp, val:UOp) -> UOp:
-  return (cur:=dst.after(loop:=UOp.loop(next(ctx))).index(UOp.const(dtypes.int, 0)).load()).end(loop, cur < val)
+  return (cur:=dst.after(loop:=UOp.loop(next(ctx))).index(UOp.const(0, dtypes.int)).load()).end(loop, cur < val)
 
 pm_host_opsel = PatternMatcher([(UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), host_wait)])
 
@@ -64,7 +64,7 @@ def encode_host_queue(q:UOp) -> UOp:
   # TODO: subset of hcq2 for now
   spins, (store,) = partition(graph_rewrite(q, pm_host_opsel, ctx=itertools.count(), walk=True, name="host opsel").src, lambda u: u.op is Ops.END)
   assert store.op is Ops.INS and store.arg == "store", f"host queue cannot encode {store.op} {store.arg}"
-  return store.src[0].after(*spins).index(UOp.const(dtypes.int, 0)).store(store.src[1])
+  return store.src[0].after(*spins).index(UOp.const(0, dtypes.int)).store(store.src[1])
 
 class CPUComputeQueue(HWQueue):
   def __init__(self, dev):

tinygrad/runtime/ops_dsp.py

@@ -13,7 +13,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat
 
 # NOTE: this just increases readability of the generated code
 dsp_string = PatternMatcher([
-  (UPat(Ops.CONST, (dtypes.int8, dtypes.uint8), name="x"), lambda ctx,x: str(x.arg)),
+  (UPat(Ops.CONST, (dtypes.int8, dtypes.uint8), name="x"), lambda ctx,x: str(x.val)),
 ])
 
 class DSPRenderer(ClangRenderer):

tinygrad/runtime/ops_python.py

@@ -102,7 +102,7 @@ class PythonProgram(Program['PythonDevice']):
         elif u.op is Ops.SPECIAL:
           if u.arg[0] == 'g': values[u] = [idxs[2-int(u.arg[-1])]] * warp_size
           elif u.arg[0] == 'l': values[u] = [x[2-int(u.arg[-1])] for x in warp]
-        elif u.op is Ops.CONST: values[u] = [u.arg] * warp_size
+        elif u.op is Ops.CONST: values[u] = [u.val] * warp_size
         elif u.op in {Ops.INDEX, Ops.SHRINK}:
           ret:list = []
           if u.src[0].addrspace == AddrSpace.ALU: