we don't float anymore

* fix pow(int, float)

* onnx

.github/workflows/benchmark.yml

@@ -521,6 +521,8 @@ jobs:
       run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
     - name: openpilot run_pickle 0.10.1 driving_vision
       run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py
+    - name: Test copy speeds
+      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})

.github/workflows/szdiff.yml

@@ -14,12 +14,15 @@ jobs:
     outputs:
       branchstat: ${{ steps.brstat.outputs.stat}}
     steps:
-      - name: Check code from PR branch 
+      - name: Check code from PR branch
         uses: actions/checkout@v6
         with:
           repository: ${{ github.event.pull_request.head.repo.full_name }}
           ref: ${{ github.event.pull_request.head.sha }}
           fetch-depth: 0
+          # PR code is only inspected with git rev-list, never executed
+          allow-unsafe-pr-checkout: true
+          persist-credentials: false
       - name: Check whether branch is up-to-date
         id: brstat
         run: |
@@ -51,6 +54,9 @@ jobs:
           repository: ${{ github.event.pull_request.head.repo.full_name }}
           ref: ${{ github.event.pull_request.head.sha }}
           path: pr
+          # PR code is only line-counted by master's sz.py, never executed
+          allow-unsafe-pr-checkout: true
+          persist-credentials: false
         # the base default to tinygrad master and cannot be other fork branch for security purpose
       - name: Checkout code from tinygrad master
         uses: actions/checkout@v6

CONTINUE.md

@@ -1,122 +0,0 @@
-# CONTINUE.md: PAD with Invalid instead of 0
-
-## Goal
-Make the low-level `Ops.PAD` pad with `Invalid` instead of `0`, while keeping
-the external `Tensor.pad` behavior unchanged.
-
-## Changes made (all 3 files are modified, see `git diff`)
-
-### 1. `tinygrad/schedule/indexing.py:92` — core change
-`convert_pad_to_where_to_keep_behavior_local` now uses `UOp.const(x.dtype, Invalid)`
-instead of `UOp.const(x.dtype, 0)` as the else value. This is what makes `Ops.PAD`
-pad with Invalid.
-
-### 2. `tinygrad/uop/symbolic.py:87-99` — Invalid propagation rules
-Added two new rules to `pm_data_invalid` so that `where(invalid_gate, a, const_b)`
-uses `b` (the const) in don't-care positions instead of poisoning to Invalid.
-This is needed so that `_pad_constant`'s mask `where(pad(ones_bool), base, value)`
-works — the mask is a `where(valid, True, Invalid)` gate, and the else `value`
-is a const.
-
-The rules are restricted to only match when the gate's valid value is a **const**
-(`UPat.cvar("x")`), to distinguish pad masks (where valid=True, a const) from
-gather masks (where valid=loaded_data, not a const). Without this restriction,
-`test_tensor_index` breaks because gather masks also create `where(cond, x, Invalid)`
-but need to keep poisoning.
-
-### 3. `tinygrad/mixin/op.py:280-289` — `_pad_constant` fix
-Swapped the `value == 0` early return for `value is Invalid` early return.
-When `value is Invalid`, just return `base` (which already has Invalid from
-`Ops.PAD`). For all other values (including 0), use the mask approach:
-`where(pad(ones_bool), base, const_value)`.
-
-## Current state
-- `test/unit/test_invalid_tensor.py` — **all 22 pass**
-- `test/unit/test_function.py` — **5 failures**, all multi-shard tests
-
-## The remaining bug: `cat` + multi-shard
-
-`cat` (op.py:716) uses `pad` + `usum` (element-wise ADD) to combine tensors:
-```python
-padded = [t.pad(...) for i,t in enumerate(tensors)]
-return padded[0].usum(*padded[1:])
-```
-
-When two shards are cat'd, each is padded and then summed. The valid masks
-are **complementary** (shard 0 valid in positions 0-1, shard 1 valid in 2-3).
-
-`_pad_constant` creates `where(mask_pad, data_pad, 0)` where:
-- `mask_pad = where(valid, True, Invalid)` — gate's valid value is const `True`
-- `data_pad = where(valid, data, Invalid)` — gate's valid value is loaded `data` (NOT const)
-
-The new const-specific rule handles the mask pad correctly. But for the data pad,
-the gate's valid value (`data`) is not a const, so the **non-const** lift-out rule
-fires: `where(valid, where(valid, data, Invalid), 0)` → `where(valid, where(valid, data, 0), Invalid)`.
-
-The `Invalid` else poisons the ADD. The binary Invalid rule lifts both gates out:
-`where(c6, data0, Invalid) + where(c8, data1, Invalid)` → `where(c6&c8, data0+data1, Invalid)`.
-
-Since `c6` and `c8` are complementary, `c6&c8` is always False → result is all Invalid → 0.
-
-### Master comparison
-On master, `convert_pad_to_where` uses `0` (not Invalid), so the ADD is just
-`where(c6, data0, 0) + where(c8, data1, 0)` with no Invalid, no lifting, works fine.
-
-### Debug output (with changes)
-```
-c16 = c6.where(c11.index(c13), 0)        # where(c6, load0, 0) — correct
-c22 = c6.where(0, c17.index(c20))         # where(c6, 0, load1) — correct
-c25 = (c6&c8).where((c16+c22), Invalid)   # WRONG: c6&c8 always False → all Invalid
-```
-
-### Master debug output
-```
-c13 = c6.where(c8.index(c10), 0)          # where(c6, load0, 0)
-c21 = c6.where(0, c14.index(c19))          # where(c6, 0, load1)
-c22 = c13+c21                              # plain ADD, no wrapper — correct
-```
-
-## Suggested fix approaches
-
-### Option A: General WHERE simplification rule
-Add a rule: `where(a, where(a, x, _), c)` → `where(a, x, c)`.
-When the outer and inner conditions are the same UOp, the inner else is
-unreachable. This would simplify `where(valid, where(valid, data, Invalid), 0)`
-→ `where(valid, data, 0)` before the lift-out rule can fire.
-Check if this rule already exists in `symbolic.py` — it may need to be added
-before the lift-out rules.
-
-### Option B: Don't use Ops.PAD for data in `_pad_constant`
-When `value is not Invalid`, avoid creating `Ops.PAD` on the data. Use `cat`
-or `expand` to create the padded tensor directly, bypassing the Invalid
-propagation entirely.
-
-### Option C: Make the lift-out rule use the outer else value
-Change the non-const lift-out rule: when `where(a, where(cond, x, Invalid), c)`
-and `c` is a const, use `c` as the else instead of `Invalid`. This is what the
-const-specific rule does, but it needs to also handle non-const gate valid values.
-
-## Test commands
-```bash
-# invalid tensor tests (currently pass)
-python -m pytest test/unit/test_invalid_tensor.py -x -q -n12
-
-# function tests (5 multi-shard failures)
-python -m pytest test/unit/test_function.py -x -q -n12
-
-# the specific failing test
-python -m pytest test/unit/test_function.py::TestFunctionMulti::test_simple_multi_sharded -x -q
-
-# debug the failing case
-DEBUG=6 python -c "
-from tinygrad import Tensor
-a = Tensor([1,2,3,4]).shard(['CPU', 'CPU:1'], axis=0)
-print(a.numpy())  # should be [1,2,3,4], gets [0,0,0,0]
-"
-```
-
-## Lint/typecheck
-```bash
-python -m mypy tinygrad/
-python -m ruff check .
-```

conftest.py

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

examples/mlperf/model_train.py

@@ -1462,6 +1462,8 @@ def train_llama3():
 
   @TinyJit
   def minibatch(tokens:Tensor):
+    for nxt in fp8_next_amax: nxt.assign(0)
+    for nxt in fp8_next_grad_amax: nxt.assign(0)
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if is_mp: tokens = tokens.shard(device)
     if not is_sharding: tokens = tokens.to(None)

examples/mlperf/models/flat_llama.py

@@ -37,8 +37,8 @@ def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
   return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
 
 def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
-           x_fp8:Tensor|None=None, x_new_amax:Tensor|None=None,
-           grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None) -> tuple[Tensor,...]:
+           x_fp8:Tensor|None=None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None,
+           next_amax_x:Tensor|None=None) -> tuple[Tensor,...]:
   if not fp8:
     if ASM_GEMM:
       from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
@@ -56,13 +56,14 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
     else:
       x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
       out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
-    return out, (amax_x.detach() if amax_x is not None else None), x_q
+    return out, x_q
   if x_fp8 is None:
-    if FUSED_INPUT_QUANTIZE and amax_x is not None:
+    if FUSED_INPUT_QUANTIZE:
       from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
-      x_fp8, _, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
+      x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
     else:
-      x_fp8, _, x_new_amax = quantize_fp8(x, amax_state=amax_x)
+      x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
+      next_amax_x.assign(new_amax_x)
   if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
     if can_use_asm_gemm(x_fp8, w.T):
@@ -73,51 +74,51 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
       else:
         out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
                        next_grad_amax_state=next_grad_amax_state)
-      return out, x_new_amax, x_fp8
-  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8
+      return out, x_fp8
+  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8
 
 def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                         grad_amax_state:Tensor, next_grad_amax_state:Tensor):
+                         next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor):
   if FUSED_ADD_NORM_MUL_QUANTIZE:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
-    x_fp8, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
+    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,
                        grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
     return out, x_normed, rrms, ret
   x_normed, rrms = rmsnorm(x, eps)
   out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state)
+                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
   return out, x_normed, rrms, ret
 
 def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                             grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
+                             next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
   if FUSED_ADD_NORM_MUL_QUANTIZE:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
-    x_fp8, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
+    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,
                        grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
     return out, h, x_normed, rrms, ret
   h = x + residual
   x_normed, rrms = rmsnorm(h, eps)
   out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state)
+                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
   return out, h, x_normed, rrms, ret
 
 def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
-                             amax_x2: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:
     from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
-    x2_fp8, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
-                                                 next_grad_amax_state=next_grad_amax_xw13)
-    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2, x_new_amax=new_amax_x2,
+    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)
+    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2,
                        grad_amax_state=grad_amax_xout, next_grad_amax_state=next_grad_amax_xout)
     return out, ret
   hidden = x_w13.shape[-1] // 2
   x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
   out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
-                     next_grad_amax_state=next_grad_amax_xout)
+                     next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
   return out, ret
 
 class FlatTransformer:
@@ -186,14 +187,14 @@ class FlatTransformer:
 
   def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
                 amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
+                next_amax_xqkv:Tensor, next_amax_xo:Tensor,
                 grad_amax_xqkv:Tensor, grad_amax_xo:Tensor, next_grad_amax_xqkv:Tensor, next_grad_amax_xo:Tensor):
     bsz, seqlen, _ = x.shape
-    amaxs, saves = [], []
+    saves = []
 
-    xqkv, x_normed, rrms, (new_amax, *s) = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
+    xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
                                                                   amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
-                                                                  next_grad_amax_state=next_grad_amax_xqkv)
-    amaxs.append(new_amax)
+                                                                  next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv)
     saves.extend([x_normed, rrms, *s, xqkv])
     if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
@@ -211,64 +212,62 @@ class FlatTransformer:
       attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
     attn = attn.reshape(bsz, seqlen, -1)
 
-    out, new_amax, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
-                               next_grad_amax_state=next_grad_amax_xo)
-    amaxs.append(new_amax)
+    out, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
+                               next_grad_amax_state=next_grad_amax_xo, next_amax_x=next_amax_xo)
     saves.extend([*s, out])
-    return out, amaxs, saves
+    return out, saves
 
   def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
-    amaxs, saves = [], []
+    saves = []
 
     if SPLIT_W13:
       h = x + residual
       x_normed, rrms = rmsnorm(h, self.norm_eps)
       saves.extend([x_normed, rrms])
       inp = x_normed * kwargs["ffn_norm"]
-      x_w1, new_amax, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
-                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"])
-      amaxs.append(new_amax)
+      x_w1, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
+                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"],
+                                  next_amax_x=kwargs["next_amax_x1"])
       saves.extend([*s, x_w1])
-      x_w3, new_amax, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
-                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"])
-      amaxs.append(new_amax)
+      x_w3, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
+                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"],
+                                  next_amax_x=kwargs["next_amax_x3"])
       saves.extend([*s, x_w3])
       if FUSED_SILU_W13 and MXFP8:
         from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
         aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
-        out, new_amax, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
-                                   w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
-                                   next_grad_amax_state=kwargs["next_grad_amax_xout"])
+        out, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
+                         w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
+                         next_grad_amax_state=kwargs["next_grad_amax_xout"], next_amax_x=kwargs["next_amax_x2"])
         out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
       else:
-        out, new_amax, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
-                                   grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"])
-      amaxs.append(new_amax)
+        out, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
+                         grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"],
+                         next_amax_x=kwargs["next_amax_x2"])
       saves.extend([*s, out])
     else:
-      x_w13, h, x_normed, rrms, (new_amax, *s) = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
+      x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
                                                                           self.norm_eps, amax_x=kwargs["amax_x13"],
+                                                                          next_amax_x=kwargs["next_amax_x13"],
                                                                           grad_amax_state=kwargs["grad_amax_xw13"],
                                                                           next_grad_amax_state=kwargs["next_grad_amax_xw13"])
-      amaxs.append(new_amax)
       saves.extend([x_normed, rrms, *s, x_w13])
-      out, (new_amax, *s) = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
+      out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
+                                                     next_amax_x2=kwargs["next_amax_x2"],
                                                      grad_amax_xw13=kwargs["grad_amax_xw13"],
                                                      next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
                                                      grad_amax_xout=kwargs["grad_amax_xout"],
                                                      next_grad_amax_xout=kwargs["next_grad_amax_xout"])
-      amaxs.append(new_amax)
       saves.extend([*s, out])
-    return out, h, amaxs, saves
+    return out, h, saves
 
   @function(precompile=True, precompile_backward=True)
   def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
-    ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
+    attn, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
+    ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
     h = h + ffn
-    amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs))
-    if save: return (h, *amaxs, *attn_saves, *ffn_saves)
-    else: return (h, *amaxs)
+    if save: return (h, *attn_saves, *ffn_saves)
+    else: return (h,)
 
   def shard(self, device:tuple[str, ...], mp:bool=False):
     from tinygrad.nn.state import get_parameters
@@ -319,21 +318,21 @@ class FlatTransformer:
     for i in range(self.n_layers):
       attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
                          amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
+                         next_amax_xqkv=na["xqkv"][i], next_amax_xo=na["xo"][i],
                          grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
                          next_grad_amax_xqkv=nga["xqkv"][i], next_grad_amax_xo=nga["xo"][i])
       ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i],
-                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i])
+                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i],
+                        next_amax_x2=na["x2"][i])
       if SPLIT_W13:
         ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
+                          next_amax_x1=na["x1"][i], next_amax_x3=na["x3"][i],
                           s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i],
                           next_grad_amax_xw1=nga["xw1"][i], next_grad_amax_xw3=nga["xw3"][i])
       else:
         ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i],
-                          next_grad_amax_xw13=nga["xw13"][i])
-      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
-      amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
-      for name, new_val in zip(amax_names, ret[:len(amax_names)]):
-        na[name][i].assign(new_val)
+                          next_grad_amax_xw13=nga["xw13"][i], next_amax_x13=na["x13"][i])
+      h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
 
     logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
     return logits
@@ -416,6 +415,9 @@ if __name__ == "__main__":
   @TinyJit
   def fwd_bwd(tokens:Tensor):
     with Timing("python forward: "):
+      for amax_dict in (model._fp8_next_amax, model._fp8_next_grad_amax):
+        for ts in amax_dict.values():
+          for nxt in ts: nxt.assign(0)
       logits = model(tokens[:, :-1], save=llama_size=="8B")
       loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     with Timing("python backward: "):

examples/mlperf/models/gpt_oss.py

@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, quantize_mxfp8
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
@@ -39,8 +39,11 @@ def quant_dequant_mx(x:Tensor) -> Tensor:
   fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
   return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
 
+def _mx_scale(e8:Tensor) -> Tensor:
+  return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
+
 def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
-  return w_q.cast(dtypes.bfloat16) * _mx_block_scale(w_scale)
+  return w_q.cast(dtypes.bfloat16) * _mx_scale(w_scale)
 
 @functools.cache
 def _dequant_fwd_fxn(wq_p, ws_p, device):
@@ -48,7 +51,7 @@ def _dequant_fwd_fxn(wq_p, ws_p, device):
 
 def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
   w_scale = Tensor(call.src[2])
-  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_block_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
+  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
 
 def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
   fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)

examples/mlperf/optim.py

@@ -96,7 +96,7 @@ class GradAccClipAdamW(Optimizer):
     up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
     new_w = w.detach() - up
     if master is not None: master.assign(new_w)
-    if self.zero: new_w = self._zero_gather(new_w)
+    if self.zero and not (MXFP8 and t.dtype in dtypes.fp8s): new_w = self._zero_gather(new_w)
     # when master is offloaded to a different device than the param, results are resharded back onto the param's (sharded) device
     offloaded = master is not None and master.device != t.device
     if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16:
@@ -106,6 +106,7 @@ class GradAccClipAdamW(Optimizer):
       if MXFP8:
         from extra.gemm.cdna_asm_gemm import quantize_mxfp8
         w_q, w_e8, _ = quantize_mxfp8(new_w.reshape(-1, new_w.shape[-1]))
+        if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
         new_e8 = w_e8.reshape(t._inv_scale.shape)
         t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
         ret = w_q.reshape(new_w.shape)

extra/f16_decompress.py

@@ -5,10 +5,10 @@ def bit_extract(x: Tensor, e: int, s: int) -> Tensor:
   return (x >> s) & mask
 
 def u16_to_f16(x: Tensor) -> Tensor:
-  sign = bit_extract(x, 15, 15).float()
+  sign = bit_extract(x, 15, 15).bool()
   exponent = bit_extract(x, 14, 10).float()
   fraction = bit_extract(x, 9, 0).float()
-  return sign.where(-1, 1) * exponent.where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
+  return sign.where(-1, 1) * exponent.bool().where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
 
 def u32_to_f16(oo: Tensor) -> Tensor:
   f1 = u16_to_f16(oo>>16)

extra/gemm/cdna_asm_gemm.py

@@ -128,6 +128,11 @@ def _mx_block_scale(e8:Tensor) -> Tensor:
   rows, scale_K = e8.shape
   return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(rows, scale_K, 1).expand(rows, scale_K, 32).reshape(rows, scale_K*32)
 
+def _mx_block_scale_3d(e8:Tensor) -> Tensor:
+  # batched (E, rows, scale_K) dequant scale 2^(e8-127) broadcast to (E, rows, scale_K*32)
+  E, rows, scale_K = e8.shape
+  return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(E, rows, scale_K, 1).expand(E, rows, scale_K, 32).reshape(E, rows, scale_K*32)
+
 counters = {"used":0, "todos":[]}
 def todo(msg:str) -> bool: counters["todos"].append(msg); return False
 def _asm_gemm_report():
@@ -169,10 +174,10 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   m = UOp.range(M, 1, AxisType.LOOP)
   n = UOp.range(N, 2, AxisType.LOOP)
   k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.flatten().index((m*UOp.const(dtypes.index, K)+k))*
-         B.flatten().index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
+         B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
   red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
-  store = C.flatten().index((m*UOp.const(dtypes.index, N)+n)).store(red).end(m, n)
+  store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n)).store(red).end(m, n)
   return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
 
 # ** bf16 A @ B.T kernel in C
@@ -275,10 +280,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=
       elif getenv("FUSED_GRAD_QUANTIZE", 0):
         grad_amax_t = Tensor(grad_amax_state, device=a.device)
         g_amax = grad_amax_t
-        g_fp8, _, new_grad_amax, _ = quantize_fp8_delayed(g_t, g_amax)
-        store_effect = next_grad_amax_state.store(new_grad_amax.uop)
-        assert g_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {g_fp8.uop.op}"
-        g_fp8 = Tensor(g_fp8.uop.replace(src=g_fp8.uop.src + (store_effect,)), device=a.device)
+        g_fp8, _ = quantize_fp8_delayed(g_t, g_amax, Tensor(next_grad_amax_state, device=a.device))
       else:
         grad_amax_t = Tensor(grad_amax_state, device=a.device)
         g_amax = grad_amax_t

extra/hcq2/hcq2.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any
 import struct, functools, time, collections, itertools
 from dataclasses import replace, dataclass
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE
 from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
@@ -49,7 +49,7 @@ def make_getaddr(u, device=None):
   return UOp(Ops.GETADDR, dtypes.uint64, src=(u,), arg=device or to_tuple(u.device)[0])
 
 def make_ins(op, *srcs):
-  return UOp(Ops.INS, dtypes.void, tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs), op)
+  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs))
 
 def make_placeholder(devs, size:int, dtype, name=None, unique=True) -> UOp:
   return UOp.param(next(UOp.unique_num) if unique else 0, dtype, shape=(size,), device=devs).rtag(name or "temp")
@@ -102,159 +102,139 @@ def stage_copy(dst:UOp, src:UOp) -> UOp|None:
 pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
 
 # *****************
-# 2.1. tag hcq calls
-
-def tag_hcq_call(ctx:itertools.count, call:UOp) -> UOp:
-  if (hcq_devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is None: return call
-
-  queue = "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0"
-  info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), to_tuple(hcq_devs), queue)
-  return call.replace(arg=replace(call.arg, aux=info)).rtag(next(ctx))
-pm_tag_hcq_calls = PatternMatcher([(UPat(Ops.LINEAR, name="l"), lambda ctx, l: l.replace(src=tuple(tag_hcq_call(ctx, s) for s in l.src)))])
-
-# *****************
-# 2.2. deps tracking
-# device.timeline_signal/value are the per-device schedule epoch. Before a schedule queue accesses memory owned by device N for the first time,
-# it waits for device[N].timeline_signal >= device[N].timeline_value - 1. This orders the schedule after all prior schedules that touched device N.
-#
-# queue.timeline_signal/value are per-queue progress counters used only inside a schedule.
-# Only the owner queue signals its queue.timeline_signal. Values are monotonic.
-#
-# At schedule end, one finalizer queue per touched device[N] waits for every active queue on device[N] to reach its schedule-local
-# final queue.timeline value, then signals device[N].timeline_signal with the schedule's reserved device epoch. After that, buffers/transients
-# for device N from this schedule are safe for the next schedule
-#
-# C programs reserve and bump timeline values, then patch command buffers with the concrete wait/signal values.
+# 2. deps
 
 class HCQDepsTracker(DepsTracker):
   @staticmethod
   def _key(buf:Any) -> tuple[Any, int, int]:
     return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
 
-def make_deps(u:UOp, dep_lanes:list[tuple[UOp, int, int]], nlanes:int) -> UOp:
-  deps:dict[UOp, list[int|None]] = collections.defaultdict(lambda: [None]*nlanes)
-  for dep, dlane, lane in dep_lanes: deps[dep][lane] = dlane
-  return u.after(*deps, arg=tuple(tuple(v) for v in deps.values()))
-
-def sched_sync(ctx:DepsTracker, call:UOp) -> UOp|None:
-  if not isinstance(call.arg.aux, HCQInfo): return None
-
+def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]:
   refs = get_call_arg_uops(call)
-  outs, _ = get_call_outs_ins(call)
-  devices, queue = call.arg.aux.device, call.arg.aux.queue
+  return [[b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs] for lane in range(len(devices))]
 
-  dep_lanes:list[tuple[UOp, int, int]] = []
-  for lane, d in enumerate(devices):
-    lane_refs = [b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs]
-    for dep, dlane in ctx.access_resources(lane_refs, outs, (call, lane)): dep_lanes.append((dep, dlane, lane))
+def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tuple[str, ...], str, int]) -> list[tuple[tuple, int, int]]:
+  dep_lanes:list[tuple[tuple, int, int]] = []
+  for lane, bufs in enumerate(bufs_by_lane):
+    dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, write if write is not None else range(len(bufs)), (key, lane))]
+  return dep_lanes
 
+def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
+  # opt1: same-queue ops are fifo-ordered
   if devices[0].split(":")[0] in {"AMD", "QCOM"} or queue.startswith("COPY"):
-    dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep.arg.aux.device[dlane], dep.arg.aux.queue) != (devices[lane], queue)]
+    dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep[0][dlane], dep[1]) != (devices[lane], queue)]
 
-  # keep latest dep per (dep device, queue, cur lane)
-  latest = {((dep.arg.aux.device[dlane], dep.arg.aux.queue), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0].tag)}
-  return make_deps(call, [(dep, dlane, lane) for (_, lane), (dep, dlane) in latest.items()], len(devices))
-pm_sched_sync = PatternMatcher([(UPat(Ops.CALL, name="call"), sched_sync)])
+  # opt2: keep latest dep per (dep device, queue, cur lane)
+  latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
+  deps:dict[tuple, list[int|None]] = collections.defaultdict(lambda: [None]*len(devices))
+  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane] = dlane
+
+  waits = []
+  for (ddevs, dqueue, dtag), lanes in deps.items():
+    sig = make_mstack([make_signal(d if dl is None else ddevs[dl], queue=dqueue, sentinel=dl is None) for dl, d in zip(lanes, devices)])
+    val = make_mstack([make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)])
+    waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(dtypes.int, 0)) + dtag)))
+  return waits, {dtag for _, _, dtag in deps}
+
+def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
+                      tracker:HCQDepsTracker) -> tuple[list[UOp], set[int]]:
+  # collect all buffers which belong to devices
+  dev_bufs:dict[str, dict[int, Any]] = collections.defaultdict(dict)
+  for call, devices in batch:
+    for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)):
+      for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b
+
+  zero, n, finalizers, waited = UOp.const(dtypes.int, 0), len(batch_info), [], set()
+  for _, devgroup in itertools.groupby(sorted(dedup([d for devs, _ in batch_info for d in devs])), key=lambda d: d.split(":")[0]):
+    devs = tuple(devgroup)
+
+    # to finalize the batch, sync all accesses from other devices to buffers that belong to this device
+    fin_deps = [dl for dl in _get_deps(tracker, [list(dev_bufs[d].values()) for d in devs], None, key=(devs, "COMPUTE:0", n)) if dl[0][2] < n]
+    waits, cur_waited = _build_wait_cmds(fin_deps, devs, "COMPUTE:0")
+    waited |= cur_waited
+
+    # wait the syncs, store the device epoch; value bumps are a separate call: no lane may bump until every lane has patched its waits
+    store = UOp(Ops.INS, arg="store", src=(make_signal(devs), (tl:=make_signal_value(devs)).index(zero)))
+    submit = make_submit(*waits, store, devs=devs, queue="COMPUTE:0")
+    upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), n) for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)])]
+    bump = UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd])
+    finalizers += [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for b in (submit, bump)]
+  return finalizers, waited
+
+def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
+  batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch]
+
+  # schedule deps
+  waited:set[int] = set()
+  deps_tracker = HCQDepsTracker()
+  call_waits:list[list[UOp]] = []
+  for tag, ((call, _), (devices, queue)) in enumerate(zip(batch, batch_info)):
+    deps = _get_deps(deps_tracker, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], key=(devices, queue, tag))
+    cmds, cur_waited = _build_wait_cmds(deps, devices, queue)
+    call_waits.append(cmds)
+    waited |= cur_waited
+
+  # build finalizers
+  finalizers, finalizer_waited = _build_finalizers(batch, batch_info, deps_tracker)
+  waited |= finalizer_waited
+
+  src = []
+  for tag, ((call, _), (devices, queue), cmds) in enumerate(zip(batch, batch_info, call_waits)):
+    # first queue use, sync prior device work with main signal
+    if batch_info.index((devices, queue)) == tag:
+      epoch = UOp(Ops.INS, arg="wait", src=(make_signal(devices), make_signal_value(devices).index(0) - 1))
+      cmds = [UOp(Ops.INS, arg="barrier", src=()), epoch] + cmds
+
+    # and make hcq call
+    info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), devices, queue)
+    cmds = [*cmds, call.replace(arg=replace(call.arg, aux=info))]
+
+    # signal queue timeline if someone waits for us
+    if tag in waited: cmds += [UOp(Ops.INS, arg="store", src=(make_signal(devices, queue), make_signal_value(devices, queue).index(0) + tag))]
+    src.append(UOp.custom_function("hcq", make_submit(*cmds, devs=devices, queue=queue).sink()).call(name="hcq", aux=info))
+  return src + finalizers
+
+def sched_hcq_batches(l:UOp) -> UOp:
+  srcs:list[UOp] = []
+  batch:list[tuple[UOp, tuple[str, ...]]] = []
+  for call in l.src:
+    if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs)))
+    else: srcs, batch = srcs + _finalize_batch(batch) + [call], []
+  return l.replace(src=tuple(srcs + _finalize_batch(batch)))
+pm_sched_hcq_batches = PatternMatcher([(UPat(Ops.LINEAR, name="l"), sched_hcq_batches)])
 
 # *****************
-# 2.3. merge into queues
+# 3. merge into queues
 
-def _merged_hcq_call(calls:list[UOp]):
-  info = replace(unwrap_after(calls[0]).arg.aux, estimates=sum((unwrap_after(c).arg.aux.estimates for c in calls), start=Estimates()))
-  cmdbuf = make_submit(*calls, devs=info.device, queue=info.queue)
-  return UOp.custom_function("hcq", cmdbuf.sink()).call(name="hcq", aux=info)
+def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
+  if len(calls) == 1: return calls[0]
+  info = replace(calls[0].arg.aux, name=f"submit {calls[0].arg.aux.queue} ({len(calls)})",
+                 estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()))
+  cmds = [cmd for c in calls for cmd in get_submit(c).src[0].src]
+  return UOp.custom_function("hcq", make_submit(*cmds, devs=info.device, queue=info.queue).sink()).call(name="hcq", aux=info)
 
 def merge_queues(linear:UOp) -> UOp:
   new_src:list[UOp] = []
-  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of calls, kept in submit order
+  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of hcq calls, kept in submit order
+  limits = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)
 
   for call in linear.src:
-    if not isinstance(unwrap_after(call).arg.aux, HCQInfo):
+    if not isinstance(info:=call.arg.aux, HCQInfo) or info.name == "hcq_finalizer": # non-hcq call or finalizer: close all open queues
       new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
       continue
 
-    devices, queue = unwrap_after(call).arg.aux.device, unwrap_after(call).arg.aux.queue
-
-    if (old:=opened_qs.pop((devices, queue), None)) is not None: new_rec = old + [call]
+    if (old:=opened_qs.pop(key:=(info.device, info.queue), None)) is not None:
+      if limits[key] and len(old) >= limits[key]: new_src, old, limits[key] = new_src + [_merged_hcq_call(old)], [], limits[key] * 2
+      new_rec = old + [call]
     else:
       # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept
-      closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devices)]
+      closing = [k for k in opened_qs if k[1] == info.queue and set(k[0]) & set(info.device)]
       new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
       new_rec = [call]
-    opened_qs[(devices, queue)] = new_rec
+    opened_qs[(info.device, info.queue)] = new_rec
   return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))
 pm_merge_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), merge_queues)])
 
-# *****************
-# 2.4. finalizer
-
-def add_finalizer(ctx:itertools.count, linear:UOp) -> UOp:
-  # collect by device type
-  parts:dict[str, list[UOp]] = collections.defaultdict(list)
-  for call in linear.src:
-    if (c:=unwrap_after(call)).src[0].op is not Ops.CUSTOM_FUNCTION or c.src[0].arg != "hcq": continue
-    parts[c.arg.aux.device[0].split(':')[0]].append(unwrap_after(get_submit(call).src[0].src[0]))
-
-  nbump = next(ctx)
-  finalizers = []
-  for calls in parts.values():
-    devs = tuple(dedup(d for call in calls for d in unwrap_after(call).arg.aux.device))
-    zero = UOp.const(dtypes.int, 0)
-    tl = make_signal_value(devs)
-
-    # split each (multi-device) call into per-device deps, then store the device timeline value into the device signal after them
-    dep_lanes = [(call, dlane, devs.index(d)) for call in calls for dlane, d in enumerate(unwrap_after(call).arg.aux.device)]
-    store = make_deps(make_signal(devs).store(tl.index(zero)), dep_lanes, len(devs))
-    submit = make_submit(store, devs=devs, queue="COMPUTE:0")
-
-    upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), nbump) for qn in dedup([unwrap_after(call).arg.aux.queue for call in calls])]
-    patches = [s.after(submit).index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd]
-    finalizers.append(UOp.custom_function("hcq", UOp.barrier(*patches).sink()).call(aux=HCQInfo("hcq finalizer", Estimates(), devs, "COMPUTE:0")))
-  return linear.replace(src=linear.src + tuple(finalizers))
-pm_add_finalizer = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), add_finalizer)])
-
-# *****************
-# 2.5. global sync
-
-def add_global_sync(ctx:set[tuple[str, ...]], submit:UOp, q:UOp) -> UOp|None:
-  if (devs:=q.arg[0]) in ctx: return None
-  ctx.add(devs)
-
-  # some devices from a command buffer might be used for the first time this schedule, so we wait for their global timeline epoch.
-  wait = (make_signal(devs).index(zero:=UOp.const(dtypes.int, 0)).load() >= make_signal_value(devs).index(zero) - 1).wait()
-  return submit.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), wait, *q.src)),))
-pm_add_global_sync = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_global_sync)])
-
-# *****************
-# 3.1. lower loads/stores
-
-def add_loads(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
-  cur_devs = q.arg[0]
-  new_src:list[UOp] = []
-  for s in q.src:
-    if s.op is Ops.AFTER:
-      for lanes, dep in zip(s.arg, s.src[1:]):
-        devs, queue = dep.arg.aux.device, dep.arg.aux.queue
-        ctx.add(dep.tag) # mark op to update signal.
-
-        sig = make_mstack([make_signal(d if dl is None else devs[dl], queue=queue, sentinel=dl is None) for dl, d in zip(lanes, cur_devs)])
-        val = make_mstack([make_signal_value(d if dl is None else devs[dl], queue=queue) for dl, d in zip(lanes, cur_devs)]).index(UOp.const(dtypes.int, 0))
-        new_src.append((sig.index(UOp.const(dtypes.int, 0)).load() >= val + dep.tag).wait())
-      s = s.src[0]
-    new_src.append(s)
-  return submit.replace(src=(q.replace(src=tuple(new_src)),))
-pm_add_inner_loads = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_loads)])
-
-def add_stores(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
-  devs, queue = q.arg
-  new_src:list[UOp] = []
-  for op in q.src:
-    new_src.append(op)
-    if (sigval:=unwrap_after(op).tag) in ctx:
-      new_src.append(make_signal(devs, queue=queue).store(make_signal_value(devs, queue=queue).index(UOp.const(dtypes.int, 0)) + sigval))
-  return submit.replace(src=(q.replace(src=tuple(new_src)),))
-pm_add_inner_stores = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_stores)])
-
 # *****************
 # 4.1. hcq lowering: programs
 
@@ -279,7 +259,7 @@ def is_value_known_at_link(val:UOp) -> bool:
   addressed_bufs = [b for g in val.toposort() if g.op is Ops.GETADDR for b in unwrap_mstack(g.buf_uop)]
 
   # addr of input params is not known at link time
-  return not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
+  return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
 
 def is_link_patch(p:UOp, jit:bool) -> bool:
   store = p.src[0] if (is_binary_patch:=p.op is Ops.END) else p
@@ -308,7 +288,7 @@ pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION
 def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
   bare = {g: g.replace(src=(unwrap_after(g.src[0]),)) for g in gaddrs}
 
-  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, to_tuple(b.tag)))
+  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(b.tag)))
   slots, table = {g:i for i,g in enumerate(order)}, make_placeholder(call.arg.aux.device, len(order), dtypes.uint64, name)
 
   reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs}
@@ -374,7 +354,7 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
       sizes[b.tag] = offs[b] + b.max_numel()
   counts = collections.Counter(b.tag for b in bufs)
   bases = {b.tag:make_placeholder(b.device, sizes[b.tag], b.dtype, b.tag) for b in bufs if counts[b.tag] > 1}
-  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.index, offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
+  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
   return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
 pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
 
@@ -391,19 +371,16 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
   if input_uops is not None: linear = graph_rewrite(linear, pm_replace_buffers, ctx=input_uops, walk=True, enter_calls=True, name="replace buffer")
 
   if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, jit)))) is None:
-    # schedule
+    # prep
     linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True)
     linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
-    linear = graph_rewrite(linear, pm_tag_hcq_calls, ctx=(enumerator:=itertools.count(0)), walk=True, name="tag hcq calls")
-    linear = graph_rewrite(linear, pm_sched_sync, ctx=HCQDepsTracker(), walk=True, name="schedule sync")
-    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True)
+
+    # schedule
+    linear = graph_rewrite(linear, pm_sched_hcq_batches, walk=True, name="schedule hcq batches")
+    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True, enter_calls=True)
     linear = graph_rewrite(linear, pm_merge_queues, walk=True, name="merge queues")
-    linear = graph_rewrite(linear, pm_add_finalizer, ctx=enumerator, walk=True, name="add finalizer")
-    linear = graph_rewrite(linear, pm_add_global_sync, ctx=set(), walk=True, name="add global sync", enter_calls=True)
 
     # lowering to hcq ir
-    linear = graph_rewrite(linear, pm_add_inner_loads, ctx=(waited:=set()), walk=True, name="add loads", enter_calls=True)
-    linear = graph_rewrite(linear, pm_add_inner_stores, ctx=waited, walk=True, name="add stores", enter_calls=True)
     linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
     linear = graph_rewrite(linear, pm_pack_placeholders, walk=True, name="pack placeholders")
 
@@ -508,7 +485,8 @@ class HCQ2Compiled(Compiled):
     self.rt_allocator = BumpAllocator(64 << 20, wrap=False)
 
   def new_buffer(self, b:UOp, jit:bool) -> Buffer:
-    if jit or b.tag in HCQ_CACHE_TAGS: return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(cpu_access=True, nolru=True))
+    if jit or b.tag in HCQ_CACHE_TAGS:
+      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=True, cpu_access=True, nolru=True))
     return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))
 
   @functools.cache

extra/hcq2/ops_amd2.py

@@ -90,7 +90,7 @@ def memory_barrier(ctx):
                  reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
     acquire_mem(ctx)))
 
-def pm4_wait(ctx, x, y): return wait_reg_mem(ctx, y, mem=make_getaddr(x.buf_uop, ctx.devs))
+def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=make_getaddr(dst, ctx.devs))
 
 def pm4_barrier(ctx): return memory_barrier(ctx)
 
@@ -138,10 +138,10 @@ def pm4_program(ctx, call, prg):
 pm_pm4_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
 
-  (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), pm4_wait),
-  (UPat(Ops.BARRIER), pm4_barrier),
-  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
-  (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
+  (UPat(Ops.INS, arg="barrier"), pm4_barrier),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])
 
 def pm4_submit(cmdbuf, devs):
@@ -184,10 +184,10 @@ def sdma_copy(ctx, call):
      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)) for off in range(0, sz, ctx.max_copy_size)]))
 
-def sdma_wait(ctx, x, y):
+def sdma_wait(ctx, 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 make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(x.buf_uop, ctx.devs)), y, 0xffffffff,
+  return make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(dst, 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, dst, val):
@@ -202,10 +202,10 @@ def sdma_timestamp(ctx, dst):
 pm_sdma_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
 
-  (UPat(Ops.BARRIER), lambda: UOp(Ops.NOOP, dtypes.void, ())),
-  (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), sdma_wait),
-  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp),
-  (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
+  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), sdma_wait),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
 ])
 
 def sdma_submit(cmdbuf, devs):

extra/llama_kernels/__init__.py

@@ -20,11 +20,6 @@ def local_abs_max(x:Tensor) -> Tensor:
   fxn = _local_abs_max_fxn(param.uop, x.device)
   return Tensor(fxn[0].uop.call(x.uop).gettuple(0))
 
-def scalar_amax(amax_buf:Tensor) -> Tensor:
-  if isinstance(amax_buf.device, tuple):
-    return local_abs_max(amax_buf).detach()
-  return amax_buf.max().detach()
-
 def shard_shape(shape:tuple, axis:int, ndev:int) -> list:
   s = list(shape)
   s[axis] //= ndev

extra/llama_kernels/cast_amax/__init__.py

@@ -3,7 +3,7 @@ import functools, pathlib
 from tinygrad import Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, alloc_local, scalar_amax, dname_of
+from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, dname_of
 
 # module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, delayed amax UOp)
 # lets cdna_asm_gemm's bwd reuse the fp8 companion produced by the fused silu_mul bwd kernel
@@ -11,13 +11,13 @@ from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like,
 _grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {}
 
 @functools.cache
-def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, grad_amax:UOp,
+def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_next:UOp, grad_amax:UOp,
                           xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
   hidden = xw13.shape[2] // 2
   n_elems = xw13.shape[0] * xw13.shape[1] * hidden
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 3 + n_elems * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(grad_xw13_fp8.base, grad_amax_buf.base, grad_amax.base,
+  mem = n_elems * 2 * 3 + n_elems * 2 + 4 + 4
+  sink = UOp.sink(grad_xw13_fp8.base, grad_amax_next.base, grad_amax.base,
                   xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
   src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
@@ -25,14 +25,14 @@ def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, grad_amax:UOp,
                                UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
 
 @functools.cache
-def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp,
+def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_out:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp,
                                 next_grad_amax_state:UOp, dname:str) -> UOp:
   # NOTE: grad_amax_state is plumbed through as an unused fwd input so the bwd kernel can read it via kernel.src
   hidden = xw13.shape[2] // 2
   n_elems = xw13.shape[0] * xw13.shape[1] * hidden
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 4
-  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
+  mem = n_elems * 2 * 2 + n_elems + 4
+  sink = UOp.sink(fp8_out.base, amax_out.base, xw13.base, amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
   src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden)
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
@@ -43,26 +43,23 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
   device = xw13.device
   axis = xw13.axis if isinstance(device, tuple) else None
   grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3,  device, axis)
-  grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32,  device, axis)
+  grad_amax_next = Tensor(next_grad_amax_state, device=device)
   grad_amax_state_t = Tensor(grad_amax_state, device=device)
   fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
   grad_amax = grad_amax_state_t.empty_like()
-  grad_xw13_fp8, grad_amax_buf, grad_amax, *_ = Tensor.custom_kernel(
-    grad_xw13_fp8, grad_amax_buf, grad_amax,
+  grad_xw13_fp8, grad_amax_next, grad_amax, *_ = Tensor.custom_kernel(
+    grad_xw13_fp8, grad_amax_next, grad_amax,
     Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
     Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
   grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16)
-  new_grad_amax = scalar_amax(grad_amax_buf)
-  store_effect = next_grad_amax_state.store(new_grad_amax.uop)
   assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}"
-  grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,))
   # Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a.
-  _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, grad_amax_state_t.uop)
+  _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8.uop, grad_amax_state_t.uop)
   return (None, None, grad_xw13_uop, None, None, None)
 
 def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor,
-                           next_grad_amax_state:Tensor) -> tuple[Tensor, Tensor]:
-  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, new_amax)
+                           next_grad_amax_state:Tensor, amax_out:Tensor) -> Tensor:
+  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns fp8.
   # grad_amax_state: delayed amax for grad_xw13 fp8 quantization in the backward.
   assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
   MBS, SEQ, H2 = xw13.shape
@@ -70,8 +67,7 @@ def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_
   HIDDEN = H2 // 2
   axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None
   fp8_out  = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,      xw13.device, axis)
-  amax_buf = alloc_local((NUM_WG,),         dtypes.float32, xw13.device, axis)
   fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device))
-  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, grad_amax_state, next_grad_amax_state,
+  fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, xw13, amax_state, grad_amax_state, next_grad_amax_state,
                                                 fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
-  return fp8_out, scalar_amax(amax_buf)
+  return fp8_out

extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp

@@ -23,14 +23,14 @@ static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
 
 // fused silu*mul backward, three outputs in a single HBM pass:
 //   1) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
-//   2) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
+//   2) fp32 grad_amax_next — scalar |grad_xw13| via global atomic max
 //   3) fp32 grad_amax_out  — delayed grad amax used for quantize/GEMM epilogue scale
 // grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's
 // buffer is built in Python as a separate effect and threaded into grad_a via .after(store).
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_bwd_w13(
     __hip_fp8_storage_t*  __restrict__ grad_xw13_fp8_out,    // fp8,  2*N_ELEMS
-    float*                __restrict__ grad_amax_buf,        // fp32, NUM_WG per-WG partials
+    float*                __restrict__ grad_amax_next,       // fp32 scalar, initialized to 0 before launch
     float*                __restrict__ grad_amax_out,        // fp32 scalar delayed grad amax
     const __hip_bfloat16* __restrict__ xw13,                 // bf16, 2*N_ELEMS
     const __hip_bfloat16* __restrict__ grad_x2,              // bf16, N_ELEMS
@@ -92,5 +92,6 @@ fused_silu_mul_bwd_w13(
     if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
     __syncthreads();
   }
-  if (tid == 0) grad_amax_buf[wg] = sdata[0];
+  if (tid == 0 && sdata[0] > *grad_amax_next)
+    atomicMax(reinterpret_cast<int32_t*>(grad_amax_next), __float_as_int(sdata[0]));
 }

extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp

@@ -24,7 +24,7 @@ static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC (so VEC loads
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_cast_amax_w13(
     __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG (per-WG amaxes)
+    float*                __restrict__ amax_out,        // fp32 scalar, initialized to 0 before launch
     const __hip_bfloat16* __restrict__ xw13,            // bf16, 2*N_ELEMS
     const float*          __restrict__ amax_state)      // fp32 scalar
 {
@@ -67,7 +67,7 @@ fused_silu_mul_cast_amax_w13(
     *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
   }
 
-  // LDS tree reduction: per-workgroup amax
+  // LDS tree reduction: per-workgroup amax, then global atomic into the scalar.
   sdata[tid] = local_max;
   __syncthreads();
   for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
@@ -75,5 +75,5 @@ fused_silu_mul_cast_amax_w13(
     __syncthreads();
   }
 
-  if (tid == 0) amax_buf[wg] = sdata[0];
+  if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast<int32_t*>(amax_out), __float_as_int(sdata[0]));
 }

extra/llama_kernels/fused_ce/__init__.py

@@ -16,7 +16,7 @@ def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp
   row_lse = (logits[b, s, v_lse].cast(dtypes.float) - row_max).exp().reduce(v_lse, arg=Ops.ADD).log() + row_max
 
   v_smooth = UOp.range(vocab, 3, axis_type=AxisType.REDUCE)
-  target = logits[b, s, targets[row].cast(dtypes.index)].cast(dtypes.float)
+  target = logits[b, s, targets[row].cast(dtypes.weakint)].cast(dtypes.float)
   mean_logits = logits[b, s, v_smooth].cast(dtypes.float).reduce(v_smooth, arg=Ops.ADD) / vocab
   loss = row_lse - (1.0 - label_smoothing) * target - label_smoothing * mean_logits
   stores = UOp.group(loss_out[row].store(loss), max_out[row].store(row_max), lse_out[row].store(row_lse))
@@ -32,7 +32,7 @@ def _custom_fused_ce_loss_bwd(d_logits:UOp, logits:UOp, lse:UOp, targets:UOp, sc
   s = row % seq
 
   prob = (logits[b, s, v].cast(dtypes.float) - lse[row]).exp()
-  target = v.eq(targets[row].cast(dtypes.index)).where(1.0 - label_smoothing, 0.0)
+  target = v.eq(targets[row].cast(dtypes.weakint)).where(1.0 - label_smoothing, 0.0)
   smooth = label_smoothing / vocab
   grad = (prob - target - smooth) * scale[0]
 

extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py

@@ -3,19 +3,19 @@ import functools, pathlib
 from tinygrad import Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip
+from extra.llama_kernels import NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, dname_of, compile_hip
 
 def _src() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8.cpp").read_text()
 def _src_bwd() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8_bwd.cpp").read_text()
 
 @functools.cache
-def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
+def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp,
                 x:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
   MBS, SEQ, HIDDEN = x.shape
   n_elems = MBS * SEQ * HIDDEN
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
+  mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + 4 + 4
+  sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_out.base,
                   x.base, weight.base, amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
                                  estimates=Estimates(ops=6*n_elems, mem=mem)))
@@ -26,13 +26,13 @@ def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
                                UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
 
 @functools.cache
-def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
+def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp,
                     x:UOp, residual:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
   MBS, SEQ, HIDDEN = x.shape
   n_elems = MBS * SEQ * HIDDEN
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
+  mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + 4 + 4
+  sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_out.base,
                   x.base, residual.base, weight.base, amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_add_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
                                  estimates=Estimates(ops=7*n_elems, mem=mem)))
@@ -85,7 +85,7 @@ def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_st
   return grad_total.uop, grad_weight_uop
 
 def _fused_bwd(gradient:UOp, kernel:UOp):
-  # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state)
+  # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state)
   _, x_normed_u, rrms_u, _, x_u, weight_u, amax_state_u = kernel.src[1:]
   grad_x, grad_w = _bwd_common(gradient, None, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
   return (None, None, None, None, grad_x, grad_w, None)
@@ -112,8 +112,9 @@ def _fused_add_bwd(*args, **kwargs):
   grad_h, grad_w = _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
   return (None, None, None, None, None, grad_h, grad_h, grad_w, None)
 
-def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor]:
-  # NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, new_amax, x_normed, rrms).
+def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype,
+                                   amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor]:
+  # NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, x_normed, rrms).
   # x_normed + rrms are saved for the rmsnorm backward (also recomputed here from x regs).
   assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
   assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
@@ -123,16 +124,15 @@ def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, e
   fp8_out      = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,       x.device, axis)
   x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
   rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
-  amax_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
   fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps)
-  fp8_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
-    fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
-  return fp8_out, scalar_amax(amax_buf), x_normed_out, rrms_out
+  fp8_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
+    fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
+  return fp8_out, x_normed_out, rrms_out
 
 def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, amax_state:Tensor,
-                                       eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+                                       eps:float, fp8_dtype, amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
   # NOTE: h = x + residual; y_normed = rmsnorm(h); fp8 = quantize(y_normed * weight).
-  # Returns (fp8, new_amax, h, x_normed, rrms). h is also written so downstream can
+  # Returns (fp8, h, x_normed, rrms). h is also written so downstream can
   # reuse it without recomputing x+residual — eliminates the separate residual-add kernel.
   assert x.dtype == dtypes.bfloat16 and residual.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
   assert x.shape == residual.shape
@@ -143,9 +143,8 @@ def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor,
   h_out        = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
   x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
   rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
-  amax_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
   fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps)
-  fp8_out, h_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
-    fp8_out, h_out, x_normed_out, rrms_out, amax_buf, x, residual, weight, amax_state,
+  fp8_out, h_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
+    fp8_out, h_out, x_normed_out, rrms_out, amax_out, x, residual, weight, amax_state,
     fxn=fxn, grad_fxn=_fused_add_bwd)
-  return fp8_out, scalar_amax(amax_buf), h_out, x_normed_out, rrms_out
+  return fp8_out, h_out, x_normed_out, rrms_out

extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp

@@ -7,7 +7,7 @@
 //   fp8 = fp8_sat(y * (FP8_MAX / amax_state))
 // Also writes:
 //   rrms[row]        — saved for the rmsnorm backward
-//   amax_buf[wg]     — per-WG |y| partials, reduced later to update amax_state
+//   amax_out         — scalar |y| via global atomic max
 //
 // Layout: one WG per row, ROWS_PER_WG rows per WG via grid-stride (ROWS = N_ELEMS / HIDDEN).
 // Each thread handles HIDDEN / THREADS_PER_WG elements per row.
@@ -48,7 +48,7 @@ fused_add_rmsnorm_mul_quantize_fp8(
     __hip_bfloat16*       __restrict__ h_out,           // bf16, ROWS*HIDDEN — x + residual (saved for downstream)
     __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN
     float*                __restrict__ rrms_out,        // fp32, ROWS
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG
+    float*                __restrict__ amax_out,        // fp32 scalar, initialized to 0 before launch
     const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
     const __hip_bfloat16* __restrict__ residual,        // bf16, ROWS*HIDDEN — added into x before rmsnorm
     const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN
@@ -60,7 +60,7 @@ fused_rmsnorm_mul_quantize_fp8(
     __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, ROWS*HIDDEN
     __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN (saved for rmsnorm bwd)
     float*                __restrict__ rrms_out,        // fp32, ROWS (fp32 to match rmsnorm_bwd.cpp expectation)
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG per-WG partials
+    float*                __restrict__ amax_out,        // fp32 scalar, initialized to 0 before launch
     const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
     const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN (per-hidden scale)
     const float*          __restrict__ amax_state)      // fp32 scalar
@@ -144,12 +144,12 @@ fused_rmsnorm_mul_quantize_fp8(
     __syncthreads();  // before next row's sum_sq reduce reuses sdata
   }
 
-  // Final per-WG amax reduce.
+  // Final per-WG amax reduce, then global atomic into the scalar.
   sdata[tid] = local_max;
   __syncthreads();
   for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
     if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
     __syncthreads();
   }
-  if (tid == 0) amax_buf[wg] = sdata[0];
+  if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast<int32_t*>(amax_out), __float_as_int(sdata[0]));
 }

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -3,14 +3,13 @@ from tinygrad import Tensor, dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import prod
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax
+from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like
 
 @functools.cache
-def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp) -> UOp:
+def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:UOp, device=None) -> UOp:
   VEC = 8
   n_elems = prod(x.shape)
   assert n_elems % (NUM_WG * THREADS_PER_WG * VEC) == 0
-  assert amax_partial.shape[0] == NUM_WG
 
   x = x.reshape(n_elems)
   fp8_out = fp8_out.reshape(n_elems)
@@ -46,8 +45,13 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_st
     lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier())
     step //= 2
 
-  amax_store = amax_partial[tid.eq(0).where(wg, UOp.invalid())].store(lds[0])
-  return amax_store.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
+  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))
+  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=()))
 
 @functools.cache
 def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp) -> UOp:
@@ -69,25 +73,19 @@ def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp):
   grad_x = (Tensor(gradient, device=device).float() * scale).cast(dtypes.bfloat16)
   return (None, None, grad_x.uop, None)
 
-def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor, Tensor, UOp]:
-  # NOTE: one-pass bf16 -> fp8 quantize with delayed scaling. Returns (fp8, inv_scale, new_amax, store_effect).
-  # Fused kernel reads x once and writes fp8 + per-WG |x| partials (then a small reduce produces scalar new_amax).
-  # store_effect writes new_amax into amax_state's buffer — the caller must thread it into a realized
-  # output via `.after(store_effect)`. Calling `amax_state.assign(new_amax)` inside a grad_fxn does
-  # NOT work because .assign mutates only the temp Tensor's .uop, not the original layer-owned buffer.
+def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, amax_out:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor]:
+  # NOTE: one-pass bf16 -> fp8 quantize with delayed scaling.
+  # Fused kernel reads x once and writes fp8 + scalar amax via global atomic max.
   assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
   axis = x.uop.axis if isinstance(x.device, tuple) else None
   fp8_out      = alloc_like(x.shape,  fp8_dtype,      x.device, axis)
   n_elems = prod(x.uop.shard_shape)
   assert n_elems % NUM_WG == 0, f"{n_elems=} must divide over {NUM_WG=}"
-  amax_partial = alloc_local((NUM_WG,), dtypes.float32, x.device, axis)
-  fxn = _custom_quantize_fp8_with_amax
-  fp8_out, amax_partial, *_ = Tensor.custom_kernel(fp8_out, amax_partial, x, amax_state,
-                                                    fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
-  new_amax = scalar_amax(amax_partial)
+  fxn = functools.partial(_custom_quantize_fp8_with_amax, device=x.device)
+  fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, x, amax_state,
+                                                fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
   inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
-  store_effect = amax_state.uop.store(new_amax.uop)
-  return fp8_out, inv_scale, new_amax, store_effect
+  return fp8_out, inv_scale
 
 def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> Tensor:
   # NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.

extra/thunder/amd/fa.py

@@ -94,7 +94,7 @@ def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, h
   B_local = B // num_devices if is_dp else B
   H_local = n_heads // num_devices if is_mp else n_heads
   H_KV_local = n_kv_heads // num_devices if is_mp else n_kv_heads
-  assert (B_local, N, H_local, H_KV_local, head_dim) == (2, 8192, 32, 8, 128)
+  assert H_local % H_KV_local == 0 and head_dim % 2 == 0 and head_dim <= 512
   single_device = xqkv.device[0] if isinstance(xqkv.device, tuple) else xqkv.device
   arch = Device[single_device].renderer.target.arch
   axis = 0 if is_dp else 2 if is_mp else None

extra/torch_backend/backend.py

@@ -565,8 +565,10 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.floor_divide": lambda x,y: x//y,
   "aten.floor_divide_.Tensor": lambda x,y: x//y,
   "aten.__lshift__.Scalar": lambda x,y: x<<y,
+  "aten.__lshift__.Tensor": lambda x,y: x<<y,
   "aten.__ilshift__.Scalar": lambda x,y: x<<y,
   "aten.__rshift__.Scalar": lambda x,y: x>>y,
+  "aten.__rshift__.Tensor": lambda x,y: x>>y,
   "aten.__irshift__.Scalar": lambda x,y: x>>y,
   # inplace ops using replace for fusion
   "aten.zero_": lambda x: x.const_like(0),

pyproject.toml

@@ -74,7 +74,6 @@ testing_minimal = [
   "torch==2.9.1",
   "pytest",
   "pytest-xdist",
-  "pytest-timeout",
   "pytest-split",
   "hypothesis>=6.148.9",
   "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
@@ -160,9 +159,6 @@ norecursedirs = [
   ".hypothesis",
   ".git",
 ]
-timeout = 300
-timeout_method = "thread"
-timeout_func_only = true
 testpaths = ["test"]
 filterwarnings = [
   # Ignore SWIG warnings from importlib

test/amd/test_asm_kernel.py

@@ -36,7 +36,7 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
   assert A.dtype == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
   threads = UOp.special(A.numel(), "lidx0")
-  var = UOp.param(2, dtypes.index, vmin_vmax=(0, 10), name="var", addrspace=AddrSpace.ALU)
+  var = UOp.param(2, dtypes.weakint, vmin_vmax=(0, 10), name="var", addrspace=AddrSpace.ALU)
   insts = [
     s_load_b128(s[4:7], s[0:1]),
     s_load_b32(s[8], s[0:1], offset=0x10), # all threads load the same variable

test/backend/test_dtype_alu.py

@@ -291,6 +291,33 @@ class TestDTypeALU(unittest.TestCase):
   @Context(EMULATED_DTYPES="long")
   def test_emulated_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
 
+  def _test_shl(self):
+    for dtype, values, distances in ((dtypes.int64, [-0x1234, 0x80000001, -1, 0x1234, 1], [0, 5, 31, 32, 62]),
+                                     (dtypes.uint64, [0x80000001, 0x80000001, 1, 0xFEDC, 1], [0, 5, 31, 32, 62]),
+                                     (dtypes.int8, [-3, 1, 7, -2, 1], [0, 1, 3, 5, 6]),
+                                     (dtypes.uint16, [3, 1, 0xFF, 7, 1], [0, 1, 7, 12, 15])):
+      with self.subTest(dtype=dtype):
+        result = Tensor(values, dtype=dtype) << Tensor(distances, dtype=dtype)
+        np.testing.assert_equal(result.numpy(), [x << d for x, d in zip(values, distances)])
+
+  def _test_shr(self):
+    for dtype, values, distances in ((dtypes.int64, [-(2**40), -1, -(2**50), -(2**40), 0x123456789ABCDEF], [0, 5, 31, 32, 63]),
+                                     (dtypes.uint64, [0xFEDCBA9876543210] * 5, [0, 5, 31, 32, 63]),
+                                     (dtypes.int8, [-128, -1, 64, -37, 1], [0, 1, 3, 5, 7]),
+                                     (dtypes.uint16, [0xFFFF] * 5, [0, 1, 8, 13, 15])):
+      with self.subTest(dtype=dtype):
+        result = Tensor(values, dtype=dtype) >> Tensor(distances, dtype=dtype)
+        np.testing.assert_equal(result.numpy(), [x >> d for x, d in zip(values, distances)])
+
+  def test_shl(self): self._test_shl()
+  def test_shr(self): self._test_shr()
+
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_shl(self): self._test_shl()
+
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_shr(self): self._test_shr()
+
   @given(ht.uint8, strat.sampled_from(integer_unary_operations))
   def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op)
 

test/backend/test_linearizer.py

@@ -290,7 +290,7 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipIf(MOCKGPU and isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CUDARenderer)), "PTX indexes differently. might be ok?")
   def test_where_fold(self):
     a = Tensor.ones(4, 4).contiguous().realize()
-    b = a.shrink(((1, 2), None)).pad(((1, 2), None))
+    b = a.shrink(((1, 2), None)).pad(((1, 2), None)).bool()
     a.assign(b.where(2, a))
     linear, var_vals = a.linear_with_vars()
     assert len(linear.src) == 1

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

test/backend/test_llama_kernels.py

@@ -49,9 +49,10 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
   with Context(DEBUG=0): Tensor.realize(x, amax_state)
 
   if delayed:
-    fp8, inv_scale, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
+    amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).realize()
+    fp8, inv_scale = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
     ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state)
-    Tensor.realize(fp8, inv_scale, new_amax)
+    Tensor.realize(fp8, inv_scale)
     Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax)
   else:
     fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE)
@@ -63,9 +64,10 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
     assert fp8.cast(dtypes.float).allclose(ref_fp8.cast(dtypes.float), atol=0, rtol=0).item(), "fp8 mismatch"
     if delayed:
       assert inv_scale.allclose(ref_inv_scale, atol=0, rtol=0).item(), "inv_scale mismatch"
-      assert new_amax.allclose(ref_new_amax, atol=0, rtol=0).item(), \
-        f"amax mismatch: got={new_amax.item()} ref={ref_new_amax.item()} diff={abs(new_amax.item()-ref_new_amax.item())}"
+      assert amax_out.allclose(ref_new_amax, atol=0, rtol=0).item(), \
+        f"amax mismatch: got={amax_out.item()} ref={ref_new_amax.item()} diff={abs(amax_out.item()-ref_new_amax.item())}"
 
+@unittest.skipUnless(Device.DEFAULT == "AMD", "requires atomic max")
 class TestQuantizeFP8(unittest.TestCase):
   def setUp(self):
     ren = Device[Device.DEFAULT].renderer
@@ -81,10 +83,11 @@ class TestQuantizeFP8(unittest.TestCase):
     x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
     x = Tensor(x, device=devs)
     amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
-    fp8, _, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
-    Tensor.realize(fp8, new_amax)
+    amax_out = Tensor.zeros((), dtype=dtypes.float32, device=devs).realize()
+    fp8, _ = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
+    Tensor.realize(fp8)
     assert fp8.uop.shape == x.uop.shape
-    assert new_amax.shape == ()
+    assert amax_out.shape == ()
 
 class TestLocalAmax(unittest.TestCase):
   def test_multi_tensor_local_shard_amax(self):

test/backend/test_ops.py

@@ -5,7 +5,6 @@ import torch
 from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import NIRRenderer
 
 TINY_BACKEND = getenv("TINY_BACKEND")
@@ -450,7 +449,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,35), (45,35), (45,35)], lambda x,y,z: x.lerp(y,z))
     helper_test_op(None, lambda x,y,z: x.lerp(y,z), vals=[[1.,2.,3.], [4.,5.,6.], 0.5])
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_tril(self):
     helper_test_op([(3,3)], lambda x: x.tril())
     helper_test_op([(3,3)], lambda x: x.tril(1))
@@ -468,7 +466,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(5,3,3)], lambda x: x.tril(1))
     helper_test_op(None, lambda x: x.tril(), vals=[[[True] * 3] * 3], forward_only=True)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_triu(self):
     helper_test_op([(3,3)], lambda x: x.triu())
     helper_test_op([(3,3)], lambda x: x.triu(1))
@@ -771,6 +768,11 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: torch.tensor([2], dtype=torch.int) ** torch.tensor(-2, dtype=torch.int),
                        lambda: Tensor([2]) ** Tensor(-2), forward_only=True)
 
+  def test_pow_int_base_float_exponent(self):
+    for exponent in (0.5, 1.5, 2.0, -1.0, 0.0):
+      helper_test_op([], lambda: torch.tensor([1, 2, 3, 4], dtype=torch.int) ** exponent,
+                         lambda: Tensor([1, 2, 3, 4], dtype=dtypes.int32) ** exponent, forward_only=True)
+
   def test_sqrt(self):
     helper_test_op([(45,65)], lambda x: x.sqrt())
     helper_test_op(None, lambda x: x.sqrt(), vals=[[0.0]])
@@ -791,7 +793,6 @@ class TestOps(unittest.TestCase):
 
     self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_xor(y), expected=RuntimeError)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_and(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -809,7 +810,6 @@ class TestOps(unittest.TestCase):
 
     self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_and(y), expected=RuntimeError)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_or(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -848,6 +848,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor << 0, lambda: (ten << 0).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor << 2, lambda: (ten << 2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor << 31, lambda: (ten << 31).cast(dtypes.int32), forward_only=True)
+    helper_test_op([], lambda: tor << torch.tensor([0,2,4]).int(),
+                   lambda: (ten << Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.__lshift__(2), lambda: ten.__lshift__(2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.bitwise_left_shift(2), lambda: ten.lshift(2).cast(dtypes.int32), forward_only=True)
 
@@ -859,6 +861,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor >> 0, lambda: (ten >> 0).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor >> 2, lambda: (ten >> 2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor >> 31, lambda: (ten >> 31).cast(dtypes.int32), forward_only=True)
+    helper_test_op([], lambda: tor >> torch.tensor([0,2,4]).int(),
+                   lambda: (ten >> Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.__rshift__(2), lambda: ten.__rshift__(2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.bitwise_right_shift(2), lambda: ten.rshift(2).cast(dtypes.int32), forward_only=True)
 
@@ -870,6 +874,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor << 2, lambda: ten << 2, forward_only=True)
     helper_test_op([], lambda: tor << 8, lambda: ten << 8, forward_only=True)
     helper_test_op([], lambda: tor << 31, lambda: ten << 31, forward_only=True)
+    helper_test_op([], lambda: tor << torch.tensor([0,2,8,31]).int(), lambda: ten << Tensor([0,2,8,31], dtype=dtypes.int), forward_only=True)
 
   def test_rshift_signed(self):
     data = [[-1, -3, 1, 7], [0, -2147483648, 2147483647, -1]]
@@ -879,6 +884,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor >> 2, lambda: ten >> 2, forward_only=True)
     helper_test_op([], lambda: tor >> 8, lambda: ten >> 8, forward_only=True)
     helper_test_op([], lambda: tor >> 31, lambda: ten >> 31, forward_only=True)
+    helper_test_op([], lambda: tor >> torch.tensor([0,2,8,31]).int(), lambda: ten >> Tensor([0,2,8,31], dtype=dtypes.int), forward_only=True)
 
   def test_idiv_shift_rewrite_negative(self):
     a = Tensor(-5).div(2, rounding_mode="trunc").item()
@@ -1223,7 +1229,6 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[False, True]])
     helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[True, False]])
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_argmin(self):
     # check if it returns the first index for multiple occurrences
     helper_test_op(None, lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[2, 2]])
@@ -1529,7 +1534,6 @@ class TestOps(unittest.TestCase):
   def test_prod_dtype_arg(self):
     with self.assertRaises(AttributeError): Tensor([1.0, 2.0]).prod(dtype="")
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_min(self):
     helper_test_op([(3,3)], lambda x: x.min())
     helper_test_op([(45,3)], lambda x: x.min())
@@ -1569,7 +1573,6 @@ class TestOps(unittest.TestCase):
   def test_any_zero_axis(self):
     helper_test_op([(1,0,3,0,5)], lambda x: x.any(axis=(1,3)), forward_only=True)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_all(self):
     helper_test_op([(3,4,5,6)], lambda x: x.all(), forward_only=True)
     helper_test_op(None, lambda x: x.all(), vals=[[True, True]], forward_only=True)
@@ -2947,7 +2950,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[...,c,:,e], lambda x: x[...,k,:,p])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_dim_collapse_int(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim collapse from int
@@ -2958,7 +2960,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[1,:,3:11:2,d,0:2], lambda x: x[1,:,3:11:2,o,0:2])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_dim_inject_none(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim injection from None
@@ -2993,7 +2994,6 @@ class TestOps(unittest.TestCase):
                             lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_list_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[((0,),)])
@@ -3005,7 +3005,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,(2,1,0),c,(-2,1,0),e], lambda x: x[i,(2,1,0),k,(-2,1,0),p])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_tuple_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[(((0,),),)], lambda x: x[(((0,),),)])

test/backend/test_pickle.py

@@ -1,7 +1,7 @@
-import unittest, pickle, types
+import unittest, pickle, types, tracemalloc
 import numpy as np
-from tinygrad import Tensor, TinyJit, Variable, dtypes
-from tinygrad.helpers import GlobalCounters, ContextVar, Context
+from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
+from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
 from tinygrad.uop.ops import PatternMatcher, UPat, UOp
 
 class TestPickle(unittest.TestCase):
@@ -78,6 +78,22 @@ class TestPickle(unittest.TestCase):
     a2:UOp = pickle.loads(s)
     self.assertListEqual(a2.base.realized.as_memoryview().cast("I").tolist(), [0, 1, 2, 3])
 
+  @unittest.skipIf(DEV.interface.startswith("MOCK"), "mock device buffers live in host RAM, not VRAM")
+  def test_pickle_oob_ram(self):
+    N, M = 8, 10**6
+    ts = [Tensor.rand(M, dtype='float32').realize() for _ in range(N)]
+    tracemalloc.start()
+    st = pickle.dumps(ts, protocol=5, buffer_callback=lambda pb: pb.release())
+    self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4)
+    tracemalloc.reset_peak()
+    def make_fake_buffers():
+      for _ in range(N):
+        Device[Device.DEFAULT].synchronize()
+        yield pickle.PickleBuffer(bytearray(M*4))
+    pickle.loads(st, buffers=make_fake_buffers())
+    self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4)
+    tracemalloc.stop()
+
   def test_pickle_unrealized_tensor(self):
     t = Tensor.ones(10, 10)
     st = pickle.dumps(t)

test/backend/test_wait_loop.py

@@ -0,0 +1,104 @@
+import unittest
+from tinygrad import Tensor, UOp
+from tinygrad.device import Device
+from tinygrad.dtype import AddrSpace, dtypes
+from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.renderer.isa.x86 import X86Renderer
+from tinygrad.uop.ops import KernelInfo
+
+def wait_loop_kernel(C:UOp) -> UOp:
+  N = 10
+
+  # a RANGE with no src is a bound-less loop header: a jump target with no induction variable.
+  # the compare and conditional backedge are expanded by the renderers from the loop RANGE/END
+  l = UOp.loop(0)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+
+  # i = 0
+  i = i.after(i[0].store(0))
+
+  # i + 1, read loop-carried through after(l)
+  inc = i.after(l)[0].load() + 1
+
+  # i = inc; END(store, l, cond): conditional backedge, loop again while inc < N (do-while)
+  # NOTE: the cond uses the computed value, not a reload of the register
+  st = i[0].store(inc)
+  i = i.after(st.end(l, inc < N))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="wait_loop"))
+
+def nested_loop_kernel(C:UOp) -> UOp:
+  r = UOp.range(4, 0)
+  l = UOp.loop(1)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+  i = i.after(i[0].store(0))
+
+  inc = i.after(l, r)[0].load() + 1
+  st = i[0].store(inc)
+
+  lend = st.end(l, inc < (r.cast(dtypes.int)+1)*3)
+  i = i.after(lend.end(r))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))
+
+def two_loops_kernel(C:UOp) -> UOp:
+  # two sequential loops on the same counter: ++ until 10, then ++ until 25
+  l1, l2 = UOp.loop(0), UOp.loop(1)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+  i = i.after(i[0].store(0))
+
+  inc1 = i.after(l1)[0].load() + 1
+  i = i.after(i[0].store(inc1).end(l1, inc1 < 10))
+
+  inc2 = i.after(l2)[0].load() + 1
+  i = i.after(i[0].store(inc2).end(l2, inc2 < 25))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="two_loops", opts_to_apply=()))
+
+def loop_in_loop_kernel(C:UOp) -> UOp:
+  # outer loop while i < 12, inner loop increments until i % 4 == 0 -> 12
+  l1, l2 = UOp.loop(0), UOp.loop(1)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+  i = i.after(i[0].store(0))
+
+  inc = i.after(l1, l2)[0].load() + 1
+  st = i[0].store(inc)
+
+  # the outer END closes the inner END, and its cond reloads the register after the inner loop (in scope at the outer level)
+  e2 = st.end(l2, inc % 4 != 0)
+  oc = i.after(e2)[0].load()
+  i = i.after(e2.end(l1, oc < 12))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="loop_in_loop", opts_to_apply=()))
+
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, X86Renderer)), "loops are not supported in LVP and X86")
+class TestWaitLoop(unittest.TestCase):
+  def test_wait_loop(self):
+    c = Tensor.empty(1, dtype=dtypes.int)
+    c = Tensor.custom_kernel(c, fxn=wait_loop_kernel)[0]
+    c.realize()
+    self.assertEqual(c.item(), 10)
+
+  def test_nested_loop_in_range(self):
+    c = Tensor.empty(1, dtype=dtypes.int)
+    c = Tensor.custom_kernel(c, fxn=nested_loop_kernel)[0]
+    c.realize()
+    self.assertEqual(c.item(), 12)
+
+  def test_two_sequential_loops(self):
+    c = Tensor.empty(1, dtype=dtypes.int)
+    c = Tensor.custom_kernel(c, fxn=two_loops_kernel)[0]
+    c.realize()
+    self.assertEqual(c.item(), 25)
+
+  def test_loop_in_loop(self):
+    c = Tensor.empty(1, dtype=dtypes.int)
+    c = Tensor.custom_kernel(c, fxn=loop_in_loop_kernel)[0]
+    c.realize()
+    self.assertEqual(c.item(), 12)
+
+if __name__ == "__main__": unittest.main()

test/external/external_benchmark_op_conv.py

@@ -24,7 +24,7 @@ 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.index, Invalid))
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
   c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
@@ -50,7 +50,7 @@ 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.index, Invalid))
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
   c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))

test/external/external_test_usb_asm24.py

@@ -1,31 +1,31 @@
 import unittest
-from tinygrad.helpers import Timing
+from tinygrad.helpers import Timing, getenv
 from tinygrad import Tensor, Device
 import numpy as np
 
 class TestDevCopySpeeds(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    cls.sz = 768
+    cls.sz = getenv("SIZE", 2e6)
     cls.dev = Device["AMD"]
     if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
 
   def testCopyCPUtoDefault(self):
     for _ in range(10):
-      t = Tensor.ones(self.sz, self.sz, device="CPU").contiguous().realize()
+      t = Tensor.ones(self.sz, device="CPU", dtype='uchar').contiguous().realize()
       with Timing(f"copyin of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"): # noqa: F821
         t.to(Device.DEFAULT).realize()
         Device[Device.DEFAULT].synchronize()
       del t
 
   def testCopyDefaulttoCPU(self):
-    t = Tensor.ones(self.sz, self.sz).contiguous().realize()
+    t = Tensor.ones(self.sz, dtype='uchar').contiguous().realize()
     for _ in range(10):
       with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
         t.to('CPU').realize()
 
   def testValidateCopies(self):
-    t = Tensor.randn(self.sz, self.sz, device="CPU").contiguous().realize()
+    t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
     x = t.to(Device.DEFAULT).realize()
     Device[Device.DEFAULT].synchronize()
 

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

test/helpers.py

@@ -55,7 +55,7 @@ def assert_jit_cache_len(fxn, expected_len):
   if linear is None or not linear.src:
     assert expected_len == 0, expected_len
     return
-  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 2 # HCQ2 merges calls on the same queue
+  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 3 # HCQ2: merged same-queue calls + finalizer + bumps
   if call_is_graph(linear.src[0]):
     assert len(linear.src) == 1, len(linear.src)
     inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION

test/mockgpu/mockgpu.py

@@ -1,4 +1,4 @@
-import ctypes, time, os, builtins, fcntl
+import ctypes, time, os, builtins, fcntl, typing
 from tinygrad.helpers import DEV
 from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.autogen import libc
@@ -9,7 +9,7 @@ start = time.perf_counter()
 
 drivers = [cls() for t in DEV.value if (cls:={"MOCKPCI+AMD": AMDriver, "MOCKKFD+AMD": AMDDriver, "MOCK+AMD": AMDDriver, "MOCKUSB+AMD": AMUSBDriver,
                                               "MOCK+NV": NVDriver}.get(f"{t.interface}+{t.device}"))]
-tracked_fds = {}
+tracked_fds: dict[int, typing.Any] = {}
 
 original_memoryview = builtins.memoryview
 class TrackedMemoryView:

test/null/test_const_folding.py

@@ -40,7 +40,7 @@ class TestWeakConstFolding(unittest.TestCase):
     self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41))
 
   def test_float_unaries(self):
-    for dtype in (dtypes.weakint, dtypes.weakfloat):
+    for dtype in (dtypes.weakfloat,):
       for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
         out = UOp.const(dtype, 4).alu(op).simplify()
         self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))

test/null/test_dtype_spec.py

@@ -224,30 +224,16 @@ class TestTypePromotion(unittest.TestCase):
     assert least_upper_dtype(dtypes.fp8e5m2, dtypes.uint64) == dtypes.fp8e5m2
 
   def test_weakint_promo(self):
-    # weakint with itself is weakint
-    assert least_upper_dtype(dtypes.weakint, dtypes.weakint) == dtypes.weakint
-    # weakint is above bool
-    assert least_upper_dtype(dtypes.weakint, dtypes.bool) == dtypes.weakint
-    # weakint defers to any concrete int type
-    assert least_upper_dtype(dtypes.weakint, dtypes.int8) == dtypes.int8
-    assert least_upper_dtype(dtypes.weakint, dtypes.uint8) == dtypes.uint8
-    assert least_upper_dtype(dtypes.weakint, dtypes.int16) == dtypes.int16
-    assert least_upper_dtype(dtypes.weakint, dtypes.int32) == dtypes.int32
-    assert least_upper_dtype(dtypes.weakint, dtypes.int64) == dtypes.int64
-    assert least_upper_dtype(dtypes.weakint, dtypes.uint64) == dtypes.uint64
-    # weakint defers to any float type
-    assert least_upper_dtype(dtypes.weakint, dtypes.float16) == dtypes.float16
-    assert least_upper_dtype(dtypes.weakint, dtypes.float32) == dtypes.float32
-    assert least_upper_dtype(dtypes.weakint, dtypes.float64) == dtypes.float64
+    with self.assertRaises(KeyError): least_upper_dtype(dtypes.weakint, dtypes.weakint)
+    with self.assertRaises(KeyError): least_upper_dtype(dtypes.weakint, dtypes.int8)
 
   def test_weakfloat_promo(self):
-    # weakfloat is a float, but like weakint it is not one of dtypes.floats
+    # weakfloat is a float, but is not one of dtypes.floats
     assert dtypes.is_float(dtypes.weakfloat) and dtypes.weakfloat not in dtypes.floats
     # weakfloat with itself is weakfloat
     assert least_upper_dtype(dtypes.weakfloat, dtypes.weakfloat) == dtypes.weakfloat
-    # weakfloat is above bool, weakint and any concrete int (they defer up to it)
+    # weakfloat is above bool and any concrete int (they defer up to it)
     assert least_upper_dtype(dtypes.weakfloat, dtypes.bool) == dtypes.weakfloat
-    assert least_upper_dtype(dtypes.weakfloat, dtypes.weakint) == dtypes.weakfloat
     assert least_upper_dtype(dtypes.weakfloat, dtypes.int32) == dtypes.weakfloat
     assert least_upper_dtype(dtypes.weakfloat, dtypes.uint64) == dtypes.weakfloat
     # weakfloat defers to any concrete float type
@@ -329,6 +315,19 @@ class TestAutoCastType(unittest.TestCase):
     assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
     assert (Tensor.ones(4, 4, dtype=dt) + True).dtype == dt
 
+  @given(strat.sampled_from(core_dtypes))
+  def test_pad_scalar(self, dt):
+    t = Tensor.ones(4, dtype=dt)
+    assert t.pad(((1, 1),), value=2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert t.pad(((1, 1),), value=2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert t.pad(((1, 1),), value=True).dtype == dt
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_sort(self, dt):
+    # sort pads with dtype.min/max, a scalar of its own dtype
+    assert Tensor([3, 1, 2], dtype=dt).sort()[0].dtype == dt
+    assert Tensor([3, 1, 2], dtype=dt).sort(descending=True)[0].dtype == dt
+
   @given(strat.sampled_from(dtype_floats))
   def test_int_div_int(self, default_float):
     dtypes.default_float = default_float
@@ -429,6 +428,9 @@ class TestAutoCastType(unittest.TestCase):
     self.check_where_alternate_input_other(3.1, True, dtypes.default_float)
     self.check_where_alternate_input_other(3, 2, dtypes.default_int)
     self.check_where_alternate_input_other(3, True, dtypes.default_int)
+
+  def test_where_non_bool_cond_raises(self):
+    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0)
     self.check_where_alternate_input_other(False, True, dtypes.bool)
 
   @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))

test/null/test_gpudims.py

@@ -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.index, 0)
+    flat = UOp.const(dtypes.weakint, 0)
     for i, idx in enumerate(idxs):
       flat = flat + idx * int(math.prod(dims[i+1:]))
     flat_p = flat.substitute({s: UOp(Ops.SPECIAL, src=s.src, arg=s.arg+"_p") for s in specials})

test/null/test_graph_rewrite.py

@@ -107,21 +107,21 @@ class TestFoldingAndReduction(unittest.TestCase):
 class TestModuloAndDivisionFolding(unittest.TestCase):
   def test_full_graph_rewrite_modulo_folding_with_define_var(self):
     # index dtype because div-mod rules only work on index
-    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.index)
+    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
     optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
     self.assertEqual(optimized_mod_uop.op, Ops.CONST)
     self.assertEqual(optimized_mod_uop.arg, 2)
 
   def test_full_graph_rewrite_division_folding_with_define_var(self):
     # index dtype because div-mod rules only work on index
-    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.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)
 
   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.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)
@@ -140,7 +140,7 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
   def test_full_graph_rewrite_modulo_large_divisor(self):
     # index dtype because div-mod rules only work on index
     x_var_uop = UOp.variable('x', 1, 5)
-    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.index) % 10).render(simplify=False), x_var_uop.render(simplify=False))
+    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))
 
   def test_full_graph_rewrite_division_with_remainder(self):
     x_var_uop = UOp.variable('x', 7, 9)

test/null/test_helpers.py

@@ -1,6 +1,7 @@
 import ctypes, gzip, unittest, timeit, pickle
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, count, all_same
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, mv_address, count, all_same
+from tinygrad.tensor import is_numpy_ndarray
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.helpers import ceildiv, ansistrip, get_shape
 from tinygrad.tensor import Tensor

test/null/test_linearizer_failures.py

@@ -8,12 +8,12 @@ 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.index, 2), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
-    c3 = ((c1*UOp.const(dtypes.index, 32))+c2)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.LOOP)
+    c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
     c4 = UOp.param(1, dtypes.float, (163840,))
-    c5 = UOp.range(UOp.const(dtypes.index, 2560), 0, AxisType.REDUCE)
-    c6 = c4.index(((((((c5//UOp.const(dtypes.index, 8))%UOp.const(dtypes.index, 8))*UOp.const(dtypes.index, 8))+(c5%UOp.const(dtypes.index, 8)))+(((c2*UOp.const(dtypes.index, 40))+(c5//UOp.const(dtypes.index, 64)))*UOp.const(dtypes.index, 64)))+(c1*UOp.const(dtypes.index, 81920))))
+    c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)
+    c6 = c4.index(((((((c5//UOp.const(dtypes.weakint, 8))%UOp.const(dtypes.weakint, 8))*UOp.const(dtypes.weakint, 8))+(c5%UOp.const(dtypes.weakint, 8)))+(((c2*UOp.const(dtypes.weakint, 40))+(c5//UOp.const(dtypes.weakint, 64)))*UOp.const(dtypes.weakint, 64)))+(c1*UOp.const(dtypes.weakint, 81920))))
     c7 = UOp.param(2, dtypes.float, (64,))
     c8 = c7.index(c3)
     c9 = ((((c6+(c8*UOp.const(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()

test/null/test_llm_server.py

@@ -1,4 +1,4 @@
-import unittest, threading, time
+import unittest, threading, time, json
 from unittest.mock import Mock
 
 class TestLLMServer(unittest.TestCase):
@@ -7,12 +7,9 @@ class TestLLMServer(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     cls.mock_tok = Mock()
-    cls.mock_tok.role = Mock(return_value=[100, 101])
     cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
     cls.mock_tok.decode = Mock(return_value="Hello")
     cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "")
-    cls.mock_tok.end_turn = Mock(return_value=[998])
-    cls.mock_tok.prefix = Mock(return_value=[1])
     cls.mock_tok.preset = "llama3"
     cls.mock_tok.bos_id = 1
     cls.mock_tok.eos_id = 999
@@ -20,12 +17,14 @@ class TestLLMServer(unittest.TestCase):
     cls.mock_tok.is_end = Mock(side_effect=lambda tid: tid in (999,))
 
     cls.mock_model = Mock()
+    cls.mock_model.max_context = 4
     cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
     cls.mock_model.get_start_pos = Mock(return_value=0)
 
-    from tinygrad.llm.cli import LLMServer
+    from tinygrad.llm.cli import FallbackTemplate
+    from tinygrad.llm.serve import LLMServer
 
-    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok)
+    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok, FallbackTemplate(cls.mock_tok))
     cls.port = cls.server.server_address[1]
     cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
     cls.server_thread.start()
@@ -131,6 +130,16 @@ class TestLLMServer(unittest.TestCase):
     self.assertIsNotNone(resp.usage.prompt_tokens)
     self.assertIsNotNone(resp.usage.completion_tokens)
 
+  def test_context_length_error(self):
+    from openai import BadRequestError
+    self.mock_tok.encode.return_value = [200, 201, 202, 203]
+    try:
+      with self.assertRaises(BadRequestError) as err:
+        self.client.chat.completions.create(model="test-model", messages=[{"role":"user", "content":"too long"}])
+      self.assertEqual(err.exception.code, "context_length_exceeded")
+    finally:
+      self.mock_tok.encode.return_value = [200, 201, 202]
+
   def test_max_tokens_streaming(self):
     self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
     stream = self.client.chat.completions.create(
@@ -149,50 +158,6 @@ class TestLLMServer(unittest.TestCase):
     self.assertEqual(resp.choices[0].finish_reason, "length")
     self.assertEqual(resp.usage.completion_tokens, 2)
 
-  def test_assistant_prefill(self):
-    """Last assistant message should be treated as prefill (not a completed turn)."""
-    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
-    captured_ids = []
-    orig_generate = self.mock_model.generate.side_effect
-    def capture_generate(ids, **kwargs):
-      captured_ids.extend(ids)
-      return orig_generate(ids, **kwargs)
-    self.mock_model.generate = Mock(side_effect=capture_generate)
-
-    resp = self.client.chat.completions.create(
-      model="test", messages=[
-        {"role": "user", "content": "Hello"},
-        {"role": "assistant", "content": "Sure"}
-      ], stream=False
-    )
-    # prefill tokens should be in ids: role("assistant") + encode("Sure") but NO end_turn after it
-    # and NO extra role("assistant") appended
-    role_tokens = self.mock_tok.role.call_args_list
-    # last role() call should be for "assistant" (the prefill message), not an extra one
-    self.assertEqual(role_tokens[-1], unittest.mock.call("assistant"))
-    # end_turn should be called once less than role() — the prefill assistant msg doesn't get end_turn
-    # NOTE: this is flaky in random order
-    #self.assertEqual(self.mock_tok.end_turn.call_count, self.mock_tok.role.call_count - 1)
-    self.assertIsNotNone(resp.choices[0].message.content)
-
-  def test_assistant_prefill_not_last(self):
-    """Assistant message that's NOT last should be a normal completed turn."""
-    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
-    self.mock_tok.role.reset_mock()
-    self.mock_tok.end_turn.reset_mock()
-    self.client.chat.completions.create(
-      model="test", messages=[
-        {"role": "user", "content": "Hello"},
-        {"role": "assistant", "content": "Sure"},
-        {"role": "user", "content": "Continue"}
-      ], stream=False
-    )
-    # all messages get end_turn, plus an extra role("assistant") at the end
-    # roles: user, assistant, user, assistant(generation prompt) = 4 role calls
-    # end_turns: user, assistant, user = 3 end_turn calls (one per message)
-    self.assertEqual(self.mock_tok.end_turn.call_count, 3)
-    self.assertEqual(self.mock_tok.role.call_count, 4)
-
   def test_models_endpoint(self):
     import requests as req
     resp = req.get(f"http://127.0.0.1:{self.port}/v1/models")
@@ -203,5 +168,103 @@ class TestLLMServer(unittest.TestCase):
     self.assertEqual(data["data"][0]["id"], "test-model")
     self.assertEqual(data["data"][0]["object"], "model")
 
+class TestLLMToolCalls(unittest.TestCase):
+  """Tool calling through the OpenAI-compatible HTTP API."""
+
+  @classmethod
+  def setUpClass(cls):
+    cls.mock_tok = Mock()
+    cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
+    cls.mock_tok.decode = Mock(return_value="")
+    cls.mock_tok.preset = "qwen2"
+    cls.mock_tok.bos_id, cls.mock_tok.eos_id, cls.mock_tok.eot_id = None, 999, None
+    cls.mock_tok.is_end = Mock(return_value=False)
+
+    cls.mock_model = Mock()
+    cls.mock_model.max_context = 4
+    cls.mock_model.get_start_pos = Mock(return_value=0)
+
+    from tinygrad.llm.serve import LLMServer
+    import jinja2
+    # .items() matches tool-aware templates and ensures OpenAI JSON argument strings are normalized before rendering the next turn.
+    template = jinja2.Template("""{% for m in messages %}{{ m.content or '' }}{% for tc in m.tool_calls or [] %}
+      {% for key, value in tc.function.arguments.items() %}{{ key }}={{ value }}{% endfor %}{% endfor %}{% endfor %}""")
+    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "tool-model", cls.mock_tok, template)
+    cls.port = cls.server.server_address[1]
+    cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
+    cls.server_thread.start()
+    time.sleep(0.1)
+
+    from openai import OpenAI
+    cls.client = OpenAI(base_url=f"http://127.0.0.1:{cls.port}/v1", api_key="test")
+
+  @classmethod
+  def tearDownClass(cls):
+    cls.server.shutdown()
+    cls.server.server_close()
+
+  def set_output(self, text:str):
+    pieces = dict(enumerate(text, 1))
+    self.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: pieces[tid] if tid is not None else "")
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter(pieces))
+
+  @staticmethod
+  def tools():
+    return [{"type":"function", "function":{"name":"read", "description":"Read a file",
+      "parameters":{"type":"object", "properties":{"path":{"type":"string"}}, "required":["path"]}}}]
+
+  def test_streaming_tool_call(self):
+    self.set_output('before<tool_call>{"name":"read","arguments":{"path":"README.md"}}</tool_call>')
+    chunks = list(self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read README.md"}],
+                                                     tools=self.tools(), stream=True))
+    self.assertEqual("".join(c.choices[0].delta.content or "" for c in chunks if c.choices), "before")
+    calls = [tc for c in chunks if c.choices for tc in c.choices[0].delta.tool_calls or []]
+    self.assertEqual(len(calls), 1)
+    self.assertEqual(calls[0].function.name, "read")
+    self.assertEqual(json.loads(calls[0].function.arguments), {"path":"README.md"})
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "tool_calls")
+
+  def test_multiple_xml_tool_calls(self):
+    self.set_output("<tool_call><function=read><parameter=path>\"a\"</parameter></function></tool_call>"
+                    "<tool_call><function=read><parameter=path>\"b\"</parameter></function></tool_call>")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read a and b"}],
+                                                   tools=self.tools())
+    self.assertEqual([json.loads(tc.function.arguments)["path"] for tc in response.choices[0].message.tool_calls], ["a", "b"])
+    self.assertEqual(response.choices[0].finish_reason, "tool_calls")
+
+  def test_multiline_tool_argument_preserves_trailing_newline(self):
+    self.set_output("<tool_call>\n<function=write>\n<parameter=content>\nfirst\nsecond\n\n</parameter>\n"
+                    "<parameter=filePath>\nout.txt\n</parameter>\n</function>\n</tool_call>")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Write out.txt"}], tools=self.tools())
+    args = json.loads(response.choices[0].message.tool_calls[0].function.arguments)
+    self.assertEqual(args, {"content":"first\nsecond\n", "filePath":"out.txt"})
+
+  def test_invalid_tool_call_becomes_content(self):
+    self.set_output("<tool_call>not a call</tool_call>")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Hello"}], tools=self.tools())
+    self.assertEqual(response.choices[0].message.content, "<tool_call>not a call</tool_call>")
+    self.assertIsNone(response.choices[0].message.tool_calls)
+    self.assertEqual(response.choices[0].finish_reason, "stop")
+
+  def test_tool_call_in_reasoning_is_not_executed(self):
+    self.set_output('<think>draft <tool_call>{"name":"wrong","arguments":{}}</tool_call></think>answer')
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Hello"}], tools=self.tools())
+    self.assertEqual(response.choices[0].message.content, "answer")
+    self.assertIsNone(response.choices[0].message.tool_calls)
+    self.assertEqual(response.choices[0].finish_reason, "stop")
+
+  def test_tool_result_round_trip(self):
+    self.set_output('<tool_call>{"name":"read","arguments":{"path":"README.md"}}</tool_call>')
+    first = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read README.md"}], tools=self.tools())
+    call = first.choices[0].message.tool_calls[0]
+    self.set_output("done")
+    second = self.client.chat.completions.create(model="tool-model", messages=[
+      {"role":"user", "content":"Read README.md"},
+      {"role":"assistant", "content":None, "tool_calls":[call.model_dump()]},
+      {"role":"tool", "tool_call_id":call.id, "content":"file contents"},
+    ], tools=self.tools())
+    self.assertEqual(second.choices[0].message.content, "done")
+    self.assertEqual(second.choices[0].finish_reason, "stop")
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_llm_tokenizer.py

@@ -1,5 +1,5 @@
-import unittest, base64, functools, sys
-from tinygrad.llm.cli import SimpleTokenizer
+import unittest, base64, functools, re, sys, time, unicodedata
+from tinygrad.llm.cli import SimpleTokenizer, FallbackTemplate
 from tinygrad.helpers import fetch
 
 @unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
@@ -46,6 +46,41 @@ class TestLLMTokenizer(unittest.TestCase):
   def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
   def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n  \n", [ 31213, 14211 ])
 
+  def test_split_regex_matches_naive_listing(self):
+    # the compacted codepoint ranges must match the same text as listing every codepoint
+    def naive(pre): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
+    r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + naive("Z"), naive("N"), naive("L")
+    naive_re = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" +
+      f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
+    sample = "hello world 한국어 中文 текст ١٢٣ 123 😊\n  \ttoday\n'équivalent ²³№ "
+    self.assertEqual(SimpleTokenizer({}, {})._split_to_word.findall(sample), naive_re.findall(sample))
+
+  def test_split_regex_speed(self):
+    # the naive listing compiles a 429KB pattern that takes 10+s to match a 225KB prompt; ranges keep it small and fast
+    tok = SimpleTokenizer({}, {})
+    self.assertLess(len(tok._split_to_word.pattern), 100_000)
+    text = "The quick brown fox jumps over the lazy dog. " * 5000
+    tok._split_to_word.findall(text)  # warmup
+    tms = []
+    for _ in range(5):
+      st = time.perf_counter()
+      words = tok._split_to_word.findall(text)
+      tms.append(time.perf_counter() - st)
+    self.assertLess(min(tms), 4)  # best-of-5 is robust to CI scheduling pauses; new code takes ~60ms
+    self.assertEqual(len(words), 50001)
+
+  def test_llama_continued_conversation(self):
+    self._test_coding(self.llama_tok, "hello <|eot_id|>world", [15339, 220, 128009, 14957])
+    self._test_coding(self.llama_tok, "hello <|eot_id|>world again", [15339, 220, 128009, 14957, 1578])
+    self._test_coding(self.llama_tok, "hello changed <|eot_id|>world again", [15339, 5614, 220, 128009, 14957, 1578])
+
+  def test_long_cached_prompt_matches_fresh_tokenization(self):
+    prefix = "system tools\n" * 700 + "<|eot_id|>"
+    first, changed = prefix + "run tower of hanoi", prefix + "run ls /"
+    expected = self.llama_tok.encode(changed)
+    self.llama_tok.encode(first)
+    self.assertEqual(self.llama_tok.encode(changed), expected)
+
   def test_tekken_from_gguf_kv(self):
     kv = {
       "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],
@@ -54,10 +89,11 @@ class TestLLMTokenizer(unittest.TestCase):
       "tokenizer.ggml.eos_token_id": 2,
     }
     tok = SimpleTokenizer.from_gguf_kv(kv)
-    self.assertEqual(tok.role("user"), [3])
+    template = FallbackTemplate(tok)
+    self.assertEqual(template.role("user"), "[INST]")
     self.assertEqual(tok.encode("hello"), [5])
-    self.assertEqual(tok.end_turn(), [4])
-    self.assertEqual(tok.role("assistant"), [])
+    self.assertEqual(template.end_turn(), "[/INST]")
+    self.assertEqual(template.role("assistant"), "")
 
   def test_stream_decoder(self):
     """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""

test/null/test_schedule.py

@@ -1472,6 +1472,18 @@ class TestSchedule(unittest.TestCase):
     x.softmax().sum().backward()
     run_linear(*check_schedule(x.grad, 4))
 
+  def test_logsumexp_backward(self):
+    Tensor.manual_seed(0)
+    x = Tensor.randn(4, 12, 64, 64).realize()
+    x.logsumexp(-1).sum().backward()
+    run_linear(*check_schedule(x.grad, 3))
+
+  def test_logcumsumexp_backward(self):
+    Tensor.manual_seed(0)
+    x = Tensor.randn(4, 512).realize()
+    x.logcumsumexp(-1).sum().backward()
+    run_linear(*check_schedule(x.grad, 3))
+
   def test_scaled_dot_product_attention_fusion(self):
     x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4))
     out = Tensor.scaled_dot_product_attention(x, y, z, attn_mask=m)

test/null/test_simplify_valid_idx.py

@@ -1,6 +1,6 @@
 import unittest, itertools
 
-from tinygrad.codegen.late.coalese import indexing_simplify
+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.symbolic import simplify_valid, sym, pm_move_where_on_load
@@ -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.index, nmax),), arg=expr)
+def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(dtypes.weakint, nmax),), arg=expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
 def Range(n, nmax): return UOp.range(nmax, n)
 
@@ -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.index, 0)))
+    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
     try:
       self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0")
     except AssertionError:
@@ -474,7 +474,7 @@ class TestImageSimplification(unittest.TestCase):
     A1 = lidx0*16 + r0*16 + lidx1*4 - 51
     valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 11)
     alu0 = lidx2 + gidx0*4 + (A1 % 16)*64 + (A1 // 16 % 8)*1024
-    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(dtypes.index, 0)))
+    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
     try:
       self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
     except AssertionError:
@@ -488,18 +488,18 @@ 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.index, 0))
+    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(dtypes.weakint, 0))
     load = get_load_image_uop((1, 48, 4), valid, idx)
     self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0")
 
 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
-    from tinygrad.codegen.late.coalese import indexing_simplify
+    from tinygrad.codegen.late.coalesce import indexing_simplify
     from tinygrad.uop.ops import graph_rewrite
     from tinygrad.uop.symbolic import sym
     buf = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(dtypes.index, 0)
+    idx = UOp.const(dtypes.weakint, 0)
     true_gate = UOp.const(dtypes.bool, True)
     index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
     # apply the optimization
@@ -516,7 +516,7 @@ 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.index, 4), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 4)
@@ -524,8 +524,8 @@ class TestRangeShrink(unittest.TestCase):
   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.index, 4), r)
-    load2 = get_gated_load_uop(r < UOp.const(dtypes.index, 8), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r)
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 8)
@@ -533,7 +533,7 @@ class TestRangeShrink(unittest.TestCase):
   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.index, 300), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 300), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 204)
@@ -541,7 +541,7 @@ class TestRangeShrink(unittest.TestCase):
   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.index, 4), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assertEqual(len(ranges), 1)
@@ -550,7 +550,7 @@ class TestRangeShrink(unittest.TestCase):
   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.index, 4), r)
+    gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
     ranges = self.get_ranges(red.sink())
     self.assertEqual(len(ranges), 1)
@@ -559,7 +559,7 @@ class TestRangeShrink(unittest.TestCase):
   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.index, 1), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 1), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 0)
 

test/null/test_tensor_uop_mixin.py

@@ -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.index, 1).stack(UOp.const(dtypes.index, 2)).shape, (2,))
+    self.assertEqual(UOp.const(dtypes.weakint, 1).stack(UOp.const(dtypes.weakint, 2)).shape, (2,))
 
 class TestTensorUOpConv2d(unittest.TestCase):
   def test_conv2d_basic(self):

test/null/test_uop_graph.py

@@ -2,7 +2,7 @@ import unittest, pytest
 from tinygrad import dtypes, Variable
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes
 from tinygrad.uop.symbolic import sym
 from test.helpers import to_uops_list
 
@@ -202,7 +202,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_same_fold(self):
     v = UOp.variable('tmp', 0, 1)
-    c0 = UOp.const(dtypes.index, 0)
+    c0 = UOp.const(dtypes.weakint, 0)
     vc = v != c0
     c1 = UOp.const(dtypes.float, 1.0)
     out = vc.where(c1, c1)
@@ -424,16 +424,16 @@ class TestUOpGraph(unittest.TestCase):
     # 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.index, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)
-    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
     c7 = UOp.param(2, dtypes.uchar, (60000,))
     c8 = c7.index(c6)
     c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -441,19 +441,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.index, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)  # c4 is a load
-    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
     c7 = UOp.param(2, dtypes.uchar, (60000,))
     c8 = c7.index(c6)
     # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
-    loaded_idx = c4.cast(dtypes.index)
-    comparison = (loaded_idx + c5) < UOp.const(dtypes.index, 60000)
+    loaded_idx = c4.cast(dtypes.weakint)
+    comparison = (loaded_idx + c5) < UOp.const(dtypes.weakint, 60000)
     c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -707,5 +707,16 @@ class TestUOpBroadcast(unittest.TestCase):
     c = a + b
     self.assertEqual(c.op, Ops.ADD)
 
+  def test_broadcast_axes(self):
+    t = Variable("t", 1, 10)
+    self.assertEqual(broadcast_axes((4, 8), (4, 8)), ())
+    self.assertEqual(broadcast_axes((8,), (4, 8)), (0,))
+    self.assertEqual(broadcast_axes((), (4, 8)), (0, 1))
+    self.assertEqual(broadcast_axes((3, 1), (4, 3, 8)), (0, 2))
+    self.assertEqual(broadcast_axes((1, 8), (1, 8)), ())
+    self.assertEqual(broadcast_axes((t, 8), (t, 8)), ())
+    self.assertEqual(broadcast_axes((1, 8), (t, 8)), (0,))
+    with self.assertRaises(RuntimeError): broadcast_axes((4, 8), (8,))
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/null/test_uop_symbolic.py

@@ -11,13 +11,13 @@ from tinygrad.uop.validate import uops_to_z3
 def check_uop_against_string(self, v:UOp, s:str):
   sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.RANGE, Ops.SPECIAL, Ops.PARAM)}
   s_eval = eval(s, sym_vars)
-  if isinstance(s_eval, int) and v.dtype==dtypes.index: s_eval = UOp.const(dtypes.index, s_eval)
+  if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(dtypes.weakint, s_eval)
   elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(dtypes.from_py(s_eval), s_eval)
   s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
   self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")
 
-def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.index): return UOp.variable(name,min_val,max_val,dtype)
-def uconst(val): return UOp.const(dtypes.index, val)
+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 usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)
 
@@ -247,12 +247,12 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (0, 9))
 
   def test_range_div_its_symbolic_bound(self):
-    a = Variable("a", 1, 10, dtypes.index)
+    a = Variable("a", 1, 10, dtypes.weakint)
     ridx0 = UOp.range(a+2, 0)
     self.helper_test_variable(ridx0//(a+2), 0, 0, "0")
 
   def test_range_mod_its_symbolic_bound(self):
-    a = Variable("a", 1, 10, dtypes.index)
+    a = Variable("a", 1, 10, dtypes.weakint)
     ridx = UOp.range(a+2, 0)
     self.helper_test_variable(ridx%(a+2), 0, 11, "r0")
 
@@ -919,8 +919,8 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond.cast(dtypes.int).ne(2), 1, 1, "True")
     self.helper_test_variable(cond.cast(dtypes.int).ne(-1), 1, 1, "True")
     # CAST(bool -> index) folds too
-    self.helper_test_variable(cond.cast(dtypes.index).ne(0), 0, 1, "(a<2)")
-    self.helper_test_variable(cond.cast(dtypes.index).ne(1), 0, 1, "((a<2)!=True)")
+    self.helper_test_variable(cond.cast(dtypes.weakint).ne(0), 0, 1, "(a<2)")
+    self.helper_test_variable(cond.cast(dtypes.weakint).ne(1), 0, 1, "((a<2)!=True)")
 
   def test_where_removal(self):
     cond = Variable("a", 0, 3) < 2
@@ -1021,7 +1021,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((numerator//denominator)<=0, 1, 1, "True")
 
   def test_symbolic_range_doesnt_collapse(self):
-    r0 = UOp.range((Variable("a", 1, 10)<5).cast(dtypes.index), 0)
+    r0 = UOp.range((Variable("a", 1, 10)<5).cast(dtypes.weakint), 0)
     self.helper_test_variable(r0, 0, 0, "r0")
 
   def test_const_reciprocal(self):
@@ -1289,16 +1289,16 @@ 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.index, (1, 1, Invalid, Invalid))
-    c2 = UOp.const(dtypes.index, (1, Invalid, 1, 1))
-    self.assertIs((c1+c2).simplify(), UOp.const(dtypes.index, (2, Invalid, Invalid, Invalid)))
+    c1 = UOp.const(dtypes.weakint, (1, 1, Invalid, Invalid))
+    c2 = UOp.const(dtypes.weakint, (1, Invalid, 1, 1))
+    self.assertIs((c1+c2).simplify(), UOp.const(dtypes.weakint, (2, Invalid, Invalid, Invalid)))
 
 class TestStoreLoadFolding(unittest.TestCase):
   """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
   def test_store_load_folding(self):
     # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
     buf = UOp.param(0, dtypes.int, (1,))
-    index = buf.index(UOp.const(dtypes.index, 0))
+    index = buf.index(UOp.const(dtypes.weakint, 0))
     # Direct: store(idx, load(idx)) -> NOOP
     self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
     # Emergent: store(idx, load(idx) + 0) -> store(idx, load(idx)) -> NOOP

test/null/test_uop_vmin_vmax.py

@@ -127,6 +127,13 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertEqual(x.vmin, 0)
     self.assertEqual(x.vmax, 10 >> 2)
 
+  def test_vmin_vmax_cast_unsigned(self):
+    # a fitting source keeps exact bounds: no wrap can occur
+    self.assertEqual(UOp.variable('x', 5, 10).cast(dtypes.uint8)._min_max, (5, 10))
+    # a possibly-negative or too-large source can wrap: conservative
+    self.assertEqual(UOp.variable('x', -1, 10).cast(dtypes.uint8)._min_max, (0, 255))
+    self.assertEqual(UOp.variable('x', 250, 260).cast(dtypes.uint8)._min_max, (0, 255))
+
   def test_vmin_vmax_xor_neg1(self):
     x = UOp.variable('x', 3, 7)
     uop = x ^ -1
@@ -160,7 +167,7 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertNotEqual(i.vmin, i.vmax)
 
   def test_vmin_vmax_invalid_vconst(self):
-    x = UOp.const(dtypes.index, (0, 4, Invalid, Invalid))
+    x = UOp.const(dtypes.weakint, (0, 4, Invalid, Invalid))
     self.assertLess(x.vmin, 0)
     self.assertGreater(x.vmax, 4)
 
@@ -364,6 +371,10 @@ class TestConstFactor(unittest.TestCase):
     uop = (x * 3) * 5
     self.assertEqual(uop.const_factor(), 15)  # Constant multipliers are combined (3 * 5 = 15)
 
+  def test_const_factor_variable_multiple_of(self):
+    x = UOp.variable('x', 16, 32, multiple_of=4)
+    self.assertEqual(x.const_factor(), 4)
+
 class TestDivides(unittest.TestCase):
   def test_divides_constant_exact(self):
     # Divides a constant by an exact divisor
@@ -402,5 +413,15 @@ class TestDivides(unittest.TestCase):
     result = uop.divides(3)
     self.assertIsNone(result)  # Cannot divide by 3, since 4 is not divisible by 3
 
+  def test_divides_variable_multiple_of_exact(self):
+    x = UOp.variable('x', 16, 32, multiple_of=4)
+    result = x.divides(4)
+    self.assertIsNotNone(result)
+
+  def test_divides_variable_multiple_of_factor(self):
+    x = UOp.variable('x', 16, 32, multiple_of=4)
+    result = x.divides(2)
+    self.assertIsNotNone(result)
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_uops.py

@@ -14,28 +14,26 @@ 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(dtypes.weakint, 1), UOp.const(dtypes.int8, 1)), None), dtypes.int8)
+    with self.assertRaises(KeyError): dtype_from_uop(Ops.ADD, (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.int8, 1)), None)
 
   def test_same_dtype_fast_path(self):
-    src = (UOp.const(dtypes.index, 1), UOp.const(dtypes.index, 2))
-    self.assertEqual(dtype_from_uop(Ops.ADD, src, None), dtypes.index)
+    src = (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.weakint, 2))
+    self.assertEqual(dtype_from_uop(Ops.ADD, src, None), dtypes.weakint)
 
   def test_where_promotion(self):
     cond = UOp.const(dtypes.bool, True)
     self.assertEqual(dtype_from_uop(Ops.WHERE, (cond, UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
     idx = UOp.range(4, 0)
-    self.assertEqual(idx.valid(idx < 4).dtype, dtypes.index)
+    self.assertEqual(idx.valid(idx < 4).dtype, dtypes.weakint)
 
   def test_const_dtype_from_value(self):
     self.assertEqual(dtype_from_uop(Ops.CONST, (), True), dtypes.bool)
-    self.assertEqual(dtype_from_uop(Ops.CONST, (), 3), dtypes.weakint)
     self.assertEqual(dtype_from_uop(Ops.CONST, (), ConstFloat(3.0)), dtypes.weakfloat)
     self.assertEqual(dtype_from_uop(Ops.CONST, (), Invalid), dtypes.bool)
     self.assertRaises(TypeError, dtype_from_uop, Ops.CONST, (), (1, 2))
 
   @Context(SPEC=2)
   def test_const_default_dtype_is_derived(self):
-    self.assertEqual(UOp(Ops.CONST, arg=3).dtype, dtypes.weakint)
     self.assertEqual(UOp(Ops.CONST, arg=ConstFloat(3.0)).dtype, dtypes.weakfloat)
     self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
     self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
@@ -81,7 +79,7 @@ class TestExecALU(unittest.TestCase):
     # Invalid poisons any binary op regardless of result dtype: a comparison must not fold to a boolean
     self.assertIs(exec_alu(Ops.CMPLT, dtypes.bool, (Invalid, 1)), Invalid)
     self.assertIs(exec_alu(Ops.CMPNE, dtypes.bool, (Invalid, 1)), Invalid)
-    self.assertIs(exec_alu(Ops.ADD, dtypes.index, (Invalid, 1)), Invalid)
+    self.assertIs(exec_alu(Ops.ADD, dtypes.weakint, (Invalid, 1)), Invalid)
 
   def test_div(self):
     self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (8, 2)), 4)
@@ -155,8 +153,8 @@ 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.index, 1)
-    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(dtypes.index, 2)).valid(gate)))
+    gate = gidx0<UOp.const(dtypes.weakint, 1)
+    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(dtypes.weakint, 2)).valid(gate)))
     val = UOp.const(dtypes.float, 42.0)
     store = UOp(Ops.STORE, src=(idx, val))
     uops = to_uops_list([store])
@@ -172,8 +170,8 @@ 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.index, 2)
-    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(dtypes.index, 1))))
+    idx = gidx0 * UOp.const(dtypes.weakint, 2)
+    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(dtypes.weakint, 1))))
     idx1 = UOp(Ops.INDEX, src=(gmem1, idx))
     val = UOp.const(dtypes.float, 42.0)
     stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
@@ -192,8 +190,8 @@ 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.index, 2)
-    gate = gidx0<UOp.const(dtypes.index, 1)
+    idx = gidx0*UOp.const(dtypes.weakint, 2)
+    gate = gidx0<UOp.const(dtypes.weakint, 1)
     idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gate)))
     idx1 = UOp(Ops.INDEX, src=(gmem1, idx.valid(gate)))
     val = UOp.const(dtypes.float, 42.0)

test/null/test_validate_oob.py

@@ -126,7 +126,7 @@ class TestValidateOOB(unittest.TestCase):
       buf0 = UOp.param(0, dtypes.int, (16,))
       buf1 = UOp.param(1, dtypes.int, (64,))
       r = UOp.range(42, 0, AxisType.GLOBAL)
-      ld0 = buf0.index(r.valid(r < 8)).load(dtype=dtypes.int).cast(dtypes.index)
+      ld0 = buf0.index(r.valid(r < 8)).load(dtype=dtypes.int).cast(dtypes.weakint)
       to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load(dtype=dtypes.int)])  # valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])  # oob
@@ -135,7 +135,7 @@ class TestValidateOOB(unittest.TestCase):
     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.index, 16),), arg="gidx0")
+      gidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.weakint, 16),), arg="gidx0")
       ld_bool = buf_bool.index(gidx).load()
       with self.assertRaises(RuntimeError):
         to_uops_list([buf_int.index(gidx.valid(ld_bool)).load()])  # gidx 0..15, buf_int size 8

test/opt/test_gen_float4.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Device, Tensor, dtypes
+from tinygrad import Device, Tensor, Variable, dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
@@ -114,5 +114,31 @@ class TestFloat4(unittest.TestCase):
 
     assert TestFloat4.count_float4(uops) == (1, 1)
 
+  def test_float4_aligned_variable(self):
+    x = Variable('x', 0, 4, multiple_of=4).bind(4)
+    a = Tensor.empty(4).realize()
+    b = Tensor.empty(12).realize().shrink(((x, x+4),))
+    c = a + b
+
+    # should float4 both
+
+    s = c.linear_with_vars()[0].src[0]
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+
+    assert TestFloat4.count_float4(uops) == (2, 1)
+
+  def test_float4_unaligned_variable(self):
+    x = Variable('x', 0, 4, multiple_of=2).bind(4)
+    a = Tensor.empty(4).realize()
+    b = Tensor.empty(12).realize().shrink(((x, x+4),))
+    c = a + b
+
+    # should float4 a but not b
+
+    s = c.linear_with_vars()[0].src[0]
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+
+    assert TestFloat4.count_float4(uops) == (1, 1)
+
 if __name__ == '__main__':
   unittest.main()

test/opt/test_tensor_cores.py

@@ -18,6 +18,9 @@ from test.backend.test_linearizer import helper_realized_ast, helper_linearizer_
 
 # NOTE: to_program always passes in Device[Device.DEFAULT].renderer explicitly for process_replay!!!
 
+def _tc_rand(*shape, dtype:DType) -> Tensor:
+  return Tensor.randint(*shape, low=dtype.min, high=dtype.max+1, dtype=dtype) if dtypes.is_int(dtype) else Tensor.rand(*shape, dtype=dtype)
+
 def run_program(prg:UOp, bufs:list[Buffer]):
   buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]
   for u,b in zip(buf_uops, bufs): buffers[u] = b
@@ -25,7 +28,7 @@ def run_program(prg:UOp, bufs:list[Buffer]):
 
 def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0,
                                          ensure_triggered:bool=True):
-  a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
+  a, b = _tc_rand(M, K, dtype=dtype_in), _tc_rand(K, N, dtype=dtype_in)
   r = a.matmul(b, dtype=dtype_out)
   sched = r.schedule_linear()
   realized_ast = sched.src[-1].src[0]
@@ -44,7 +47,7 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d
     except KernelOptError: pass
 
 def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0, use_tensor_cores:int=1):
-  a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
+  a, b = _tc_rand(M, K, dtype=dtype_in), _tc_rand(K, N, dtype=dtype_in)
   np_a, np_b = a.numpy(), b.numpy()
   r = a.matmul(b, dtype=dtype_out)
   if dtype_in == dtypes.bfloat16: r = r.float()
@@ -57,9 +60,11 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
   run_program(ast, bufs)
   if dtype_in == dtypes.half: tc_atol, tc_rtol = 1e-2, 1e-3
   elif dtype_in == dtypes.bfloat16: tc_atol, tc_rtol = (1e-1, 2e-2) if dtype_out == dtypes.bfloat16 else (1e-2, 1e-2)
+  elif not dtypes.is_float(dtype_in): tc_atol, tc_rtol = 0, 0
   else: tc_atol, tc_rtol = 5e-3, 1e-4
   c = bufs[0].numpy().reshape((M,N))
-  np.testing.assert_allclose(c, np_a @ np_b, atol=tc_atol, rtol=tc_rtol)
+  ref = (np_a.astype(np.int32) @ np_b.astype(np.int32)) if not dtypes.is_float(dtype_in) else (np_a @ np_b)
+  np.testing.assert_allclose(c, ref, atol=tc_atol, rtol=tc_rtol)
 
 class TestTensorCores(unittest.TestCase):
   # TODO: don't skip bf16 for real device (METAL, AMD)
@@ -75,7 +80,7 @@ class TestTensorCores(unittest.TestCase):
   def test_tensor_cores_codegen(self):
     for tc in Device[Device.DEFAULT].renderer.tensor_cores:
       n, m, k = tc.dims
-      a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
+      a, b = _tc_rand(m, k, dtype=tc.dtype_in), _tc_rand(k, n, dtype=tc.dtype_in)
       r = a.matmul(b, dtype=tc.dtype_out)
       prg = to_program(replace_opts(r.schedule_linear().src[-1].src[0],
                         [Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))]), Device[Device.DEFAULT].renderer)

test/unit/test_dtype_weak.py

@@ -1,10 +1,7 @@
-import unittest
-from unittest.mock import patch
+import tempfile, unittest
 
 from tinygrad import Tensor, dtypes
-from tinygrad.uop import Ops
 from tinygrad.uop.ops import UOp
-from tinygrad.uop.spec import spec_tensor
 
 
 class TestWeakPromotion(unittest.TestCase):
@@ -14,28 +11,20 @@ class TestWeakPromotion(unittest.TestCase):
     with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).randn_like()
 
   def test_sum_stays_weak(self):
-    for weak, value in ((dtypes.weakint, 1), (dtypes.weakfloat, 1.0)):
+    for weak, value in ((dtypes.weakfloat, 1.0),):
       self.assertEqual(Tensor.const(weak, value).expand(3).sum().dtype, weak)
     self.assertEqual((Tensor.const(dtypes.weakfloat, 1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float16)
 
-  def test_storage_width(self):
-    t = Tensor.const(dtypes.weakint, 2)
-    for fn in (lambda: t.bitcast(dtypes.int32), lambda: Tensor.const(dtypes.int32, 2).bitcast(dtypes.weakint), t.element_size, t.nbytes):
-      with self.assertRaises(RuntimeError): fn()
-
   def test_materialize_at_default_dtype(self):
-    for weak, value, strong in ((dtypes.weakint, 3, dtypes.default_int), (dtypes.weakfloat, 0.5, dtypes.default_float)):
+    for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),):
       t = Tensor.const(weak, value)
       self.assertEqual(t.dtype, weak)
       self.assertEqual(t.data().itemsize, strong.itemsize)
       self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize)
-      realized = t.clone("CPU").realize()
-      self.assertEqual((realized.dtype, realized.uop.buffer.dtype), (strong, strong))
-    with patch.object(dtypes, "default_int", dtypes.int64):
-      self.assertEqual(Tensor.const(dtypes.weakint, 3).numpy().dtype.itemsize, dtypes.int64.itemsize)
+      with self.assertRaises(RuntimeError): t.clone("CPU")
 
   def test_uop_scalar_const_unchanged(self):
-    for dtype, value in ((dtypes.index, 1), (dtypes.int32, 1), (dtypes.float32, 0.5)):
+    for dtype, value in ((dtypes.weakint, 1), (dtypes.int32, 1), (dtypes.float32, 0.5)):
       out = UOp.variable("x", 0.0 if dtype == dtypes.float32 else 0, 10.0 if dtype == dtypes.float32 else 10, dtype) + value
       self.assertEqual((out.dtype, out.src[1].dtype), (dtype, dtype))
 
@@ -49,7 +38,6 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertEqual(((t_bool + 1) + t_i8).dtype, dtypes.int8)
     self.assertEqual(((t_bool + 1) + t_u16).dtype, dtypes.uint16)
     self.assertEqual((Tensor(3) + t_i8).dtype, dtypes.int8)
-    self.assertEqual(Tensor([2], dtype=dtypes.uint8).pad(((1, 1),), value=1).dtype, dtypes.uint8)
     # zeros/ones are full with a python fill value, so they are weak too (jnp.zeros pins float32; deliberate divergence)
     self.assertEqual((Tensor.zeros(3) + t_f16).dtype, dtypes.float16)
 
@@ -58,23 +46,13 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertEqual((t_i8 + 1).dtype, dtypes.int8)
     self.assertEqual((t_f16 + 0.5).dtype, dtypes.float16)
     self.assertEqual((t_f32 + t_f16).dtype, dtypes.float32)
+    self.assertEqual(Tensor([2], dtype=dtypes.uint8).pad(((1, 1),), value=1).dtype, dtypes.uint8)
 
   @unittest.expectedFailure  # TODO: dot of a weak const tensor defers to the other operand once python scalars are weak consts
   def test_dot_defers_weak(self):
     weak = Tensor([True, False]).where(Tensor(1), 2)
     self.assertEqual(weak.dot(Tensor([1, 1], dtype=dtypes.int8)).dtype, dtypes.int8)
 
-  @unittest.expectedFailure  # TODO: Tensor(3).uop becomes CONST(weakint); Tensor.dtype is always uop.dtype; buffers lower to the default
-  def test_dtype_is_uop_dtype(self):
-    for value, weak, lowered in ((3, dtypes.weakint, dtypes.default_int), (0.5, dtypes.weakfloat, dtypes.default_float)):
-      t = Tensor(value)
-      self.assertEqual((t.uop.dtype, t.dtype), (weak, weak))
-      self.assertEqual(t.numpy().dtype.itemsize, lowered.itemsize)
-      realized = t.clone("CPU").realize()
-      self.assertEqual((realized.dtype, realized.uop.buffer.dtype), (lowered, lowered))
-    with patch.object(dtypes, "default_int", dtypes.int64):
-      self.assertEqual(Tensor(3).clone("CPU").realize().uop.buffer.dtype, dtypes.int64)
-
   def test_integer_values(self):
     x = Tensor.full((1,), 1, dtype=dtypes.int64, device="CPU")
     self.assertEqual((x + 2**40).item(), 2**40 + 1)
@@ -94,15 +72,6 @@ class TestWeakPromotion(unittest.TestCase):
     for out in (Tensor(2).exp(), Tensor(2).cos(), Tensor(2).sigmoid()):
       self.assertEqual((out.dtype, (out + t_f16).dtype), (dtypes.weakfloat, dtypes.float16))
 
-  @unittest.expectedFailure  # TODO: where of weak consts stays weak and resolves per consumer
-  def test_where_and_shared_literal(self):
-    gate, weak = Tensor([True, False], device="CPU"), Tensor(2)
-    weak_where = gate.where(weak, 3)
-    self.assertEqual(weak_where.dtype, dtypes.weakint)
-    self.assertEqual((weak_where + Tensor([1, 1], dtype=dtypes.int64, device="CPU")).tolist(), [3, 4])
-    self.assertEqual((weak + Tensor([1], dtype=dtypes.int32, device="CPU")).item(), 3)
-    self.assertEqual((weak + Tensor([1], dtype=dtypes.int64, device="CPU")).item(), 3)
-
   def test_null_lowering(self):
     for t in (Tensor.full((1,), 1, dtype=dtypes.int64, device="NULL") + 2**40,
               Tensor.full((1,), 1.0, dtype=dtypes.float64, device="NULL") + (1.0 + 2**-40)):
@@ -110,14 +79,47 @@ class TestWeakPromotion(unittest.TestCase):
       self.assertNotIn(t.uop.buffer.dtype, dtypes.weaks)
 
 
-class TestWeakSpec(unittest.TestCase):
-  def test_weak_operand_allowed(self):
-    x = UOp.variable("x", 0, 10, dtypes.int64)
-    weak = UOp.const(dtypes.weakint, 3)
-    for u in (x.alu(Ops.ADD, weak), x.alu(Ops.CMPLT, weak), x.alu(Ops.SHL, weak)):
-      self.assertIs(spec_tensor.rewrite(u), True)
-    gate = UOp.variable("gate", False, True, dtypes.bool)
-    self.assertIs(spec_tensor.rewrite(UOp(Ops.WHERE, dtypes.int8, (gate, UOp.const(dtypes.int8, 1), weak))), True)
+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)
+    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
+    fdst = Tensor.zeros(2, dtype=dtypes.float32, device="CPU").contiguous().realize()
+    fdst[0:1] = w05                                                                    # weakfloat defers to float
+    self.assertEqual(fdst.tolist(), [0.5, 0.0])
+    with tempfile.TemporaryDirectory() as td:                                          # the DISK path checks the same
+      ddst = Tensor.empty(2, dtype=dtypes.int32, device=f"DISK:{td}/t")
+      with self.assertRaises(RuntimeError): ddst.assign(w05.expand(2))
+
+  def test_weak_has_no_storage(self):
+    import numpy as np
+    with self.assertRaises(RuntimeError): Tensor(np.ones(2, dtype=np.float32), dtype=dtypes.weakfloat)
+    with self.assertRaises(RuntimeError): Tensor(bytes(8), dtype=dtypes.weakfloat)
+
+class TestWeakMaterializationEntries(unittest.TestCase):
+  # everything that creates storage from a weak value raises
+  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))
+      self.assertEqual(weak_val().dtype, weak)
+      self.assertEqual(weak_val().to("CPU").dtype, weak)
+      self.assertEqual(weak_val().data().format, strong.fmt)
+      self.assertEqual(weak_val().numpy().dtype.itemsize, strong.itemsize)
+      self.assertEqual(weak_val().tolist(), [value])
+      self.assertEqual(weak_val().cast(strong).realize().uop.buffer.dtype, strong)
+      for entry in (lambda t: t.contiguous(), lambda t: t.realize(), lambda t: t.clone(),
+                    lambda t: t.to("CPU:1").realize(), lambda t: t.as_param(0)):
+        with self.assertRaises(RuntimeError): entry(weak_val())
+
+  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),))
+      self.assertEqual(empty.data().format, strong.fmt)
+      self.assertEqual(empty.numpy().dtype.itemsize, strong.itemsize)
+      self.assertEqual(empty.tolist(), [])
 
 
 if __name__ == "__main__":

test/unit/test_gradient.py

@@ -1,8 +1,8 @@
-import unittest
+import unittest, math
 import numpy as np
 from tinygrad import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UOp, KernelInfo
+from tinygrad.uop.ops import UOp, KernelInfo, Ops
 
 class TestTensorGradient(unittest.TestCase):
   def test_example(self):
@@ -98,6 +98,33 @@ class TestTensorGradient(unittest.TestCase):
     x = Tensor.randn(4, 4)
     np.testing.assert_allclose(x.pad(((1,0),(0,0))).gradient(x, gradient=g2)[0].numpy(), np.zeros((4, 4)))
 
+  def test_implicit_broadcast_where_gradient(self):
+    # WHERE with a bare ()-shape branch: the scalar's gradient counts the positions where it is selected
+    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0)
+    dw = Tensor(cond.uop.alu(Ops.WHERE, x.uop, w.uop)).sum().gradient(w)[0]
+    self.assertEqual(dw.shape, ())
+    self.assertEqual(dw.item(), 1.0)
+    dw = Tensor(cond.uop.alu(Ops.WHERE, w.uop, x.uop)).sum().gradient(w)[0]
+    self.assertEqual(dw.item(), 2.0)
+
+  def test_implicit_broadcast_alu_gradient(self):
+    # MUL with a bare ()-shape src, no EXPAND in the graph
+    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0)
+    m = x.uop.alu(Ops.MUL, w.uop)
+    self.assertIs(m.src[1], w.uop)
+    dw = Tensor(m).sum().gradient(w)[0]
+    self.assertEqual(dw.shape, ())
+    self.assertEqual(dw.item(), 6.0)
+
+  def test_implicit_broadcast_intermediate_accumulation(self):
+    # s is used directly and through an implicit broadcast edge, each edge's gradient reduces to s's shape before they sum
+    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5)
+    s = p.sin()
+    z = Tensor(x.uop.alu(Ops.MUL, s.uop)).sum() + s
+    dp = z.gradient(p)[0]
+    self.assertEqual(dp.shape, ())
+    self.assertAlmostEqual(dp.item(), 7*math.cos(0.5), places=5)
+
   def test_bare_const_skipped_by_backward(self):
     Tensor.manual_seed(0)
     w = Tensor(1.0)

test/unit/test_hashing.py

@@ -71,7 +71,7 @@ class TestKeccak(unittest.TestCase):
   def test_variable_bs(self):
     data = Tensor([b"abc", b"abc", b"def"], dtype=dtypes.uint8).repeat(2048, 1)
     bs = UOp.variable("bs", 1, 4096).bind(3)
-    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(3, 32)
+    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(3, 32).realize()
     self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
     self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
     self.assertEqual(bytes(out[2].tolist()), bytearray.fromhex("8e0d8f672252acb0 ffc5093db8653b18 1513bf9a2097e737 b4f73533dcaf46df"))

test/unit/test_helpers.py

@@ -1,7 +1,7 @@
 import unittest, gc
 import numpy as np
-from tinygrad.helpers import polyN, is_numpy_ndarray, disable_gc
-from tinygrad.tensor import Tensor
+from tinygrad.helpers import polyN, disable_gc
+from tinygrad.tensor import Tensor, is_numpy_ndarray
 
 class TestPolyN(unittest.TestCase):
   def test_tensor(self):

test/unit/test_invalid_tensor.py

@@ -3,6 +3,7 @@ from tinygrad import Tensor
 from tinygrad.device import Buffer
 from tinygrad.dtype import Invalid, dtypes
 from tinygrad.engine.realize import run_linear
+from tinygrad.uop.ops import Ops, UOp
 
 class TestInvalidTensor(unittest.TestCase):
   def _invalid_test_helper(self, out, expected):
@@ -132,5 +133,15 @@ class TestInvalidTensor(unittest.TestCase):
     out = Tensor([1.0, 2.0, 3.0, 4.0])[idx]
     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)))
+    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)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_linalg.py

@@ -17,6 +17,7 @@ class TestLinAlg(unittest.TestCase):
     for size in sizes:
       a = Tensor.randn(size).realize()
       U,S,V = a.svd()
+      Tensor.realize(U,S,V)
       b_shape,m,n = size[0:-2],size[-2],size[-1]
       k = min(m,n)
       s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
@@ -29,6 +30,7 @@ class TestLinAlg(unittest.TestCase):
     with Context(CHECK_OOB=0):  # sometimes this is slow in CI
       a = Tensor.randn(size).realize()
       U,S,V = a.svd(full_matrices=False)
+      Tensor.realize(U,S,V)
       b_shape,m,n = size[0:-2],size[-2],size[-1]
       k = min(m,n)
       s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
@@ -61,6 +63,7 @@ class TestLinAlg(unittest.TestCase):
     for size in sizes:
       a = Tensor.randn(size).realize()
       Q,R = a.qr()
+      Tensor.realize(Q,R)
       orthogonality_helper(Q)
       reconstruction_helper([Q,R],a)
 
@@ -73,9 +76,10 @@ class TestLinAlg(unittest.TestCase):
     reconstruction_helper([Q,R], a)
 
   def test_svd_identity(self):
-    for a in (Tensor.eye(2), Tensor.zeros(2, 2)):
+    for a in (Tensor.eye(2).clone(), Tensor.zeros(2, 2)):
       a = a.realize()
       U,S,V = a.svd()
+      Tensor.realize(U,S,V)
       assert not np.isnan(U.numpy()).any()
       assert not np.isnan(S.numpy()).any()
       assert not np.isnan(V.numpy()).any()
@@ -85,6 +89,7 @@ class TestLinAlg(unittest.TestCase):
   def test_svd_identity_4x4(self):
     a = Tensor.eye(4).clone()
     U,S,V = a.svd()
+    Tensor.realize(U,S,V)
     assert not np.isnan(U.numpy()).any()
     assert not np.isnan(S.numpy()).any()
     assert not np.isnan(V.numpy()).any()

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.index, offset))
+      buf.src = (src, UOp.const(dtypes.weakint, offset))
       buf.dtype = dtypes.uint8
     else:
       buf.op = Ops.BUFFER

test/unit/test_multitensor.py

@@ -63,6 +63,13 @@ class TestMultiTensor(unittest.TestCase):
     np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0))).numpy(), [2, 3, 4, 5])
     np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5])
 
+  def test_add_rank_expand_shard(self):
+    # a sharded src keeps its own rank under implicit broadcast, its shard axis right-aligns into the output
+    a = Tensor([1.,2.,3.,4.]).shard(devices_2, 0)
+    b = Tensor([[10.,20.,30.,40.]]).shard(devices_2, None)
+    self.assertEqual((a+b).uop.axis, 1)
+    np.testing.assert_equal((a+b).numpy(), [[11.,22.,33.,44.]])
+
   def test_shard_reduce(self):
     self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=1), [3.,3.], n=6)
     self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=0), [2.,2.,2.], n=6)
@@ -602,8 +609,8 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
       np.testing.assert_allclose((a+a).numpy(), (b+b).numpy(), rtol=1e-7, atol=1e-3)
       np.testing.assert_equal((a+1).numpy(), (b+1).numpy())
       np.testing.assert_equal((1+a).numpy(), (1+b).numpy())
-      np.testing.assert_allclose((a.where(a+a, a)).numpy(), (b.where(b+b, b)).numpy(), rtol=1e-7, atol=1e-3)
-      np.testing.assert_allclose((a.where(1, 0)).numpy(), (b.where(1, 0)).numpy(), rtol=1e-7, atol=1e-3)
+      np.testing.assert_allclose((a.bool().where(a+a, a)).numpy(), (b.bool().where(b+b, b)).numpy(), rtol=1e-7, atol=1e-3)
+      np.testing.assert_allclose((a.bool().where(1, 0)).numpy(), (b.bool().where(1, 0)).numpy(), rtol=1e-7, atol=1e-3)
 
       # reduce
       np.testing.assert_allclose(a.max().numpy(), b.max().numpy(), rtol=1e-7, atol=1e-3)

test/unit/test_randomness.py

@@ -119,7 +119,7 @@ class TestRandomness(unittest.TestCase):
     self.assertRaises(AssertionError, lambda: Tensor(2).multinomial(1, replacement=False))
     self.assertRaises(AssertionError, lambda: Tensor([1, 9]).multinomial(0, replacement=False))
     def _check_with_torch(w, num_samples, replacement):
-      tiny_res = Tensor(w).multinomial(num_samples, replacement=replacement)
+      tiny_res = Tensor(w).multinomial(num_samples, replacement=replacement).realize()
       torch_res = torch.tensor(w).multinomial(num_samples, replacement=replacement)
       self.assertEqual(tiny_res.shape, torch_res.shape)
       if torch_res.ndim == 1:

tinygrad/callify.py

@@ -61,7 +61,7 @@ def _make_buffer_view(src:UOp) -> UOp|None:
     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(dtypes.index, offset)), src.numel())
+  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(dtypes.weakint, offset)), src.numel())
 
 def contiguous_mops_to_view(c:UOp, src:UOp):
   """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""
@@ -180,8 +180,9 @@ def finalize_after(ctx:AllocCtx, x:UOp):
 def replace_input_buffer(ctx:AllocCtx, b:UOp):
   ctx.replacements.append(b)
   return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
-                   b._min_max if b.op is Ops.BIND else None, b.src[0].expr if b.op is Ops.BIND else None,
-                   b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
+                   b._min_max if b.op is Ops.BIND else None, name=b.src[0].expr if b.op is Ops.BIND else None,
+                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL,
+                   multiple_of=b.src[0].arg.multiple_of if b.op is Ops.BIND else None)
 
 pm_finalize_call = PatternMatcher([
   (UPat(Ops.AFTER, name="x"), finalize_after),
@@ -193,7 +194,7 @@ pm_replace_buf = PatternMatcher([
   (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
    replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
   # replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
-  (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.index)), name="b"), replace_input_buffer),
+  (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer),
   # strip value from BIND for cache key normalization, so different values hit same cache
   (UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer),
 ])

tinygrad/codegen/__init__.py

@@ -17,14 +17,14 @@ from tinygrad.uop.movement import mop_cleanup
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
 from tinygrad.codegen.decomp.transcendental import get_transcendental_patterns
-from tinygrad.codegen.late.coalese import indexing_simplify
+from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
 from tinygrad.schedule.rangeify import pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
-from tinygrad.codegen.late.coalese import memory_coalesing, pm_simplify_add_image
+from tinygrad.codegen.late.coalesce import memory_coalescing, pm_simplify_add_image
 from tinygrad.helpers import all_same, flatten, argsort, partition
 from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element
 from tinygrad.schedule.rangeify import BufferizeOpts
@@ -39,8 +39,8 @@ pm_number_params = PatternMatcher([
 ])
 
 pm_no_index = PatternMatcher([
-  (UPat(GroupOp.ALU.union({Ops.CONST}), dtype=dtypes.index, name="x"), lambda x: x.replace(dtype=dtypes.int)),
-  (UPat(Ops.CAST, dtype=dtypes.index, src=(UPat.var("x"),)), lambda x: x.cast(dtypes.int)),
+  (UPat(GroupOp.ALU.union({Ops.CONST}), dtype=dtypes.weakint, name="x"), lambda x: x.replace(dtype=dtypes.int)),
+  (UPat(Ops.CAST, dtype=dtypes.weakint, src=(UPat.var("x"),)), lambda x: x.cast(dtypes.int)),
 ])
 
 def build_range_map(sink:UOp) -> dict[int, int]:
@@ -110,7 +110,7 @@ def broadcast_and_devec_wmma(b:UOp):
                   for u,shp in zip(b.src, shaped_aligned)]
   src = []
   for idx in itertools.product(*[range(i) for i in b.shape[:-1]]):
-    idx_c = [UOp.const(dtypes.index, i) for i in idx]
+    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
     src.append(b.replace(src=tuple([x.index(*idx_c) for x in src_reshaped])))
   return UOp.stack(*src).reshape(b.shape)
 
@@ -135,7 +135,7 @@ def do_devectorize(b:UOp):
   if not all_same([x.shape for x in b.src]): return None
   src = []
   for idx in itertools.product(*[range(x) for x in b.shape]):
-    idx_c = [UOp.const(dtypes.index, i) for i in idx]
+    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
     src.append(b.replace(src=tuple([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)
 
@@ -145,7 +145,7 @@ def do_stack_wmma(u:UOp):
   src = []
   for b in u.src:
     if b.op != Ops.STACK:
-      src.append(UOp.stack(*[b.index(UOp.const(dtypes.index, i)) for i in range(b.max_numel())]))
+      src.append(UOp.stack(*[b.index(UOp.const(dtypes.weakint, i)) for i in range(b.max_numel())]))
     else:
       src.append(b)
   return u.replace(src=tuple(src))
@@ -171,13 +171,10 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
   # RESHAPE a void is removed (hack for AFTER)
   (UPat(Ops.RESHAPE, dtype=dtypes.void, name="x"), lambda x: x.src[0]),
   # reshape of a single element shaped value to scalar is an index
-  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(dtypes.index, 0)) if x.marg == () and x.src[0].shape == (1,) else None),
+  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(dtypes.weakint, 0)) if x.marg == () and x.src[0].shape == (1,) else None),
   # EXPAND on scalar -> STACK
   (UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"),
    lambda x,out: UOp.stack(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None),
-  # INDEX on INDEX is INDEX
-  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
-   lambda idx1, idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:])),
 ])
 
 def fix_group_for_reduce(x:UOp):
@@ -318,11 +315,11 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   # simplify indexing
   sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
 
-  # some coalesing misses without this
+  # some coalescing misses without this
   sink = graph_rewrite(sink, sym, name="early symbolic")
 
-  # do memory coalesing (late)
-  sink = memory_coalesing(sink, ren)
+  # do memory coalescing (late)
+  sink = memory_coalescing(sink, ren)
   sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
 
   # extra symbolic before decomp. crashes without this?

tinygrad/codegen/decomp/dtype.py

@@ -33,11 +33,14 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
     case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
     case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
     case Ops.SHL:
-      lo, hi = shl(a0, b0_mod:=b0 & 31), shl(a1, b0_mod) | shr(shr(a0, 1), 31 - b0_mod)
+      a0u, a1u, n = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), (b0 & 31).cast(dtypes.uint)
+      lo, hi = (a0u << n).bitcast(dt), ((a1u << n) | ((a0u >> 1) >> (31 - n))).bitcast(dt)
       return (b0 >= 32).where(zero, lo), (b0 >= 32).where(lo, hi)
     case Ops.SHR:
-      lo, hi = shr(a0, b0_mod:=b0 & 31) | shl(shl(a1, 1), 31 - b0_mod), shr(a1, b0_mod)
-      return (b0 >= 32).where(hi, lo), (b0 >= 32).where(zero, hi)
+      a0u, a1u, n = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), (b0 & 31).cast(dtypes.uint)
+      lo, hi = ((a0u >> n) | ((a1u << 1) << (31 - n))).bitcast(dt), a1 >> (b0 & 31)
+      fill = a1 >> 31 if dt == dtypes.int else zero  # vacated high word: sign bits when signed, else 0
+      return (b0 >= 32).where(hi, lo), (b0 >= 32).where(fill, hi)
     case Ops.ADD: return (low:=a0+b0), (a1 + b1).replace(dtype=dt) + (low.bitcast(dtypes.uint) < a0.bitcast(dtypes.uint)).cast(dt)
     case Ops.SUB: return a0 - b0, a1 - b1 - (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)).cast(dt)
     case Ops.MUL:

tinygrad/codegen/gpudims.py

@@ -57,7 +57,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
 
   # get the idxs
   ki: KernelInfo = s.arg
-  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.index)]
+  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.weakint)]
   elif ki.dont_use_locals:
     assert not local_dims, "can't use locals if there's no local dims"
     idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)

tinygrad/codegen/late/coalesce.py

@@ -97,16 +97,16 @@ pm_simplify_add_image = PatternMatcher([
   (UPat.var("x", dtype=dtypes.float).cast(dtypes.half).cast(dtypes.float), lambda x: x),
 ])
 
-def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
+def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
   if getenv("DMC"): return sink
 
   # collect
   memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp], dict[int, list[UOp]]] = defaultdict(dict)
   for u in sink.toposort():
-    # TODO: this should handle images too, it's just memory coalesing
+    # TODO: this should handle images too, it's just memory coalescing
     if u.op in {Ops.LOAD, Ops.STORE}:
-      assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalesing does not support gated loads/stores"
-      assert u.src[0].op is Ops.INDEX, f"memory coalesing should be on INDEX, not {u.src[0].op}"
+      assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores"
+      assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}"
       buf, idx_u = u.src[0].src
       if buf.addrspace == AddrSpace.REG: continue
       idx, valid = idx_u.get_idx(), idx_u.get_valid()
@@ -141,12 +141,12 @@ def memory_coalesing(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(dtypes.index, full_grp[0])
+        offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(dtypes.weakint, 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(dtypes.index, len(grp)))) if len(grp) > 1 else buf.index(offset)
+        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(dtypes.weakint, len(grp)))) if len(grp) > 1 else buf.index(offset)
         if op == Ops.STORE:
           datas = []
           for i,g in enumerate(grp):
@@ -158,8 +158,8 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
           ld = idx.load()
           for i,g in enumerate(grp):
             for oo in offsets[g]:
-              replacements[oo] = ld.index(UOp.const(dtypes.index, i)) if len(grp) > 1 else ld
+              replacements[oo] = ld.index(UOp.const(dtypes.weakint, i)) if len(grp) > 1 else ld
         full_grp = full_grp[length:]
 
   # apply
-  return sink.substitute(replacements, name="memory coalesing")
+  return sink.substitute(replacements, name="memory coalescing")

tinygrad/codegen/late/linearizer.py

@@ -2,7 +2,7 @@ import heapq
 from typing import Any
 from collections import defaultdict
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
-from tinygrad.dtype import AddrSpace
+from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.helpers import prod, getenv, TUPLE_ORDER
 
 def linearize(sink:UOp) -> list[UOp]:
@@ -85,9 +85,9 @@ pm_add_control_flow = PatternMatcher([
 ])
 
 def do_split_ends(e:UOp):
-  ret = e.src[0]
-  for r in sorted(UOp.sink(*e.src[1:]).ranges, key=lambda x: x.arg, reverse=True): ret = ret.end(r)
-  return ret
+  ret, backedge = e.src[0], tuple(x for x in e.src[1:] if x.dtype in (dtypes.void, dtypes.bool))
+  for r in sorted(UOp.sink(*[x for x in e.src[1:] if x not in backedge]).ranges, key=lambda x: x.arg, reverse=True): ret = ret.end(r)
+  return ret.end(*backedge) if len(backedge) else ret
 
 pm_split_ends = PatternMatcher([
   # split the ends

tinygrad/codegen/opt/heuristic.py

@@ -2,7 +2,7 @@ import itertools
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, IMAGE
 from tinygrad.uop.ops import Ops, resolve, AxisType
-from tinygrad.codegen.late.coalese import image_valid_dims
+from tinygrad.codegen.late.coalesce import image_valid_dims
 from tinygrad.codegen.opt.postrange import Scheduler
 
 def hand_coded_optimizations(k:Scheduler) -> Scheduler:

tinygrad/codegen/opt/postrange.py

@@ -21,8 +21,9 @@ class Scheduler:
 
   @property
   def rngs(self):
-    # always in order by axistype
-    return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
+    # always in order by axistype. void RANGEs are loops, not opt axes
+    return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.dtype is not dtypes.void and u.vmax > 0],
+                  key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
   @property
   def shape_len(self) -> int: return len(self.rngs)
   @property

tinygrad/codegen/opt/tc.py

@@ -103,7 +103,7 @@ amd_rdna3 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(16,16,
   opts=("l0","l0","l0","l0","l1","u1","u1","u1"),
   swizzle=((('l4', 'u0', 'u1', 'u2', 'l0'), ('r1', 'r2', 'r3'), ('l1', 'l2', 'l3', 'r0')),
            (('l0', 'l1', 'l2', 'l3', 'l4'), ('r1', 'r2', 'r3'), ('u0', 'u1', 'u2', 'r0'))))
-  for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float)]]
+  for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float),(dtypes.int8,dtypes.int32)]]
 amd_rdna4 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(8,8,8), dtype_in=di, dtype_out=do,
   opts=("l0","l0","l0","l0","u1","u1","u1","l1"),
   swizzle=((('u0', 'u1', 'u2', 'l4', 'r2'), ('r0', 'r1', 'r3'), ('l0', 'l1', 'l2', 'l3')),

tinygrad/codegen/simplify.py

@@ -9,7 +9,9 @@ def flatten_range(r:UOp) -> UOp|None:
   off = range_start[r.op]
   rngs = r.src[off:]
   if not len(rngs): return None
-  return r.replace(src=r.src[:off]+tuple(UOp.sink(*rngs).ranges))
+  # ranges in the cond should not be ended
+  backedge = tuple(x for x in rngs if x.dtype in (dtypes.void, dtypes.bool))
+  return r.replace(src=r.src[:off]+tuple(UOp.sink(*[x for x in rngs if x not in backedge]).ranges)+backedge)
 
 pm_flatten_range = PatternMatcher([
   # real ranges only
@@ -19,6 +21,7 @@ pm_flatten_range = PatternMatcher([
 # index/range arithmetic uses FLOORDIV/FLOORMOD prior to late rewrite
 def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.FLOORDIV, Ops.FLOORMOD} for u in x.backward_slice)
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
+  if not all(r.op is Ops.RANGE for r in u.ended_ranges): return None
   reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
   # on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations
   for r0, r1 in (zip(u.ended_ranges, u.ended_ranges[1:]) if u.op is Ops.END else itertools.permutations(u.ended_ranges, 2)):
@@ -149,5 +152,5 @@ def no_load(u:UOp) -> bool: return not any(x.op is Ops.INDEX for x in u.backward
 pm_load_collapse = PatternMatcher([
   (UPat(Ops.REDUCE, arg=(Ops.ADD, 0), src=(UPat.var("u"), UPat()), name="red"), reduce_load_collapse),
   # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
-  ((UPat.var("x", dtypes.index)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
+  ((UPat.var("x", dtypes.weakint)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
 ])

tinygrad/device.py

@@ -100,8 +100,8 @@ class MultiBuffer:
 
 class Buffer:
   profile_events:list[ProfileEvent] = []
-  def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None, initial_value:bytes|None=None,
-               uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False):
+  def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None,
+               initial_value:bytes|pickle.PickleBuffer|None=None, uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False):
     assert isinstance(dtype, DType)
     self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0
     self._bufs: dict[str, Any] = {}
@@ -113,6 +113,7 @@ class Buffer:
       if initial_value is not None:
         self.allocate()
         self.copy_from(Buffer("PYTHON", self.size, self.dtype, opaque=memoryview(bytearray(initial_value))))
+        if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
     else:
       assert base._base is None, "base can't have a base"
       assert device == base.device, "base must have the same device"
@@ -171,12 +172,13 @@ class Buffer:
       self.allocator.free(self._buf, self.nbytes, self.options)
     elif self._base is not None: self._base.allocated_views -= 1
     self._bufs.clear()
-  def __reduce__(self):
-    buf = None
+  def __reduce_ex__(self, protocol):
+    buf:bytearray|pickle.PickleBuffer|None = None
     if self._base is not None:
       return self.__class__, (self.device, self.size, self.dtype, None, None, None, 0, self.base, self.offset, self.is_allocated())
     if self.device == "NPY": return self.__class__, (self.device, self.size, self.dtype, self._buf, self.options, None, self.uop_refcount)
-    if self.is_allocated(): buf = bytearray(self.as_memoryview())
+    if self.is_allocated():
+      buf = pickle.PickleBuffer(self.as_memoryview()) if protocol >= 5 else bytearray(self.as_memoryview())
     return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount)
   @property
   def trace_num(self) -> int:

tinygrad/dtype.py

@@ -89,7 +89,7 @@ class dtypes:
   def is_float(x: DType) -> bool: return x in (dtypes.floats + (dtypes.weakfloat,))
   @staticmethod # static methods on top, or bool in the type info will refer to dtypes.bool
   @functools.cache
-  def is_int(x: DType) -> bool: return x in (dtypes.ints + (dtypes.weakint, dtypes.index))
+  def is_int(x: DType) -> bool: return x in (dtypes.ints + (dtypes.weakint,))
   @staticmethod
   @functools.cache
   def is_unsigned(x: DType) -> bool: return x in dtypes.uints
@@ -111,8 +111,7 @@ class dtypes:
     return {dtypes.float16: (5, 10), dtypes.bfloat16: (8, 7), dtypes.float32: (8, 23), dtypes.float64: (11, 52),
             dtypes.fp8e4m3: (4, 3), dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3fnuz: (4, 3), dtypes.fp8e5m2fnuz: (5, 2)}[dtype]
   void: Final[DType] = DType.new(-1, 0, "void", None)
-  weakint: Final[DType] = DType.new(0, 800, "weakint", None)
-  index: Final[DType] = DType.new(0, 800, "index", None)  # NOTE: not in the promo lattice: index math never mixes dtypes
+  weakint: Final[DType] = DType.new(0, 800, "weakint", None)  # NOTE: not in the promo lattice: index math never mixes dtypes
   bool: Final[DType] = DType.new(0, 1, "bool", '?')
   int8: Final[DType] = DType.new(1, 8, "signed char", 'b')
   uint8: Final[DType] = DType.new(2, 8, "unsigned char", 'B')
@@ -154,7 +153,7 @@ class dtypes:
   uints = (uint8, uint16, uint32, uint64)
   sints = (int8, int16, int32, int64)
   ints = uints + sints
-  weaks = (weakint, weakfloat)
+  weaks = (weakfloat,)
   all = floats + ints + (bool,) # noqa: A003
 
 if (env_default_float := getenv("DEFAULT_FLOAT", "")):
@@ -164,12 +163,13 @@ if (env_default_float := getenv("DEFAULT_FLOAT", "")):
 DTypeLike = str|DType
 def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) else getattr(dtypes, dtype.lower())
 def strong_dtype(dtype:DType) -> DType:
-  return dtypes.default_int if dtype == dtypes.weakint else dtypes.default_float if dtype == dtypes.weakfloat else dtype
+  # TODO: weakint
+  return dtypes.default_float if dtype == dtypes.weakfloat else dtype
 
 # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html
 # we don't support complex type
-promo_lattice = { dtypes.bool: [dtypes.weakint], dtypes.weakint: [dtypes.int8, dtypes.uint8],
-  dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
+# TODO: weakint
+promo_lattice = { dtypes.bool: [dtypes.int8, dtypes.uint8], dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
   dtypes.int64: [dtypes.uint64], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32],
   dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.weakfloat],
   dtypes.weakfloat: [dtypes.fp8e4m3, dtypes.fp8e5m2, dtypes.fp8e4m3fnuz, dtypes.fp8e5m2fnuz],
@@ -185,8 +185,8 @@ def least_upper_dtype(*ds:DType) -> DType:
   return min(set.intersection(*[_get_recursive_parents(d) for d in ds]))
 def least_upper_float(dt:DType) -> DType: return dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.default_float)
 
-DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "weak", "index", "_"))}
-INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "weakint":"weakint", "index":"index", "weakfloat":"weakfloat"}
+DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "weak", "_"))}
+INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "weakint":"weakint", "weakfloat":"weakfloat"}
 
 @functools.cache
 def can_lossless_cast(dt0:DType, dt1:DType) -> bool:
@@ -194,7 +194,7 @@ def can_lossless_cast(dt0:DType, dt1:DType) -> bool:
   # similar to https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
   if dt0 == dt1 or dt0 == dtypes.bool: return True
   match dt1:
-    case dtypes.weakint | dtypes.index: return dt0 in dtypes.ints
+    case dtypes.weakint: return dt0 in dtypes.ints
     case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16, *dtypes.fp8s,
       dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
     case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s, dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)

tinygrad/helpers.py

@@ -81,7 +81,6 @@ def to_be32(val:Any) -> Any: return ((val & 0xFF) << 24) | (((val >> 8) & 0xFF)
 def to_be64(val:Any) -> Any: return to_be32(val >> 32) | (to_be32(val & 0xFFFFFFFF) << 32)
 def getbits(value: int, start: int, end: int): return (value >> start) & ((1 << (end - start + 1)) - 1)
 def i2u(bits: int, value: int): return value if value >= 0 else (1<<bits)+value
-def is_numpy_ndarray(x) -> bool: return str(type(x)) == "<class 'numpy.ndarray'>"
 def merge_dicts(ds:Iterable[dict[T,U]]) -> dict[T,U]:
   kvs = set([(k,v) for d in ds for k,v in d.items()])
   if len(kvs) != len(set(kv[0] for kv in kvs)): raise RuntimeError(f"{kvs} contains different values for the same key")
@@ -413,7 +412,7 @@ def diskcache_get(table:str, key:dict|str|int) -> Any:
   if (val:=res.fetchone()) is not None: return pickle.loads(val[0])
   return None
 
-_db_tables = set()
+_db_tables: set[str] = set()
 def diskcache_put(table:str, key:dict|str|int, val:Any, prepickled=False):
   if CACHELEVEL < 1: return val
   if isinstance(key, (str,int)): key = {"key": key}

tinygrad/llm/cli.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
-import sys, argparse, codecs, typing, re, unicodedata, json, uuid, time, pathlib
+import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time
+from typing import TYPE_CHECKING
 from tinygrad import nn
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored, Context, fetch, profile_marker, getenv
-from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
+from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, Context, fetch, profile_marker, getenv
 from tinygrad.llm.model import Transformer
+if TYPE_CHECKING:
+  import jinja2
 
 class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3",
@@ -18,7 +20,11 @@ class SimpleTokenizer:
 
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
     # 0x323b0 is one past the max codepoint in unicode categories L/N/Z (0x323af is max L)
-    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
+    # compact adjacent codepoints into ranges: listing them all makes re spend seconds on large prompts
+    def ucat_range(pre:str) -> str:
+      cps = enumerate(cp for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
+      runs = [list(g) for _, g in itertools.groupby(cps, lambda e: e[1]-e[0])]
+      return "".join(re.escape(chr(g[0][1])) + (f"-{re.escape(chr(g[-1][1]))}" if len(g) > 1 else "") for g in runs)
     r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
     self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
       f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
@@ -64,25 +70,6 @@ class SimpleTokenizer:
     dec = codecs.getincrementaldecoder('utf-8')('replace')
     def _decode(tid:int|None=None) -> str: return dec.decode(self._tok2bytes[tid]) if tid is not None else dec.decode(b'', final=True)
     return _decode
-  def role(self, role:str):
-    if self.preset == 'olmo': return self.encode("<|" + role + "|>\n")  # OLMoE Instruct format
-    if self.preset == 'kimi-k2': return self.encode("<|im_" + role + "|>" + role + "<|im_middle|>")
-    if self.preset == 'qwen2': return self.encode("<|im_start|>" + role + "\n")
-    if self.preset == 'glm4': return self.encode("<|" + role + "|>")
-    if self.preset == 'tekken':
-      if role == 'user': return self.encode("[INST]")
-      if role == 'assistant': return []
-      raise ValueError(f"Unsupported role '{role}' for tokenizer preset '{self.preset}'")
-    return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
-  def end_turn(self):
-    if self.preset == 'olmo': return self.encode("\n")
-    if self.preset == 'kimi-k2': return [self.eos_id]
-    if self.preset == 'qwen2': return [self.eos_id] + self.encode("\n")
-    if self.preset == 'glm4': return []
-    if self.preset == 'tekken': return self.encode("[/INST]")
-    return [self.eos_id]
-  def prefix(self) -> list[int]:
-    return ([] if self.bos_id is None else [self.bos_id]) + (self.encode("<sop>") if self.preset == 'glm4' else [])
   def is_end(self, token_id:int) -> bool: return token_id in (self.eos_id, self.eot_id)
 
 models = {
@@ -105,81 +92,41 @@ models = {
   "glm-4.7-flash": "https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/resolve/main/GLM-4.7-Flash-Q4_K_M.gguf",
 }
 
-# *** simple OpenAI API compatible server with web interface on http://localhost:8000/ ***
+class FallbackTemplate:
+  # minimal jinja2.Template-compatible chat template without jinja2, no tool calling support
+  def __init__(self, tok:SimpleTokenizer): self.tok = tok
+  def role(self, role:str) -> str:
+    if self.tok.preset == 'olmo': return "<|" + role + "|>\n"  # OLMoE Instruct format
+    if self.tok.preset == 'kimi-k2': return "<|im_" + role + "|>" + role + "<|im_middle|>"
+    if self.tok.preset == 'qwen2': return "<|im_start|>" + role + "\n"
+    if self.tok.preset == 'glm4': return "<|" + role + "|>"
+    if self.tok.preset == 'tekken':
+      if role == 'user': return "[INST]"
+      if role == 'assistant': return ""
+      raise ValueError(f"Unsupported role '{role}' for tokenizer preset '{self.tok.preset}'")
+    return "<|start_header_id|>" + role + "<|end_header_id|>\n\n"
+  def end_turn(self) -> str:
+    if self.tok.preset == 'olmo': return "\n"
+    if self.tok.preset == 'kimi-k2': return self.tok.decode([self.tok.eos_id])
+    if self.tok.preset == 'qwen2': return self.tok.decode([self.tok.eos_id]) + "\n"
+    if self.tok.preset == 'glm4': return ""
+    if self.tok.preset == 'tekken': return "[/INST]"
+    return self.tok.decode([self.tok.eos_id])
+  def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True) -> str:
+    out = self.tok.decode([] if self.tok.bos_id is None else [self.tok.bos_id]) + ("<sop>" if self.tok.preset == 'glm4' else "")
+    for msg in messages:
+      out += self.role(msg["role"])
+      content = msg.get("content")
+      if isinstance(content, str): out += content
+      elif isinstance(content, list):
+        for c in content:
+          if c["type"] == "text": out += c["text"]
+          else: raise RuntimeError(f"unhandled type: {c['type']}")
+      elif content is not None: raise RuntimeError(f"unknown content type: {type(content)}")
+      out += self.end_turn()
+    return out + self.role("assistant") if add_generation_prompt else out
 
-class Handler(HTTPRequestHandler):
-  server: LLMServer
-  def log_request(self, code='-', size='-'): pass
-  def do_GET(self):
-    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
-    cache_start_pos = model.get_start_pos(ids)
-    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  "
-               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}
-    yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
-    out: list[int] = []
-    finish_reason = "stop"
-    st = time.perf_counter()
-    dec = tok.stream_decoder()
-    for next_id in model.generate(ids, temperature=temperature):
-      if len(out) == 0: stderr_log(f"prefill:{(len(ids)-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)
-      yield {"choices": [{"index":0, "delta":{"content":dec(next_id)}, "finish_reason":None}], **tmpl}
-      if max_tokens is not None and len(out) >= max_tokens:
-        finish_reason = "length"
-        break
-    if (tail := dec()): yield {"choices": [{"index":0, "delta":{"content":tail}, "finish_reason":None}], **tmpl}
-    yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
-    if include_usage:
-      yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + 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")
-
-  def do_POST(self):
-    tok = self.server.tok
-    raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
-    body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
-    if DEBUG >= 1: print(json.dumps(body, indent=2))
-    if self.path == "/v1/chat/completions":
-      # extract tokens, last assistant message is treated as prefill
-      ids: list[int] = tok.prefix()
-      for i, msg in enumerate(body["messages"]):
-        ids += tok.role(msg["role"])
-        content = msg["content"]
-        if isinstance(content, str): ids += tok.encode(content)
-        elif isinstance(content, list):
-          for c in content:
-            if c["type"] == "text": ids += tok.encode(c["text"])
-            else: raise RuntimeError(f"unhandled type: {c['type']}")
-        else: raise RuntimeError(f"unknown content type: {type(content)}")
-        if msg["role"] == "assistant" and i == len(body["messages"]) - 1: break
-        ids += tok.end_turn()
-      else: ids += tok.role("assistant")
-
-      # reply
-      max_tokens = body.get("max_completion_tokens") or body.get("max_tokens")
-      chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False),
-                              max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0)))
-      if body.get("stream"): self.stream_json(chunks)
-      else:
-        out, finish_reason = [], "stop"
-        for c in chunks:
-          if c["choices"] and c["choices"][0].get("delta", {}).get("content"): out.append(c["choices"][0]["delta"]["content"])
-          if c["choices"] and c["choices"][0].get("finish_reason"): finish_reason = c["choices"][0]["finish_reason"]
-        self.send_data(json.dumps({**c, "object":"chat.completion",
-          "choices":[{"index":0, "message":{"role":"assistant","content":"".join(out)}, "finish_reason":finish_reason}]}).encode())
-    else:
-      raise RuntimeError(f"unhandled path {self.path}")
-
-class LLMServer(TCPServerWithReuse):
-  def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer):
-    self.model, self.model_name, self.tok = model, model_name, tok
-    super().__init__(server_address, Handler)
+from tinygrad.llm.serve import LLMServer
 
 def main():
   parser = argparse.ArgumentParser()
@@ -194,11 +141,26 @@ def main():
   model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
   model_name = kv.get('general.name') or kv.get('general.basename') or args.model
   file_sizes = [y.nbytes() for y in UOp.sink(*[x.uop for x in nn.state.get_parameters(model)]).toposort() if y.op is Ops.BUFFER]
-  print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params")
+  print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params, "
+        f"max context {args.max_context} on {nn.state.get_parameters(model)[0].device}")
 
   # get tokenizer
   tok = SimpleTokenizer.from_gguf_kv(kv)
 
+  # use the model's chat template if jinja2 is available (enables model-specific formatting)
+  template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
+  if (ct := kv.get('tokenizer.chat_template')) is not None:
+    try:
+      import jinja2
+      env = jinja2.Environment()
+      env.filters['tojson'] = lambda obj, **kwargs: json.dumps(obj, **kwargs)  # jinja2's tojson escapes <>& for HTML safety
+      env.globals['raise_exception'] = lambda msg: (_ for _ in ()).throw(RuntimeError(msg))
+      env.globals['strftime_now'] = lambda fmt: time.strftime(fmt)
+      env.globals['bos_token'] = tok.decode([tok.bos_id]) if tok.bos_id is not None else ""
+      env.globals['eos_token'] = tok.decode([tok.eos_id])
+      template = env.from_string(ct)
+    except ImportError: print("warning: jinja2 is not installed, the model's chat template is disabled")
+
   # warmup the JIT
   if args.warmup or args.serve:
     # run 2 tokens through the model twice to capture the JIT before serving
@@ -206,7 +168,7 @@ def main():
       for _ in range(2): list(zip(range(2), model.generate([0])))
 
   # start server
-  if args.serve: LLMServer(('', args.serve), model, model_name, tok).serve_forever()
+  if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever()
 
   # do benchmark
   if args.benchmark is not None:
@@ -224,16 +186,19 @@ def main():
     exit(0)
 
   # interactive chat
-  ids: list[int] = tok.prefix()
+  messages: list[dict] = []
   while 1:
-    try:
-      ids += tok.role("user") + tok.encode(input('>>> ')) + tok.end_turn() + tok.role("assistant")
-    except EOFError:
-      break
-    dec = tok.stream_decoder()
+    try: messages.append({"role":"user", "content":input('>>> ')})
+    except EOFError: break
+    ids = tok.encode(template.render(messages=messages, add_generation_prompt=True))
+    reply, dec = "", tok.stream_decoder()
     for next_id in model.generate(ids):
-      sys.stdout.write(dec(next_id) if not tok.is_end(next_id) else dec() + "\n\n")
+      if tok.is_end(next_id):
+        sys.stdout.write(dec() + "\n\n")
+        break
+      reply += (piece := dec(next_id))
+      sys.stdout.write(piece)
       sys.stdout.flush()
-      if tok.is_end(next_id): break
+    messages.append({"role":"assistant", "content":reply})
 
 if __name__ == "__main__": main()

tinygrad/llm/serve.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+import json, pathlib, re, time, typing, uuid
+from typing import TYPE_CHECKING
+from tinygrad.helpers import DEBUG, colored, stderr_log
+from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
+if TYPE_CHECKING:
+  from tinygrad.llm.cli import SimpleTokenizer
+  from tinygrad.llm.model import Transformer
+
+def parse_tool_call(s:str) -> tuple[str, typing.Any]|None:
+  s = s.strip()
+  if s.startswith("{"):  # hermes JSON format: {"name": ..., "arguments": {...}}
+    try:
+      call = json.loads(s)
+      return call["name"], call.get("arguments", call.get("parameters", {}))
+    except (json.JSONDecodeError, KeyError): return None
+  # XML format: <function=name>\n<parameter=key>\nvalue\n</parameter>...</function>
+  if (fm := re.match(r"<function=([^>]+)>\s*(.*?)\s*(?:</function>)?$", s, re.DOTALL)):
+    args = {}
+    for pm in re.finditer(r"<parameter=([^>]+)>(.*?)</parameter>", fm.group(2), re.DOTALL):
+      value = re.sub(r"^\r?\n|\r?\n\Z", "", pm.group(2))
+      try: args[pm.group(1)] = json.loads(value)
+      except json.JSONDecodeError: args[pm.group(1)] = value
+    return fm.group(1), args
+  return None
+
+def normalize_messages(messages:list[dict]) -> None:
+  # chat templates expect tool_call arguments as dicts (OpenAI clients send JSON strings)
+  for m in messages:
+    for tc in m.get("tool_calls") or []:
+      if "function" in tc and isinstance(args := tc["function"].get("arguments"), str):
+        try: tc["function"]["arguments"] = json.loads(args)
+        except json.JSONDecodeError: pass
+
+class StreamRouter:
+  # routes streamed output text to (field, text) deltas, keeping tool_call regions in .buf for the final parse
+  def __init__(self):
+    self.buf = ""
+    self.mode = "undecided"  # output inside a think block is sent as reasoning_content
+  def split(self, tag:str, final:bool) -> tuple[str, bool]:
+    # split buf on the first full tag, holding back a partial tag at the end unless final
+    if tag in self.buf:
+      before, self.buf = self.buf.split(tag, 1)
+      return before, True
+    hold = max((i for i in range(1, min(len(self.buf), len(tag))+1) if tag.startswith(self.buf[-i:])), default=0) if not final else 0
+    emit, self.buf = self.buf[:len(self.buf)-hold], self.buf[len(self.buf)-hold:]
+    return emit, False
+  def route(self, piece:str, final:bool=False) -> typing.Iterator[tuple[str, str]]:
+    self.buf += piece
+    if self.mode == "undecided":  # decide whether the output starts with a think block
+      if not final and len(self.buf) < len("<think>") and "<think>".startswith(self.buf): return
+      self.mode, self.buf = ("reasoning", self.buf[len("<think>"):]) if self.buf.startswith("<think>") else ("content", self.buf)
+    if self.mode == "reasoning":
+      emit, done = self.split("</think>", final)
+      if emit: yield "reasoning_content", emit
+      if not done: return
+      self.mode = "content"
+    if self.mode == "tool": return
+    emit, found = self.split("<tool_call>", final)
+    if emit: yield "content", emit
+    if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
+
+class Handler(HTTPRequestHandler):
+  server: LLMServer
+  def log_request(self, code='-', size='-'): pass
+  def do_GET(self):
+    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
+    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:{(len(ids)-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": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + 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")
+
+  def do_POST(self):
+    request_st = time.perf_counter()
+    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  ")
+    raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
+    body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
+    if DEBUG >= 1: print(json.dumps(body, indent=2))
+    if self.path == "/v1/chat/completions":
+      # render and tokenize
+      normalize_messages(body["messages"])
+      rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True)
+      ids: list[int] = self.server.tok.encode(rendered)
+      stderr_log(f"prep:{(time.perf_counter()-request_st)*1e3:5.0f} ms  {colored('--', 'BLACK')}  ")
+      if len(ids) >= self.server.model.max_context:
+        stderr_log(f"{colored('context length exceeded', 'red')}  in:{len(ids):5d}  max:{self.server.model.max_context:5d}\n")
+        return self.send_data(json.dumps({"error":{"message":f"prompt has {len(ids)} tokens, but the model context is "
+          f"{self.server.model.max_context}", "type":"invalid_request_error", "param":"messages", "code":"context_length_exceeded"}}).encode(),
+          status_code=400)
+
+      # reply
+      max_tokens = body.get("max_completion_tokens") or body.get("max_tokens")
+      chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False),
+                              max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0)))
+      if body.get("stream"): self.stream_json(chunks)
+      else:
+        out, reasoning, tool_calls, finish_reason = [], [], [], "stop"
+        for c in chunks:
+          if not c["choices"]: continue
+          choice = c["choices"][0]
+          if (delta := choice.get("delta", {})):
+            if delta.get("content"): out.append(delta["content"])
+            if delta.get("reasoning_content"): reasoning.append(delta["reasoning_content"])
+            tool_calls += [{k:v for k, v in tc.items() if k != "index"} for tc in delta.get("tool_calls", [])]
+          if choice.get("finish_reason"): finish_reason = choice["finish_reason"]
+        message: dict[str, typing.Any] = {"role":"assistant", "content":"".join(out) or None}
+        if reasoning: message["reasoning_content"] = "".join(reasoning)
+        if tool_calls: message["tool_calls"] = tool_calls
+        self.send_data(json.dumps({**c, "object":"chat.completion",
+          "choices":[{"index":0, "message":message, "finish_reason":finish_reason}]}).encode())
+    else:
+      raise RuntimeError(f"unhandled path {self.path}")
+
+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
+    super().__init__(server_address, Handler)

tinygrad/mixin/elementwise.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import argfix, polyN
 from tinygrad.mixin.creation import CreationMixin
 
 if TYPE_CHECKING:
-  from tinygrad.uop.ops import UOp
+  from tinygrad.uop.ops import UOp, sint
 
 
 class ElementwiseMixin(CreationMixin):
@@ -18,9 +18,11 @@ class ElementwiseMixin(CreationMixin):
   def ufix(self, x: 'Self|ConstType|UOp') -> Self:
     return x if isinstance(x, type(self)) else self._wrap_uop(self._uop.ufix(x))
 
-  # implemented in OpMixin, broadcasting needs the movement ops
   def _broadcasted(self, y: 'Self|ConstType|UOp', reverse: bool = False) -> tuple[Self, Self]:
-    raise NotImplementedError
+    y = self.ufix(y)
+    x, y = (self, y) if not reverse else (y, self)
+    if x.dtype == y.dtype: return x, y
+    return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype)
 
   def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self:
     lhs, rhs = self._broadcasted(x, reverse)
@@ -49,6 +51,7 @@ class ElementwiseMixin(CreationMixin):
     """
     Returns a contiguous tensor.
     """
+    if self.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {self.dtype}")
     uop = self._uop
     if uop.op is Ops.CONTIGUOUS or self.device is None or uop.has_buffer_identity(): return self._wrap_uop(uop)
     return self._wrap_uop(uop.alu(Ops.CONTIGUOUS, **kwargs))
@@ -413,10 +416,19 @@ class ElementwiseMixin(CreationMixin):
     m = a.maximum(b)
     return ((a-m).exp() + (b-m).exp()).log() + m
 
-  def where(self, x: Self | ConstType, y: Self | ConstType) -> Self:
-    ref: Self = x if isinstance(x, type(self)) else y if isinstance(y, type(self)) else \
-      self.cast(least_upper_dtype(dtypes.from_py(x), dtypes.from_py(y)))
-    return self.alu(Ops.WHERE, ref.ufix(x), ref.ufix(y))
+  def where(self, x: 'Self | ConstType | sint', y: 'Self | ConstType | sint') -> Self:
+    """
+    Returns a tensor of elements selected from either `x` or `y`, depending on `self`.
+    `output_i = x_i if self_i else y_i`.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    cond = Tensor([[True, True, False], [True, False, False]])
+    print(cond.where(1, 3).numpy())
+    ```
+    """
+    ref = x if isinstance(x, type(self)) else y if isinstance(y, type(self)) else self
+    x, y = ref.ufix(x)._broadcasted(y)
+    return self.alu(Ops.WHERE, x, y)
 
   def masked_fill(self, mask:Self, value:Self|PyConst) -> Self:
     """
@@ -547,9 +559,7 @@ class ElementwiseMixin(CreationMixin):
     # TODO: int pow
     if not base.is_floating_point() and isinstance(x, ConstType) and not (isinstance(x, int) and x >= 0):
       raise RuntimeError("base needs to be float")
-    ret = base.alu(Ops.POW, exponent)
-    # NOTE: pow(int, float) -> int
-    return ret.round().cast(self.dtype) if not reverse and not dtypes.is_float(self.dtype) and dtypes.is_float(exponent.dtype) else ret
+    return base.alu(Ops.POW, exponent)
 
   def __pow__(self, x: Self | ConstType) -> Self:
     return self.pow(x)

tinygrad/mixin/gradient.py

@@ -1,18 +1,13 @@
 from typing import cast
 import math, dataclasses
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
-  def broadcast_to_input(x:UOp) -> UOp: return x._broadcast_to(ret.src[0].shape)
-  if op == Ops.ADD: return (broadcast_to_input(ctx),)
-  if op == Ops.MAX:
-    assert ret.op is Ops.REDUCE, "only works on REDUCE"
-    mask = ret.src[0].eq(broadcast_to_input(ret)).cast(ctx.dtype)
-    count = mask._rop(Ops.ADD, tuple(range(ret.arg[1])))
-    return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
-  if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
+  if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),)
+  if op == Ops.MAX: return (((mask:=ret.src[0].eq(ret).cast(ctx.dtype))/mask._rop(Ops.ADD, tuple(range(ret.arg[1])))) * ctx,)
+  if op == Ops.MUL: return (ctx * ret / ret.src[0],)
 
 def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]:
   """Remove unused PARAMs from body and return compacted (body, args)."""
@@ -67,9 +62,7 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.CONTIGUOUS), lambda ctx: (ctx,)),
   (UPat(Ops.CONTIGUOUS_BACKWARD), lambda ctx: (ctx.contiguous(),)),
   (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
-  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret:
-    (ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, tuple(range(len(ret.marg))))
-     .reshape(ret.src[0].shape).cast(ctx.dtype), None)),
+  (UPat(Ops.EXPAND), lambda ctx: (ctx, None)),
   (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[0]-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
   (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
@@ -119,6 +112,9 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
     assert len(lgrads) == len(t0.src), f"got {len(lgrads)} gradient, expected {len(t0.src)}"
     for k,v in zip(t0.src, lgrads):
       if v is None: continue
+      # a shaped edge's gradient is summed to its source's shape
+      if k._shape is not None and v._shape is not None and k._shape != v._shape:
+        v = v.cast(sum_acc_dtype(v.dtype))._rop(Ops.ADD, broadcast_axes(k.shape, v.shape)).reshape(k.shape).cast(v.dtype)
       if k in grads and grads[k].op is not Ops.NOOP:
         if v.op is Ops.TUPLE and grads[k].op is Ops.TUPLE:
           grads[k] = UOp.maketuple(*(p + n if (p.op is not Ops.NOOP and n.op is not Ops.NOOP) else

tinygrad/mixin/movement.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Self, Sequence
 from tinygrad.uop import Ops
 from tinygrad.helpers import prod, argfix, argsort, flatten, dedup, make_tuple, ceildiv, round_up, all_int
-from tinygrad.uop.ops import resolve, smax, _align_left, _broadcast_shape
+from tinygrad.uop.ops import resolve, smax, _align_left, _broadcast_shape, broadcast_axes
 
 if TYPE_CHECKING:
   from tinygrad.uop.ops import sint
@@ -125,7 +125,7 @@ class MovementMixin:
       raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}")
     # EXPAND only adds dims on the left. squeeze 1s that need expanding, EXPAND on left, permute back.
     n_left = len(new_shape) - len(self.shape)
-    expand_at = tuple(i for i, s in enumerate(self.shape) if resolve(s == 1, default=False) and resolve(new_shape[n_left+i] != 1))
+    expand_at = tuple(i-n_left for i in broadcast_axes(self.shape, new_shape) if i >= n_left)
     kept = tuple(i for i in range(len(self.shape)) if i not in expand_at)
     squeezed = self.reshape(tuple(self.shape[i] for i in kept))
     expanded = squeezed._mop(Ops.EXPAND, arg=new_shape[:n_left] + tuple(new_shape[n_left+i] for i in expand_at))

tinygrad/mixin/op.py

@@ -11,7 +11,7 @@ from tinygrad.helpers import all_int, argfix, argsort, ceildiv, flatten, flat_to
 from tinygrad.helpers import resolve_pool_pads, round_up, IMAGE, FLOAT16, WINO
 
 if TYPE_CHECKING:
-  from tinygrad.uop.ops import sint, UOp
+  from tinygrad.uop.ops import sint
 
 ReductionStr = Literal["mean", "sum", "none"]
 
@@ -110,8 +110,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
       consecutive = dims == list(range(dims[0], dims[0] + len(dims)))
       if v is None and len(dims) > 1 and consecutive and all_int(ishp := tuple(x.shape[d] for d in dims)):
         strides = tuple(prod(ishp[i+1:]) for i in range(len(dims)))
-        try: linear_idx = type(self).usum(*[t._broadcast_to(big_shape) * s for t, s in zip(tensors, strides)])
-        except ValueError as err: raise IndexError(f"cannot broadcast indices: {err}") from err
+        linear_idx = type(self).usum(*[t * s for t, s in zip(tensors, strides)])
         valid = type(self).uprod(*[(t >= 0) & (t < s) for t, s in zip(tensors, ishp)])
         pre, post = x.shape[:dims[0]], x.shape[dims[-1]+1:]
         x = x.reshape(pre + (prod(ishp),) + post)[tuple([slice(None)] * len(pre)) + (valid.where(linear_idx, 0),)]
@@ -284,9 +283,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     X = self.shrink(tuple((-smin(pB,0),smin(pA+s,s)) for (pB,pA),s in zip(pX, self.shape))) if has_neg else self
     pads = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) if has_neg else pX
     base = MovementMixin.pad(X, pads)
-    if value is Invalid: return base
-    if value != 0: base = base.cast(least_upper_dtype(base.dtype, dtypes.from_py(value)))
-    return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, base.const_like(value))
+    if value == 0: return base
+    return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, value)
 
   def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
     # shrink first for negative pads, then wrap the non-negative remainder
@@ -357,17 +355,6 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if mode in {"reflect", "replicate"}: return self._pad_reflect_replicate(pX, mode)
     raise NotImplementedError(f"{mode=} is not supported")
 
-  def _broadcasted(self, y:Self|ConstType|UOp, reverse:bool=False) -> tuple[Self, Self]:
-    if not isinstance(y, type(self)): y = self.ufix(y)
-    x, y = (self, y) if not reverse else (y, self)
-    # ValueError: unsized ptr has shape (-1,) which can't broadcast; RuntimeError: shape mismatch
-    try:
-      out_shape = _broadcast_shape(x.shape, y.shape)
-      x, y = x._broadcast_to(out_shape), y._broadcast_to(out_shape)
-    except (RuntimeError, ValueError): pass
-    if x.dtype == y.dtype: return x, y
-    return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype)
-
   def dot(self, w:Self, dtype:DTypeLike|None=None) -> Self:
     """
     Performs dot product between two tensors.
@@ -658,7 +645,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     print(t.logsumexp(axis=1).numpy())
     ```
     """
-    m = self.max(axis=axis, keepdim=True)
+    m = self.max(axis=axis, keepdim=True).detach()
     return (self - m).exp().sum(axis=axis, keepdim=keepdim).log() + (m if keepdim else m.squeeze(axis))
 
   def _softmax(self, axis, dtype:DTypeLike|None=None) -> tuple[Self, Self, Self]:
@@ -841,8 +828,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     x = self.transpose(axis, -1)
     last_dim_size = x.shape[-1]
     x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None))
-    x_cummax, _ = x.cummax(-1)
-    mask = type(self).ones(last_dim_size, last_dim_size, buffer=False).tril()
+    x_cummax = x.cummax(-1)[0].detach()
+    mask = type(self).ones(last_dim_size, last_dim_size, buffer=False, dtype=dtypes.bool).tril()
     ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax
     return ret.transpose(-1, axis)
 

tinygrad/mixin/rand.py

@@ -11,7 +11,7 @@ class RandMixin(OpMixin):
   @staticmethod
   def _threefry_random_bits(key, counts0, counts1):
     x = (counts1.cast(dtypes.uint64) << 32) | counts0.cast(dtypes.uint64)
-    x = x.threefry((key[1]._broadcast_to(x.shape).cast(dtypes.uint64) << 32) | key[0]._broadcast_to(x.shape).cast(dtypes.uint64))
+    x = x.threefry((key[1].cast(dtypes.uint64) << 32) | key[0].cast(dtypes.uint64))
     return (x & 0xffffffff).cast(dtypes.uint32).cat(((x >> 32) & 0xffffffff).cast(dtypes.uint32))
 
   @classmethod

tinygrad/nn/__init__.py

@@ -348,12 +348,12 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
       # each device owns [offset, offset+local_vocab_size) of the global vocabulary
       dnum = UOp.variable("_device_num", 0, ndev-1)
       offset = dnum * local_vocab_size
-      global_token_id = idx_flat[i].cast(dtypes.index)
+      global_token_id = idx_flat[i].cast(dtypes.weakint)
       local_token_id = (global_token_id - offset).clip(0, grad_weight.shape[0]-1)
       in_range = (global_token_id >= offset) & (global_token_id < (offset + local_vocab_size)) & j_ok
       grad_val = in_range.where(grad_emb_flat[i, j_idx].load().cast(dtypes.float), 0.0)
     else:
-      local_token_id = idx_flat[i].clip(0, grad_weight.shape[0]-1).cast(dtypes.index)
+      local_token_id = idx_flat[i].clip(0, grad_weight.shape[0]-1).cast(dtypes.weakint)
       grad_val = j_ok.where(grad_emb_flat[i, j_idx].load().cast(dtypes.float), 0.0)
     # atomic scatter-add: grad_weight[token_id, j] += grad_emb_flat[i, j]
     if device in ("CPU", "NULL"): atomic_arg = "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"

tinygrad/nn/onnx.py

@@ -2,9 +2,9 @@
 from typing import Any, Sequence, cast, Literal, NamedTuple, Generator
 import dataclasses, functools, io, math, types, warnings, pathlib, sys, os, struct, enum
 from tinygrad.nn.state import TensorIO
-from tinygrad.tensor import Tensor
+from tinygrad.tensor import Tensor, is_numpy_ndarray
 from tinygrad.mixin.op import ReductionStr
-from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN, Context
+from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, get_single_element, polyN, Context
 from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate, least_upper_dtype, DTYPES_DICT
 from tinygrad.device import Device
 from tinygrad.uop.ops import sint, _broadcast_shape
@@ -617,6 +617,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y
   def Sub(x:Tensor|int,y:Tensor): return x - y # some test has input as int
   def Div(x:Tensor,y:Tensor): return x.div(y, rounding_mode='trunc' if dtypes.is_int(x.dtype) else None)
+  # ONNX Pow is (T, T1) -> T, the output takes the base dtype while Tensor.pow promotes base and exponent
+  def Pow(x:Tensor,y:Tensor): return x.pow(y).round().cast(x.dtype) if dtypes.is_int(x.dtype) else x.pow(y)
   def Less(x:Tensor,y:Tensor): return x < y
   def LessOrEqual(x:Tensor,y:Tensor): return x <= y
   def Greater(x:Tensor,y:Tensor): return x > y
@@ -1297,7 +1299,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   return {
     # Tensor ops
-    **{op: getattr(Tensor, op.lower()) for op in ("Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan",
+    **{op: getattr(Tensor, op.lower()) for op in ("Neg", "Reciprocal", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan",
     "Asin", "Acos", "Atan", "Relu", "Sigmoid", "MatMul", "Floor", "Ceil", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh",
     "Tanh", "Softsign", "Asinh", "Acosh", "Atanh", "Elu", "Celu", "Selu", "Round", "Erf")},
     # Implemented ops

tinygrad/nn/state.py

@@ -70,6 +70,7 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=No
   nn.state.safe_save({'t':t}, "test.safetensor")
   ```
   """
+  if any(v.dtype in dtypes.weaks for v in tensors.values()): raise ValueError("safe_save requires concrete dtypes")
   headers, offset = {}, 0
   if metadata: headers['__metadata__'] = metadata
   for k,v in tensors.items():

tinygrad/renderer/__init__.py

@@ -39,9 +39,10 @@ class Estimates:
           mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.scalar().itemsize)
       if u.op is Ops.RANGE:
         mult_stack.append(mults)
-        mults *= cast(sint, u.src[0].ssimplify())
-        # SPECIAL are already counted in mults
-        mults = mults.substitute({x:x.const_like(0) for x in mults.toposort() if x.op is Ops.SPECIAL}) if isinstance(mults, UOp) else mults
+        if u.dtype is not dtypes.void:  # unbounded loop, unknown trip count
+          mults *= cast(sint, u.src[0].ssimplify())
+          # SPECIAL are already counted in mults
+          mults = mults.substitute({x:x.const_like(0) for x in mults.toposort() if x.op is Ops.SPECIAL}) if isinstance(mults, UOp) else mults
       elif u.op is Ops.END: mults = mult_stack.pop(-1)
       elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these
       elif u.op is Ops.PARAM and u.arg.addrspace == AddrSpace.ALU and u.expr == 'core_id': mults *= int(u.vmax) + 1

tinygrad/renderer/cstyle.py

@@ -12,9 +12,11 @@ base_rewrite = PatternMatcher([
   # local/reg buffers
   (UPat(Ops.BUFFER, name="x"), lambda ctx,x: ctx.render_buffer(x)),
 
-  # range/if/endif
+  # range/loop/if/endif
+  (UPat(Ops.RANGE, dtypes.void, name="x"), lambda ctx,x: "for (;;) {"),
   (UPat(Ops.RANGE, name="x"),
    lambda ctx,x: f"for ({ctx.render_dtype(x.dtype)} {ctx[x]} = 0; {ctx[x]} < {ctx[x.src[0]]}; {ctx[x]}++) {{"),
+  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE), UPat(name="c", dtype=dtypes.bool))), lambda ctx,c: f"  if (!({ctx[c]})) {{ break; }}\n}}"),
   (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"),
   (UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"),
 
@@ -99,7 +101,8 @@ def uops_to_dtypes(uops:list[UOp]) -> list[tuple[DType, int]]:
   return dedup((u.dtype, u.max_numel()) for u in uops if u.addrspace in (AddrSpace.ALU, None) and u.dtype != dtypes.void and u._shape is not None)
 
 def _wmma_name(u:UOp) -> str:
-  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.scalar().name}"
+  # sanitize spaces in DType.name (int8 = "signed char")
+  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.scalar().name}".replace(" ", "_")
 
 # (name, dims, dtype_in, dtype_out, device, threads, upcast_sizes)
 def wmma_args(uops:list[UOp]):
@@ -226,14 +229,14 @@ class CStyleLanguage(Renderer):
 
       if u.op in {Ops.ENDIF, Ops.END}: depth -= 1
       if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
-        (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG) or \
+        (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \
         (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
         (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
         r[u] = l
       else:
         if u.op not in {Ops.RANGE, Ops.STORE, Ops.BUFFER} and u.dtype != dtypes.void:
           l = f"{self.render_type(u)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "")
-        kernel.append("  "*depth + l)
+        kernel.append("\n".join("  "*depth + line for line in l.split("\n")))
         if prefix: c[prefix] += 1  # if it was used, increment
       if u.op in {Ops.IF, Ops.RANGE}: depth += 1
     del self.r
@@ -565,6 +568,11 @@ class HIPRenderer(CStyleLanguage):
       # #define __WMMA_16_16_16_half_half __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12
       elif self.tensor_cores == tc.amd_rdna4:
         prefix.append(f"#define __{name} __builtin_amdgcn_wmma_{type_map[dtype_out]}_16x16x16_{type_map[dtype_in]}_w32_gfx12")
+      elif dtype_out == dtypes.int32:
+        prefix.append("typedef int wmma_int4 __attribute__((ext_vector_type(4)));\n"+
+          f"static inline __attribute__((device)) int8 __{name}"+"""(signed_char16 a, signed_char16 b, int8 c) {
+  return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, __builtin_bit_cast(wmma_int4, a),
+    true, __builtin_bit_cast(wmma_int4, b), c, false);\n}""")
       elif dtype_out == dtypes.float:
         prefix.append(f"#define __{name} __builtin_amdgcn_wmma_f32_16x16x16_{'f16' if dtype_in == dtypes.half else 'bf16'}_w32")
       else: prefix.append(f"static inline __attribute__((device)) half8 __{name}"+"""(half16 a, half16 b, half8 c) {
@@ -589,3 +597,8 @@ class QCOMCLRenderer(OpenCLRenderer):
   def supported_dtypes(self):
     return {d for d in Renderer.supported_dtypes(self)
             if (d != dtypes.float16 or (bool(IMAGE) and bool(FLOAT16))) and d not in dtypes.fp8s+(dtypes.bfloat16,dtypes.double)}
+
+  # QCOM's load vectorizer emits invalid IR for vectorized bool loads ("Range types must match load type"), type bool buffers as uchar
+  def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.ALU, mutable=True, override_ptr=False, shape=None):
+    if dtype == dtypes.bool and addrspace == AddrSpace.GLOBAL: dtype = dtypes.uint8
+    return super()._render_dtype(dtype, sz, addrspace, mutable, override_ptr, shape)

tinygrad/renderer/isa/x86.py

@@ -86,7 +86,7 @@ class X86GroupOp:
                 X86Ops.CMPi, X86Ops.IMULi, X86Ops.LEA}
 
   # X86Ops whose second src can read from memory NOTE: some of these are TwoAddress so the second src is actually the first
-  ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.SHL, X86Ops.SHR, X86Ops.SAR, X86Ops.IMUL, X86Ops.CMP,
+  ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.IMUL, X86Ops.CMP,
                 X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VADDPS, X86Ops.VADDPD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VSUBPS, X86Ops.VSUBPD,
                 X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VMULPS, X86Ops.VMULPD, X86Ops.VDIVSS, X86Ops.VDIVSD, X86Ops.VDIVPS, X86Ops.VDIVPD,
                 X86Ops.VPADDB, X86Ops.VPADDW, X86Ops.VPADDD, X86Ops.VPADDQ, X86Ops.VPSUBB, X86Ops.VPSUBW, X86Ops.VPSUBD, X86Ops.VPSUBQ,
@@ -99,8 +99,9 @@ class X86GroupOp:
 
   # X86Ops that can write to memory
   WriteMem = {X86Ops.MOVm, X86Ops.MOVi, X86Ops.VMOVSSm, X86Ops.VMOVSDm, X86Ops.VMOVUPSm, X86Ops.VMOVDm, X86Ops.VMOVQm,
-              X86Ops.ADDi, X86Ops.SUBi, X86Ops.ANDi, X86Ops.ORi, X86Ops.XORi, X86Ops.SHLi, X86Ops.SHRi, X86Ops.SARi, X86Ops.SETNE,
-              X86Ops.SETE, X86Ops.SETL, X86Ops.SETB, X86Ops.VCVTPS2PH, X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ}
+              X86Ops.ADDi, X86Ops.SUBi, X86Ops.ANDi, X86Ops.ORi, X86Ops.XORi, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR,
+              X86Ops.SARi, X86Ops.SETNE, X86Ops.SETE, X86Ops.SETL, X86Ops.SETB,
+              X86Ops.VCVTPS2PH, X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ}
 
   # X86Ops that read flags
   ReadFlags = {X86Ops.CMOVB, X86Ops.CMOVL, X86Ops.CMOVE, X86Ops.CMOVNE, X86Ops.SETB, X86Ops.SETL, X86Ops.SETE, X86Ops.SETNE, X86Ops.JB, X86Ops.JL,
@@ -272,6 +273,11 @@ def idiv(ctx:IselContext, x:UOp) -> UOp:
   # this move "cleanses" the register constraints (rax/rdx) of idiv as that only applies on definition and not on the uses of idiv
   return x.ins(X86Ops.MOV, src=(idiv,))
 
+# a variable shift count implicitly reads cl so it goes in rcx, the shifted value can't be in rcx
+def shift(x:UOp, op:X86Ops) -> UOp:
+  val = x.ins(X86Ops.MOV, src=(x.src[0],), tag=tuple(r for r in WGPR if r is not RCX))
+  return x.ins(op, src=(val, x.ins(X86Ops.MOV, src=(x.src[1],), tag=(RCX,))))
+
 # a memory address operand is (base, index, displacement, size). size is the element size, it scales the index and is the memory operand width.
 # it is materialized as an immediate so the address stays correct if the base register is ever spilled and refilled
 def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]:
@@ -452,9 +458,9 @@ isel_matcher = PatternMatcher([
   (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.cvar("c"), lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar("c"))), lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   # scalar int binary with register
-  (UPat.var("a", dtypes.ints) << UPat.var("b"), lambda a,b: a.ins(X86Ops.SHL, src=(a, b))),
-  (UPat.var("a", dtypes.uints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SHR, src=(a, b))),
-  (UPat.var("a", dtypes.sints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SAR, src=(a, b))),
+  ((UPat(dtype=dtypes.ints) << UPat()).named("x"), lambda x: shift(x, X86Ops.SHL)),
+  ((UPat(dtype=dtypes.uints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SHR)),
+  ((UPat(dtype=dtypes.sints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SAR)),
   (UPat.var("a", dtypes.ints) + UPat.var("b"), lambda a,b: a.ins(X86Ops.ADD, src=(a, b))),
   (UPat.var("a", dtypes.ints) * UPat.var("b"), lambda a,b: a.ins(X86Ops.IMUL, src=(a, b))),
   (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.var("b"), lambda a,b: a.ins(X86Ops.AND, src=(a, b))),
@@ -651,7 +657,8 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
   if x.arg in X86GroupOp.WriteMem:
     if len(x.src) > 4: address, rest = x.src[:4], x.src[4:]
     else: address, rest = (x, None, None, None), x.src
-    return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *rest[:1]))
+    imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,)
+    return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *imm_uop))
 
   if x.arg in X86GroupOp.Rm1st:
     if len(x.src) > 3: address, rest = x.src[:4], x.src[4:]
@@ -704,8 +711,9 @@ encodings = {
   # int division
   X86Ops.IDIV: lambda x: encode(x, 0xF7, reg=7), X86Ops.DIV: lambda x: encode(x, 0xF7, reg=6),
   # scalar int binary
-  X86Ops.SHLi: lambda x: encode(x, 0xC1, reg=4),
-  X86Ops.SHRi: lambda x: encode(x, 0xC1, reg=5), X86Ops.SARi: lambda x: encode(x, 0xC1, reg=7),
+  X86Ops.SHL: lambda x: encode(x, 0xD3, reg=4), X86Ops.SHLi: lambda x: encode(x, 0xC1, reg=4),
+  X86Ops.SHR: lambda x: encode(x, 0xD3, reg=5), X86Ops.SHRi: lambda x: encode(x, 0xC1, reg=5),
+  X86Ops.SAR: lambda x: encode(x, 0xD3, reg=7), X86Ops.SARi: lambda x: encode(x, 0xC1, reg=7),
   X86Ops.ADD: lambda x: encode(x, 0x03), X86Ops.ADDi: lambda x: encode(x, 0x81, reg=0),
   X86Ops.SUB: lambda x: encode(x, 0x2B), X86Ops.SUBi: lambda x: encode(x, 0x81, reg=5),
   X86Ops.AND: lambda x: encode(x, 0x23), X86Ops.ANDi: lambda x: encode(x, 0x81, reg=4),
@@ -784,6 +792,7 @@ class X86Renderer(ISARenderer):
   post_regalloc_matcher = post_regalloc_matcher
   code_for_op = {x: lambda: None for x in (Ops.SQRT, Ops.AND, Ops.OR, Ops.SHL, Ops.SHR, Ops.NEG, Ops.SUB, Ops.FDIV, Ops.CMPLT, Ops.CMPEQ)}
   def __init__(self, target:Target):
+    if target.arch.split(",")[0] != "x86_64": raise RuntimeError(f"X86Renderer only supports x86_64, got {target.arch}")
     super().__init__(target)
     from tinygrad.runtime.support.compiler_cpu import X86Compiler
     self.compiler = X86Compiler()

tinygrad/renderer/llvmir.py

@@ -35,7 +35,7 @@ def lcast(input_type:DType, output_type:DType):
 
 def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
   dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.ushort: "bf16.1k" if cdna else "bf16", dtypes.bfloat16: "bf16.1k" if cdna else "bf16",
-            dtypes.fp8e4m3: ".fp8.fp8", dtypes.fp8e5m2: ".bf8.bf8"}
+            dtypes.fp8e4m3: ".fp8.fp8", dtypes.fp8e5m2: ".bf8.bf8", dtypes.int8: "iu8", dtypes.int32: "i32"}
   # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
   N,M,K = wmma.arg[0]
   if cdna:
@@ -44,9 +44,10 @@ def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
            f".{N}x{M}x{K}{dt_map[wmma.arg[1]]}(" + ", ".join([f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
   # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
   # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101)
+  args = [f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]
+  if wmma.arg[1] == dtypes.int8: args = ["i1 true", args[0], "i1 true", args[1], args[2]]  # iu8 flags A/B signed
   return f"  {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype]}.16x16x16." + \
-    f"{dt_map[wmma.src[0].dtype]}(" + ", ".join([f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]) + (", i1 false)" \
-      if wmma.dtype != dtypes.float else ")")
+    f"{dt_map[wmma.arg[1]]}(" + ", ".join(args) + (", i1 false)" if wmma.dtype != dtypes.float else ")")
 
 # llvm ops, lop[<dtype>][<op>]
 unsigned_lop = { Ops.ADD: "add", Ops.MUL: "mul", Ops.CDIV: "udiv", Ops.CMOD: "urem",
@@ -95,6 +96,11 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.WHERE, name="x"), lambda ctx,x:
    f"  {ctx[x]} = select {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}, {ldt(x.src[2].dtype)} {ctx[x.src[2]]}"),
 
+  # loop (a RANGE with no src is an unbounded loop header)
+  (UPat(Ops.RANGE, dtypes.void, name="l"), lambda ctx,l: f"  br label %loop_{ctx[l][1:]}\nloop_{ctx[l][1:]}:"),
+  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void, name="l"), UPat(name="c"))), lambda ctx,l,c:
+    f"  br i1 {ctx[c]}, label %loop_{ctx[l][1:]}, label %loop_exit_{ctx[l][1:]}\nloop_exit_{ctx[l][1:]}:"),
+
   # range
   (UPat(Ops.RANGE, name="r"), lambda ctx,r:
    f"  br label %loop_entry_{range_str(r)}\n"
@@ -261,6 +267,9 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
       ])
     if target.arch in {"gfx1100", "gfx1151"}:
       self.extra_matcher += PatternMatcher([
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
+          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)
             for j in range(x.max_numel()*2)))),

tinygrad/renderer/ptx.py

@@ -116,6 +116,9 @@ string_rewrite = PatternMatcher([
   # simple
   (UPat(Ops.BUFFER, name="x"), lambda ctx, x: [] if x.addrspace == AddrSpace.REG else [
     f".shared .align 16 .b8 local{x.arg.slot}[{x.max_numel()*x.dtype.itemsize}];", f"mov.u64 {ctx.r[x]}, local{x.arg.slot}[0];"]),
+  (UPat(Ops.RANGE, dtypes.void, name="l"), lambda ctx, l: f"WAITLOOP_{ctx.uops.index(l)}:"),
+  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void, name="l"), UPat(name="c"))), lambda ctx, l, c:
+    f"@{ctx.r[c]} bra WAITLOOP_{ctx.uops.index(l)};"),
   (UPat(Ops.RANGE, name="r"), lambda ctx, r: [
     f"mov.u32 {ctx.r[r]}, -1;",
     f"bra END_{ctx.r[r][1:]};",
@@ -211,6 +214,7 @@ class PTXRenderer(Renderer):
       prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.END: ("pred", "pred"), Ops.RANGE: ("ridx", None),
         Ops.CONST: ("const", None), Ops.BUFFER: ("local", "u64"), Ops.INDEX: ("bidx", "u64"), Ops.SHRINK: ("bidx", "u64"),
         Ops.PARAM: ("dat", "u64" if u.addrspace is AddrSpace.GLOBAL else None), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None))
+      if u.op is Ops.RANGE and u.dtype == dtypes.void: prefix = None  # loop headers don't have a register
       if prefix: r[u] = ssa(prefix, u, dtype)
 
       l: str|list[str]|None = string_rewrite.rewrite(u, ctx=self)

tinygrad/runtime/ops_python.py

@@ -48,7 +48,6 @@ class PythonProgram:
     st = time.perf_counter()
     warp = list(itertools.product(*[range(x) for x in local_size[::-1]]))
     warp_size = len(warp)
-    void_ops = {Ops.END, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.GROUP, Ops.STORE}
     for idxs in itertools.product(*[range(x) for x in global_size[::-1]]):
       values: dict[UOp, Any] = {}
       pbufs: list[memoryview] = list(bufs)
@@ -57,11 +56,15 @@ class PythonProgram:
       i = 0
       while i < len(self.uops):
         u = self.uops[i]
-        src_values = [values[v] for v in u.src if v.op not in void_ops]
-        src_dtypes = [v.dtype for v in u.src if v.op not in void_ops]
+        src_values = [values[v] for v in u.src if v.dtype is not dtypes.void]
+        src_dtypes = [v.dtype for v in u.src if v.dtype is not dtypes.void]
         if getenv("TRACE"): print(i, u.op, u.dtype, u.arg, src_values, src_dtypes)
         if u.op is Ops.END:
-          i = self.uop_to_index[u.src[1]]
+          if len(u.src) == 3:
+            # conditional backedge on a loop: jump back while the condition is true
+            if values[u.src[2]][0]: i = self.uop_to_index[u.src[1]]
+            else: i += 1
+          else: i = self.uop_to_index[u.src[1]]
           continue
         if u.op is Ops.IF:
           exec_masks.append([x and y for x,y in zip(exec_masks[-1], src_values[0])])
@@ -71,7 +74,7 @@ class PythonProgram:
           exec_masks.pop()
           i += 1
           continue
-        if u.op in (Ops.BARRIER, Ops.SINK, Ops.NOOP, Ops.GROUP):
+        if u.op in (Ops.BARRIER, Ops.SINK, Ops.NOOP, Ops.GROUP) or (u.op is Ops.RANGE and u.dtype == dtypes.void):
           # in the python emulator, the warp is always in sync
           i += 1
           continue

tinygrad/runtime/ops_qcom.py

@@ -278,7 +278,7 @@ class QCOMProgram(HCQProgram):
   def _parse_lib(self, lib):
     # Extract image binary
     self.image_size = _read_lib(lib, 0x100)
-    self.image = bytearray(lib[(image_offset:=_read_lib(lib, 0xc0)):image_offset+self.image_size])
+    self.image = lib[(image_offset:=_read_lib(lib, 0xc0)):image_offset+self.image_size]
 
     # Parse image descriptors
     image_desc_off = _read_lib(lib, 0x110)

tinygrad/schedule/indexing.py

@@ -1,8 +1,8 @@
 from typing import Iterator
 import functools, itertools
 from dataclasses import dataclass, field, replace
-from tinygrad.dtype import dtypes, AddrSpace, Invalid
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
+from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches, broadcast_axes
 from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
@@ -52,16 +52,22 @@ class IndexingContext:
   range_idx: Iterator[int] = field(default_factory=itertools.count)
   def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP) -> UOp:
     if isinstance(s, UOp) and s.op is Ops.RANGE: return s
-    # if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
-    return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
+    # if a range has a 1 src, it's the same as UOp.const(dtypes.weakint, 0)
+    return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.weakint, 0)
+
+def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
+  if x.op not in GroupOp.Broadcastable: return rngs
+  baxes, nleft = broadcast_axes(src.shape, x.shape), len(x.shape)-len(src.shape)
+  return tuple(r.const_like(0) if j in baxes else r for j,r in enumerate(rngs) if j >= nleft)
 
 def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
   new_srcs = []
   for i, s in enumerate(x.src):
     new_src = s
+    src_rngs = broadcast_rngs(x, s, ctx.range_map[x][0]) if x in ctx.range_map else ()
     # shape args of movement ops are at src[1:] and should not be indexed
     if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
-      if x in ctx.range_map and not (x.op in GroupOp.Movement and i > 0): new_src = new_src.index(*ctx.range_map[x][0])
+      if x in ctx.range_map and not (x.op in GroupOp.Movement and i > 0): new_src = new_src.index(*src_rngs)
     elif s in ctx.realize_map:
       realized_ranges = ctx.realize_map[s]
       assert isinstance(realized_ranges, list), "realize map must contain range list"
@@ -77,7 +83,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
         opts = BufferizeOpts(device=s.device, removable=removable) if len(ctx.range_map[s][1]) == len(realized_ranges) else \
                BufferizeOpts(device=s.device, addrspace=AddrSpace.LOCAL, removable=removable)
         new_src = UOp(Ops.STAGE, src=(new_src,)+closed_ranges, arg=opts)
-        if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][0]) if i in realized_ranges])
+        if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(src_rngs) if i in realized_ranges])
     new_srcs.append(new_src)
   return new_srcs
 
@@ -89,7 +95,7 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
   if x not in ctx.range_map: return None
   bx = create_bufferize_and_index_based_on_ranges(ctx, x)
   valid: UOp = UOp.const(dtypes.bool, True).uprod([r.get_valid() for r in ctx.range_map[x][0]])
-  return valid.where(bx.src[0], UOp.const(x.dtype, Invalid))
+  return valid.where(bx.src[0], UOp.const(x.dtype, 0))
 
 def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
   if x.arg[1] == 0: return None
@@ -136,7 +142,7 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U
   for s,src in list(zip(out_shape, urngs.src))[::-1]:
     axes_in.append(acc*src)
     acc *= s
-  combined_axes = UOp.const(dtypes.index, 0).usum(axes_in)
+  combined_axes = UOp.const(dtypes.weakint, 0).usum(axes_in)
   axes_out:list[UOp] = []
   for s in in_shape[::-1]:
     axes_out.append(combined_axes % s)
@@ -188,15 +194,21 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # treat MSTACK/MSELECT like SINK
     if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
 
-    if x.dtype == dtypes.index: continue  # TODO: why do I need this?
+    if x.dtype == dtypes.weakint: continue  # TODO: why do I need this?
     ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
+    # ranges the consumers iterate that this node broadcasts over
+    ended = [rctx.range_map[c][0][i] for c in consumer_map[x] if c in rctx.range_map and c.op in GroupOp.Broadcastable
+             for i in broadcast_axes(x.shape, c.shape)]
+    broadcast_ending_ranges = list(UOp.sink(*ended).ranges)
+    # fusion decision: REDUCE before the broadcast
+    if x.op is Ops.REDUCE: ending_ranges[x] += broadcast_ending_ranges
 
     # *** the ranges on the output are
     #  1. new if this op is realized
     #  2. from the single consumer if this op only has one consumer
     #  3. potentially new if this op has 2+ consumers
 
-    consumer_rngs = [rctx.range_map[c][0] for c in consumer_map[x] if c in rctx.range_map]
+    consumer_rngs = [broadcast_rngs(c, x, rctx.range_map[c][0]) for c in consumer_map[x] if c in rctx.range_map]
     if x in rctx.realize_map:
       # if this is in the realize_map, we create new ranges (at the output)
       out_rngs = tuple(rctx.new_range(s) for s in x.shape)
@@ -248,6 +260,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
       if len(_realize_axis):
         rctx.realize_map[x] = _realize_axis
         out_rngs = tuple([(rctx.new_range(x.shape[i]) if i in _realize_axis else r) for i,r in enumerate(out_rngs)])
+    ending_ranges[x] += broadcast_ending_ranges
 
     # TODO: some ops don't have shape, enable this after the `.st` property is removed
     #assert len(out_rngs) == len(x.shape), \

tinygrad/schedule/memory.py

@@ -56,7 +56,7 @@ def memory_plan_rewrite(linear:UOp, held_bufs:set[UOp]|None=None) -> UOp:
   arenas = {key: UOp.new_buffer(key[0], sz, dtypes.int8) for key, sz in arena_sizes.items()}
   replace_map:dict[UOp, UOp] = {}
   for buf_uop, offset in offsets.items():
-    replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(dtypes.index, offset)), buf_uop.max_numel())
+    replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(dtypes.weakint, offset)), buf_uop.max_numel())
 
   if DEBUG >= 1 and (omem:=sum(nbytes.values()) / 1e6) != (nmem:=sum(arena_sizes.values()) / 1e6):
     print(f"memory reduced from {omem:.2f} MB -> {nmem:.2f} MB, {len(first_appearance)} -> {len(arenas)} bufs")

tinygrad/schedule/multi.py

@@ -1,5 +1,5 @@
 from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST
-from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite
+from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite, broadcast_axes, _broadcast_shape
 from tinygrad.dtype import dtypes
 from tinygrad.schedule.allreduce import handle_allreduce
 
@@ -47,20 +47,17 @@ def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]:
   assert all_same(devices), f"all buffers must have the same device {devices}"
   dcount = len(devices[0])
 
+  out_shape = _broadcast_shape(*[x.shape for x in msrcs])
   srcs:list[UOp] = []
   for mlb in msrcs:
-    if mlb.axis is None:
-      # no axis, shard it
-      assert mlb.op is not Ops.MULTI
-      srcs.append(mlb._shard(axis, dcount))
+    src_axis = axis - (len(out_shape)-len(mlb.shape))
+    if mlb.axis == src_axis:
+      # same axis, just copy through
+      srcs.append(mlb.src[0])
     else:
-      assert mlb.op is Ops.MULTI
-      if mlb.axis == axis:
-        # same axis, just copy through
-        srcs.append(mlb.src[0])
-      else:
-        # axis mismatch, copy to all devices, and shard it correctly
-        srcs.append(copy_multi(mlb, mlb.device)._shard(axis, dcount))
+      # otherwise every device gets the full copy, sharded iff this src has the axis (broadcast srcs stay whole)
+      full = mlb if mlb.axis is None else copy_multi(mlb, mlb.device)
+      srcs.append(full if axis in broadcast_axes(mlb.shape, out_shape) else full._shard(src_axis, dcount))
   return srcs
 
 def alu_multi(root:UOp):
@@ -148,7 +145,7 @@ def rewrite_into_function(call:UOp):
 
 def param_to_multi(p:UOp):
   if p.axis is None: return None
-  return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.name, p.arg.addrspace).multi(p.axis)
+  return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.multiple_of, p.arg.name, p.arg.addrspace).multi(p.axis)
 
 # NOTE: this is the same pattern as unrolled ranges
 multi_pm = PatternMatcher([

tinygrad/schedule/rangeify.py

@@ -78,7 +78,7 @@ def split_reduceop(reduce:UOp, x:UOp):
   # split is moved to the end to provide maximum locality for the second phase reduce.
 
   # get expanded by rangeifying the UOp x
-  indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else UOp.const(dtypes.index, 0) for i,s in enumerate(x.shape)])
+  indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else UOp.const(dtypes.weakint, 0) for i,s in enumerate(x.shape)])
   range_nums = [y.arg[0] for y in indexed.substitute({x.base:UOp(Ops.NOOP, x.base.dtype)}, extra_pm=pm_mops).ranges]
   is_expanded = [i not in range_nums for i in range(len(x.shape))]
 
@@ -470,7 +470,7 @@ to_define_global = PatternMatcher([
   (UPat(Ops.STORE, name="x"), find_bufs),
   (UPat((Ops.BUFFER, Ops.MSTACK, Ops.MSELECT), name="buf"), debuf),
   (UPat(Ops.PARAM, name="v"), lambda v:
-   UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype)
+   UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype, multiple_of=v.arg.multiple_of)
    if v.arg.name is not None and v.arg.vmin_vmax is not None else None),
 
   # this renumbers the params

tinygrad/tensor.py

@@ -1,12 +1,12 @@
 # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
 from __future__ import annotations
 import time, functools, sys, inspect, pathlib, hashlib, weakref
-from typing import Any, Callable, cast, get_args, ParamSpec, TypeVar, Generic, TYPE_CHECKING
+from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
-from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst
-from tinygrad.helpers import all_int, getenv, fully_flatten, fetch, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey
+from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst
+from tinygrad.helpers import all_int, getenv, fully_flatten, fetch, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc
-from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape
+from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable
 from tinygrad.mixin.rand import RandMixin
 from tinygrad.schedule import create_linear_with_vars
 from tinygrad.device import Buffer, canonicalize_device
@@ -34,6 +34,8 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
 
 # **** Tensor helper functions ****
 
+def is_numpy_ndarray(x) -> "TypeGuard[numpy.ndarray]": return str(type(x)) == "<class 'numpy.ndarray'>"
+
 def _fromnp(x: 'numpy.ndarray') -> UOp:
   ret = UOp.new_buffer("NPY", x.size, _from_np_dtype(x.dtype))
   # fake realize
@@ -69,28 +71,27 @@ class Tensor(RandMixin):
 
     # create a UOp from the different types of inputs
     if isinstance(data, UOp):
-      # if data is dtype.index that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
-      if data.dtype == dtypes.index: data = _index_to_concrete_int(data)
+      # if data is dtype.weakint that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
+      if data.dtype == dtypes.weakint: data = _index_to_concrete_int(data)
     elif data is None:
       data = UOp.const(_dtype or dtypes.default_float, 0)
     elif isinstance(data, get_args(ConstType)):
       data = UOp.const(_dtype or dtypes.from_py(data), data)
-    elif isinstance(data, bytes): data = UOp._frompy(data, _dtype or dtypes.uint8, _device)
-    elif isinstance(data, (list, tuple)):
-      if _dtype is None:
-        if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): _dtype = dtypes.bool
-        else: _dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float  # NOTE: this works because all_int([True, False]) is True
-      data = UOp._frompy(data, _dtype, _device)
-    elif is_numpy_ndarray(data):
-      import numpy as np
-      assert isinstance(data, np.ndarray), f"expected np.ndarray, got {data}"
-      if data.shape == ():
-        data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item())
-      else:
+    elif is_numpy_ndarray(data) and data.shape == ():
+      data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item())
+    else:
+      if _dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {_dtype}")
+      if isinstance(data, bytes): data = UOp._frompy(data, _dtype or dtypes.uint8, _device)
+      elif isinstance(data, (list, tuple)):
+        if _dtype is None:
+          if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): _dtype = dtypes.bool
+          else: _dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float  # NOTE: this works because all_int([True, False]) is True
+        data = UOp._frompy(data, _dtype, _device)
+      elif is_numpy_ndarray(data):
         data = _fromnp(data.astype(npdtype) if _dtype is not None and (npdtype:=_to_np_dtype(_dtype)) is not None else data)
-    elif isinstance(data, pathlib.Path):
-      _dtype = _dtype or dtypes.uint8
-      data = UOp.new_buffer(f"DISK:{data.resolve()}", data.stat().st_size // _dtype.itemsize, _dtype)
+      elif isinstance(data, pathlib.Path):
+        _dtype = _dtype or dtypes.uint8
+        data = UOp.new_buffer(f"DISK:{data.resolve()}", data.stat().st_size // _dtype.itemsize, _dtype)
 
     # by this point, it has to be a UOp
     if not isinstance(data, UOp): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}")
@@ -178,6 +179,7 @@ class Tensor(RandMixin):
 
   def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
     """Creates the LINEAR UOp needed to realize these Tensor(s), with Variables."""
+    if any(t.dtype in dtypes.weaks for t in (self,)+lst): raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
     big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
     _apply_map_to_tensors(becomes_map, name="buffers")
     return create_linear_with_vars(big_sink)
@@ -205,14 +207,16 @@ class Tensor(RandMixin):
     return self
 
   def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor:
+    if self.dtype in dtypes.weaks: raise RuntimeError("cannot assign into a weak tensor; it has no storage")
     is_disk = isinstance(self.device, str) and self.device.startswith(("DISK", "TINYFS"))
     if not isinstance(x, Tensor): x = Tensor(x, device="CPU" if is_disk else self.device, dtype=self.dtype)
     if self.uop is x.uop: return self  # a self assign is a NOOP
     # broadcast x (shape only, dtype must match)
     x = x._broadcast_to(self.shape)
+    if x.dtype in dtypes.weaks: x = x.cast(least_upper_dtype(self.dtype, x.dtype))
+    if x.dtype != self.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
     if not is_disk and x.uop.device is not None and self.device is not None and self.device != x.device:
       raise RuntimeError(f"assign device mismatch {self.device} != {x.device}")
-    if not is_disk and self.dtype != x.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
     if isinstance(self.device, tuple) and x.uop.device is not None and self.uop.axis != x.uop.axis:
       raise RuntimeError(f"multi axis mismatch {self.uop.axis} != {x.uop.axis}")
 
@@ -253,6 +257,7 @@ class Tensor(RandMixin):
     print(np.frombuffer(t.data(), dtype=np.int32))
     ```
     """
+    if self.dtype in dtypes.weaks: return self.cast(strong_dtype(self.dtype)).data()
     if 0 in self.shape: return memoryview(bytearray(0)).cast(self.dtype.fmt)  # type: ignore[arg-type,return-value]
     assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
     buf = self._buffer()
@@ -276,6 +281,7 @@ class Tensor(RandMixin):
     print(t.tolist())
     ```
     """
+    if self.dtype in dtypes.weaks: return self.cast(strong_dtype(self.dtype)).tolist()
     # TODO: remove half once minimum python supports it
     if self.dtype in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s): return self.cast(dtypes.float32).tolist()
     if 0 in self.shape:
@@ -293,6 +299,7 @@ class Tensor(RandMixin):
     print(repr(t.numpy()))
     ```
     """
+    if self.dtype in dtypes.weaks: return self.cast(strong_dtype(self.dtype)).numpy()
     assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
     import numpy as np
     if self.dtype in { dtypes.bfloat16, *dtypes.fp8s }: return self.float().numpy()
@@ -450,7 +457,10 @@ class Tensor(RandMixin):
   def _rop(self, op:Ops, axis:tuple[int, ...]) -> Tensor: return self._apply_uop(UOp._rop, op=op, axis=axis)
 
   def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None:
-    if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
+    if self.dtype in dtypes.weaks: raise RuntimeError("cannot setitem into a weak tensor; it has no storage")
+    if isinstance(v, Tensor):
+      if v.dtype in dtypes.weaks: v = v.cast(least_upper_dtype(self.dtype, v.dtype))
+      if v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
     # raise if mutation would diverge from eager (allow only pure views of a realized buffer; exclude +=/-= RHS via v_uop/v_bw)
     v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {})
     if self.uop.op_in_backward_slice_with_self(Ops.BUFFER):
@@ -479,32 +489,6 @@ class Tensor(RandMixin):
   def __delitem__(self, indices) -> None:
     raise TypeError("Tensor does not support deleting items")
 
-  # ***** broadcasted elementwise ops *****
-
-  def where(self:Tensor, x:Tensor|ConstType|sint, y:Tensor|ConstType|sint) -> Tensor:
-    """
-    Returns a tensor of elements selected from either `x` or `y`, depending on `self`.
-    `output_i = x_i if self_i else y_i`.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    cond = Tensor([[True, True, False], [True, False, False]])
-    print(cond.where(1, 3).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    Tensor.manual_seed(42)
-    cond = Tensor.randn(2, 3)
-    print(cond.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print((cond > 0).where(cond, -float("inf")).numpy())
-    ```
-    """
-    if isinstance(x, Tensor): x, y = x._broadcasted(y)
-    elif isinstance(y, Tensor): y, x = y._broadcasted(x)
-    else: x, y = self.ufix(x)._broadcasted(y)
-    out_shape = _broadcast_shape(self.shape, x.shape)
-    return self.cast(dtypes.bool)._broadcast_to(out_shape)._apply_uop(UOp.where, x._broadcast_to(out_shape), y._broadcast_to(out_shape))
-
   # ***** op wrappers *****
 
   # unlike Tensors, UOps are immutable, so these don't go in mixin

tinygrad/uop/divandmod.py

@@ -11,6 +11,8 @@ def fold_divmod_general(d: UOp) -> UOp|None:
   if y.vmin==y.vmax==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.FLOORDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
   # x//y is constant
   if (xdiv:=x//y).vmin == xdiv.vmax: return x - xdiv.vmin*y if d.op is Ops.FLOORMOD else xdiv.const_like(xdiv.vmin)
+  # PARAM // c is irreducible
+  if x.op is Ops.PARAM and y.op is Ops.CONST and x.arg.multiple_of % y.arg == 0: return d.const_like(0) if d.op is Ops.FLOORMOD else None
 
   # split uops for the rest of the processing
   x_peeled, const = x.pop_const()
@@ -98,9 +100,9 @@ div_and_mod_symbolic = PatternMatcher([
   # (x//c+a)//d -> (x+a*c)//(c*d) for c>0, d>0
   ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d) if d.vmin>0 else None),
   # (x+c)//d -> (x+c%d)//d + c//d for d>0 (split out the multiple of d in the constant)
-  ((UPat.var("x", dtypes.index)+UPat.cvar("c"))//UPat.cvar("d"),
+  ((UPat.var("x", dtypes.weakint)+UPat.cvar("c"))//UPat.cvar("d"),
     lambda x,c,d: (x+c.arg%d.arg)//d + c.arg//d.arg if c.arg%d.arg!=c.arg and d.arg>0 else None),
 
   # ** 2. Slow Rules **
-  (UPat((Ops.FLOORDIV, Ops.FLOORMOD), dtypes.index, name="d"), lambda d: fold_divmod_general(d)),
+  (UPat((Ops.FLOORDIV, Ops.FLOORMOD), dtypes.weakint, name="d"), lambda d: fold_divmod_general(d)),
 ])

tinygrad/uop/movement.py

@@ -16,5 +16,11 @@ mop_cleanup = PatternMatcher([
    lambda src,stk: src if stk.shape == src.shape and list(range(len(stk.src))) == [x.src[1].arg for x in stk.src] else None),
   # const INDEX into STACK is src
   (UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="a"), UPat.cvar("i")), name="idx", allow_any_len=True),
-    lambda a,i,idx: a.src[i.arg] if len(idx.src) <= 2 else a.src[i.arg].index(*idx.src[2:])),
+   lambda a,i,idx: a.src[i.arg] if len(idx.src) <= 2 else a.src[i.arg].index(*idx.src[2:])),
+  # INDEX on INDEX is INDEX
+  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
+   lambda idx1,idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:]) if all(x.shape == () for x in idx1.src[1:]+idx2.src[1:]) else None),
+  # INDEX on shaped INDEX (TODO: this can be more generic)
+  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx1_arg"))),), allow_any_len=True, name="idx2"),
+   lambda buf,idx1_arg,idx2: buf.index(idx1_arg.index(*idx2.src[1:])) if len(idx1_arg.shape) == len(idx2.src[1:]) else None),
 ])