broadcast

backward through clone crashes now

.github/workflows/test.yml

@@ -1003,6 +1003,15 @@ jobs:
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
           python -m pytest -n=auto test/backend/test_ops.py --durations=20
+      - name: Run test_ops (IMAGE)
+        if: matrix.backend == 'ir3'
+        shell: bash
+        env:
+          IMAGE: 1
+          DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
+        run: |
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
   qcomclcompiletests:
     name: Compile-only (QCOM CL)
     runs-on: ubuntu-24.04-arm
@@ -1026,3 +1035,11 @@ jobs:
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
           python -m pytest -n=auto test/backend/test_ops.py --durations=20
+      - name: Run test_ops (IMAGE)
+        shell: bash
+        env:
+          IMAGE: 1
+          DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
+        run: |
+          DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20

examples/mlperf/model_train.py

@@ -1442,7 +1442,7 @@ def train_llama3():
 
   from tinygrad.nn.state import get_state_dict
   model_state = get_state_dict(model)
-  for wname in ["wqkv", "wo", "w13", "w2"]:
+  for wname in model._fp8_inv_scale:
     w = model_state[wname]
     w._inv_scale = model._fp8_inv_scale[wname]
     if optim.master_params:

examples/mlperf/models/flat_llama.py

@@ -105,13 +105,16 @@ class FlatTransformer:
     scaled_std = 0.02 / math.sqrt(2 * n_layers)
 
     # Attention
-    self._init_inv_scales = []  # populated by lin_per_layer
-    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
-    self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
+    self.wqkv, s_qkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
+    self.wo, s_o = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
 
     # FeedForward
-    self.w13 = self.lin_per_layer(dim, hidden_dim * 2)
-    self.w2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
+    if SPLIT_W13:
+      self.w1, s_1 = self.lin_per_layer(dim, hidden_dim)
+      self.w3, s_3 = self.lin_per_layer(dim, hidden_dim)
+    else:
+      self.w13, s_13 = self.lin_per_layer(dim, hidden_dim * 2)
+    self.w2, s_2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
 
     self.norm_eps = norm_eps
     self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
@@ -125,35 +128,34 @@ class FlatTransformer:
     self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
 
     def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().requires_grad_(False)
-    names = ["xqkv", "xo", "x13", "x2"]
+    names = ["xqkv", "xo", "x2"]
+    names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
     self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
-    grad_names = ["xqkv", "xo", "xw13", "xout"]
-    if SPLIT_W13: grad_names.append("xw3")
+    grad_names = ["xqkv", "xo", "xout"]
+    grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
     self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
-    w_names = ["wqkv", "wo", "w13", "w2"]
-    self._fp8_inv_scale = {wname: inv_scales.float().contiguous().requires_grad_(False)
-                           for wname, inv_scales in zip(w_names, self._init_inv_scales)}
-    del self._init_inv_scales
+    w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
+    w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
+    self._fp8_inv_scale = {name: s.float().contiguous().requires_grad_(False) for name, s in w_scales}
 
   def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
     if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
     else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
     amax = w.abs().flatten(1).max(1).detach()
     scale = FP8_MAX / (amax + 1e-8)
-    self._init_inv_scales.append((amax + 1e-8) / FP8_MAX)
-    return (w * scale.reshape(-1, 1, 1)).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE)
+    inv_scale = (amax + 1e-8) / FP8_MAX
+    return (w * scale.reshape(-1, 1, 1)).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale
 
-  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
+  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,
                 grad_amax_xqkv:Tensor, grad_amax_xo:Tensor):
     bsz, seqlen, _ = x.shape
-    new_amaxs, saves = [], []
+    amaxs, saves = [], []
 
-    xqkv, x_normed, rrms, ret = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
-                                                     amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv)
-    saves.extend([x_normed, rrms])
-    new_amaxs.extend(ret[:1])
-    saves.extend(ret[1:] + [xqkv])
+    xqkv, x_normed, rrms, (new_amax, *s) = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
+                                                                  amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv)
+    amaxs.append(new_amax)
+    saves.extend([x_normed, rrms, *s, xqkv])
     xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
     xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
     xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
@@ -170,65 +172,45 @@ 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, *ret = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo)
-    new_amaxs.extend(ret[:1])
-    saves.extend(ret[1:] + [out])
-    return (out, *new_amaxs, *saves)
+    out, new_amax, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo)
+    amaxs.append(new_amax)
+    saves.extend([*s, out])
+    return out, amaxs, saves
 
-  def feed_forward(self, x:Tensor, residual:Tensor, ffn_norm:Tensor, w13:Tensor, w2:Tensor,
-                   amax_x13:Tensor, amax_x2:Tensor, s_13:Tensor, s_2:Tensor,
-                   grad_amax_xw13:Tensor, grad_amax_xout:Tensor,
-                   w1:Tensor|None=None, w3:Tensor|None=None, grad_amax_xw3:Tensor|None=None):
-    new_amaxs, saves = [], []
+  def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
+    amaxs, saves = [], []
 
     if SPLIT_W13:
-      assert w1 is not None and w3 is not None and grad_amax_xw3 is not None
       h = x + residual
       x_normed, rrms = rmsnorm(h, self.norm_eps)
       saves.extend([x_normed, rrms])
-      inp = x_normed * ffn_norm
-      # separate w1 and w3 matmuls
-      x_w1, *ret1 = matmul(inp, w1, amax_x=amax_x13, w_inv_scale=s_13, grad_amax_state=grad_amax_xw13)
-      new_amaxs.extend(ret1[:1])
-      saves.extend(ret1[1:] + [x_w1])
-      x_w3, *ret3 = matmul(inp, w3, amax_x=amax_x13, w_inv_scale=s_13, grad_amax_state=grad_amax_xw3)
-      saves.extend(ret3[1:] + [x_w3])
-      # silu * mul + w2 matmul
-      out, *ret2 = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout)
-      new_amaxs.extend(ret2[:1])
-      saves.extend(ret2[1:] + [out])
-      return (out, h, *new_amaxs, *saves)
-
-    x_w13, h, x_normed, rrms, ret = add_norm_quantize_matmul(x, residual, ffn_norm, w13, s_13, self.norm_eps,
-                                                             amax_x=amax_x13, grad_amax_state=grad_amax_xw13)
-    saves.extend([x_normed, rrms])
-    new_amaxs.extend(ret[:1])
-    saves.extend(ret[1:] + [x_w13])
-
-    out, ret = silu_w13_quantize_matmul(x_w13, w2, s_2, amax_x2=amax_x2, grad_amax_xw13=grad_amax_xw13, grad_amax_xout=grad_amax_xout)
-    new_amaxs.extend(ret[:1])
-    saves.extend(ret[1:] + [out])
-    return (out, h, *new_amaxs, *saves)
+      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"])
+      amaxs.append(new_amax)
+      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"])
+      amaxs.append(new_amax)
+      saves.extend([*s, x_w3])
+      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"])
+      amaxs.append(new_amax)
+      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"],
+                                                                          self.norm_eps, amax_x=kwargs["amax_x13"],
+                                                                          grad_amax_state=kwargs["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"],
+                                                     grad_amax_xw13=kwargs["grad_amax_xw13"], grad_amax_xout=kwargs["grad_amax_xout"])
+      amaxs.append(new_amax)
+      saves.extend([*s, out])
+    return out, h, amaxs, saves
 
   @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor,
-                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                ffn_norm:Tensor, w13:Tensor, w2:Tensor,
-                amax_xqkv:Tensor, amax_xo:Tensor,
-                amax_x13:Tensor, amax_x2:Tensor,
-                s_qkv:Tensor, s_o:Tensor, s_13:Tensor, s_2:Tensor,
-                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor,
-                grad_amax_xw13:Tensor, grad_amax_xout:Tensor,
-                w1:Tensor|None=None, w3:Tensor|None=None, grad_amax_xw3:Tensor|None=None):
-    attn, *attn_ret = self.attention(x, freqs_cis, attention_norm, wqkv, wo,
-                                     amax_xqkv=amax_xqkv, amax_xo=amax_xo, s_qkv=s_qkv, s_o=s_o,
-                                     grad_amax_xqkv=grad_amax_xqkv, grad_amax_xo=grad_amax_xo)
-    attn_amaxs, attn_saves = attn_ret[:2], attn_ret[2:]
-    ffn, h, *ffn_ret = self.feed_forward(x, attn, ffn_norm, w13, w2,
-                                                   amax_x13=amax_x13, amax_x2=amax_x2, s_13=s_13, s_2=s_2,
-                                                   grad_amax_xw13=grad_amax_xw13, grad_amax_xout=grad_amax_xout,
-                                                   w1=w1, w3=w3, grad_amax_xw3=grad_amax_xw3)
-    ffn_amaxs, ffn_saves = ffn_ret[:2], ffn_ret[2:]
+  def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict):
+    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)
     h = h + ffn
     return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)
 
@@ -241,11 +223,10 @@ class FlatTransformer:
       self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
       self.wo.shard_(device, axis=2).realize()            # (n_layers, dim, in) shard in
       if SPLIT_W13:
-        self.w1 = self.w13[:, :self.hidden_dim, :].contiguous()
-        self.w3 = self.w13[:, self.hidden_dim:, :].contiguous()
         self.w1.shard_(device, axis=1).realize()
         self.w3.shard_(device, axis=1).realize()
-      self.w13.shard_(device, axis=1).realize()           # (n_layers, hidden*2, dim) shard out
+      else:
+        self.w13.shard_(device, axis=1).realize()           # (n_layers, hidden*2, dim) shard out
       self.w2.shard_(device, axis=2).realize()            # (n_layers, dim, hidden) shard in
       self.attention_norm.shard_(device, axis=None).realize()
       self.ffn_norm.shard_(device, axis=None).realize()
@@ -265,18 +246,19 @@ class FlatTransformer:
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
     a, ga, s = self._fp8_amax, self._fp8_grad_amax, self._fp8_inv_scale
     for i in range(self.n_layers):
-      split_kwargs = dict(w1=self.w1[i], w3=self.w3[i], grad_amax_xw3=ga["xw3"][i]) if SPLIT_W13 else {}
-      h, *ret = self.run_layer(h, freqs_cis,
-                               self.attention_norm[i], self.wqkv[i], self.wo[i],
-                               self.ffn_norm[i], self.w13[i], self.w2[i],
-                               amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i],
-                               amax_x13=a["x13"][i], amax_x2=a["x2"][i],
-                               s_qkv=s["wqkv"][i], s_o=s["wo"][i],
-                               s_13=s["w13"][i], s_2=s["w2"][i],
-                               grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
-                               grad_amax_xw13=ga["xw13"][i], grad_amax_xout=ga["xout"][i],
-                               **split_kwargs)
-      for name, new_val in zip(["xqkv", "xo", "x13", "x2"], ret[:5]):
+      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],
+                         grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["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])
+      if SPLIT_W13:
+        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
+                          s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["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])
+      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs)
+      amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
+      for name, new_val in zip(amax_names, ret[:len(amax_names)]):
         a[name][i].assign(new_val)
 
     logits = matmul(self.norm(h), self.output[0], fp8=False)[0]

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -18,6 +18,7 @@ export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh

@@ -16,7 +16,8 @@ export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
 export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -18,6 +18,7 @@ export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh

@@ -10,9 +10,19 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-0}
+export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-0}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0}
+export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh

@@ -19,6 +19,7 @@ export FP8=1
 export ALLREDUCE_CAST=1
 export FAST_CE=1
 export FUSED_INPUT_QUANTIZE=1
+export FUSED_GRAD_QUANTIZE=1
 export FUSED_ADD_NORM_MUL_QUANTIZE=1
 export FUSED_SILU_W13=1
 export FUSED_PAD_GRAD_ACCUM=1

extra/gemm/cdna_asm_gemm.py

@@ -2713,12 +2713,20 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
     gbase = gradient.base if hasattr(gradient, "base") else gradient
     mailbox_entry = _grad_fp8_mailbox.pop(gbase, None) or _grad_fp8_mailbox.pop(gradient, None)
     if mailbox_entry is not None:
-      g_fp8_u, inv_scale_u, _new_amax_u, store_effect = mailbox_entry
+      g_fp8_u, inv_scale_u = mailbox_entry
       g_fp8 = Tensor(g_fp8_u, device=a.device)[:a.shape[0]]
       g_scale = Tensor(inv_scale_u, device=a.device)
     else:
       assert grad_amax_state is not None, "fp8 matmul bwd needs either a mailbox entry or a grad_amax_state"
-      g_fp8, g_scale, _, store_effect = quantize_fp8_delayed(g_t, Tensor(grad_amax_state, device=a.device))
+      if getenv("FUSED_GRAD_QUANTIZE", 0):
+        g_fp8, g_scale, _, store_effect = quantize_fp8_delayed(g_t, Tensor(grad_amax_state, device=a.device))
+        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)
+      else:
+        grad_amax_t = Tensor(grad_amax_state, device=a.device)
+        g_fp8, g_scale, new_grad_amax = quantize_fp8(g_t, amax_state=grad_amax_t)
+        store_effect = grad_amax_state.store(new_grad_amax.uop)
+        g_fp8 = Tensor(g_fp8.contiguous().uop.after(store_effect), device=a.device)
     # dgrad: uses g_scale * x_scale * w_scale
     grad_a = asm_gemm(g_fp8, b_t, x_scale=g_scale * s_x_t, w_scale=s_w_t)
     # wgrad: no w_scale
@@ -2729,8 +2737,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
     else:
       g_fp8_T = g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1)
     grad_b = asm_gemm(g_fp8_T, a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t)
-    # Attach the delayed-amax store effect (if any) to grad_a so realizing grads commits the amax update.
-    ret = (None, grad_a.uop.after(store_effect), grad_b.uop, None, None)
+    ret = (None, grad_a.uop, grad_b.uop, None, None)
     if len(inputs) == 6: ret = ret + (None,)
     return ret
   else:

extra/hcq2/graph/hcq.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+import time
+from typing import cast
+from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer
+from tinygrad.dtype import dtypes
+from tinygrad.engine.jit import GraphRunner
+from tinygrad.engine.realize import get_call_outs_ins, get_runtime
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, graph_rewrite
+from extra.hcq2.hcq2 import HCQ2Compiled, HCQ2DeviceCtx, HCQ2LowerCtx, prep_runtime, pm_lower_kernargs, pm_lower_ops
+from extra.hcq2.hcq2 import pm_split_into_queues, pm_add_barriers, pm_add_signals, build_host_program
+
+# **************** insert deps ****************
+
+def insert_deps(ctx:HCQ2Graph, linear:UOp) -> UOp:
+  src = []
+  for j, call in enumerate(linear.src):
+    call = call.replace(tag=j)
+    _, _, bufs, _ = ctx.calls[j]
+    outs, ins = get_call_outs_ins(call)
+    deps = ctx._access_resources([bufs[i] for i in outs + ins], list(range(len(outs))), call)
+    src.append(UOp(Ops.AFTER, call.dtype, (call, *deps), tag=call.tag))
+  return linear.replace(src=tuple(src))
+pm_insert_deps = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), insert_deps)])
+
+def replace_params(ctx:HCQ2Graph, call:UOp) -> UOp|None:
+  if not any(x.op is Ops.PARAM for x in call.src[1:]): return None
+  return call.replace(src=tuple(ctx.input_addrs_uop[x.arg] if x.op is Ops.PARAM else x for x in call.src))
+pm_replace_params = PatternMatcher([(UPat(Ops.CALL, name="call", allow_any_len=True), replace_params)])
+
+# **************** graph-only passes ****************
+
+def alloc_queue_sig(ctx:HCQ2Graph, q:UOp) -> None:
+  if q.arg in ctx.queue_sigs: return None
+  buf = Buffer(q.arg[0], 0x100, dtypes.uint8, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+  ctx.queue_sig_bufs.append(buf)
+  ctx.queue_sigs[q.arg] = UOp.from_buffer(buf, q.arg[0])
+  return None
+pm_alloc_queue_sigs = PatternMatcher([(UPat(Ops.LINEAR, src=UPat({Ops.PROGRAM, Ops.COPY}), name="q"), alloc_queue_sig)])
+
+def lower_queue_deps(ctx:HCQ2Graph, after:UOp) -> UOp:
+  wrapper, deps, call_idx = after.src[0], after.src[1:], after.tag
+  def store(q_arg, v): return ctx.queue_sigs[q_arg].store(UOp.const(dtypes.uint32, v))
+  waits = tuple(UOp(Ops.WAIT, dtypes.void, (ctx.queue_sigs[dep.src[0].arg], UOp.const(dtypes.uint32, dep.tag),
+                                            store(dep.src[0].arg, dep.tag))) for dep in deps)
+  return wrapper.replace(src=tuple(q.replace(src=(*waits, *q.src, store(q.arg, call_idx))) for q in wrapper.src))
+pm_lower_queue_deps = PatternMatcher([(UPat(Ops.AFTER, src=UPat(Ops.LINEAR), name="after"), lower_queue_deps)])
+
+def optimize_queue_deps(ctx:HCQ2Graph, queue:UOp) -> UOp|None:
+  src, seen, pending, queue_sig = [], {}, {}, ctx.queue_sigs[queue.arg]
+  for x in queue.src:
+    if x.op is Ops.WAIT:
+      sig, val = x.src[0], x.src[1]
+      if sig is queue_sig or seen.get(sig, -1) >= val.arg: continue
+      if (old:=pending.get(sig)) is None or old.src[1].arg < val.arg: pending[sig] = x
+      continue
+    for wait in pending.values():
+      src.append(wait)
+      seen[wait.src[0]] = wait.src[1].arg
+    pending.clear()
+    src.append(x)
+  src += pending.values()
+  return queue.replace(src=tuple(src)) if tuple(src) != queue.src else None
+pm_optimize_queue_deps = PatternMatcher([
+  (UPat(Ops.LINEAR, src=UPat({Ops.BARRIER, Ops.WAIT, Ops.STORE, Ops.PROGRAM, Ops.COPY}), name="queue"), optimize_queue_deps),
+])
+
+def drop_dead_stores(ctx:HCQ2Graph, outer:UOp) -> UOp:
+  live = {u.src[2] for u in outer.toposort() if u.op is Ops.WAIT}
+  return outer.replace(src=tuple(q.replace(src=tuple(x for x in q.src if x.op is not Ops.STORE or x in live)) for q in outer.src))
+pm_drop_dead_stores = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"), drop_dead_stores)])
+
+def add_queue_sig_resets(ctx:HCQ2Graph, outer:UOp) -> UOp|None:
+  if not ctx.queue_sig_bufs: return None
+  resets = tuple(ctx.hcq_ctx.host_param(sig).index(UOp.const(dtypes.int, 0), ptr=True).cast(dtypes.uint64.ptr())
+                 .store(UOp.const(dtypes.uint64, 0)) for sig in ctx.queue_sig_bufs)
+  return outer.replace(src=tuple(c.replace(src=c.src + resets) if c.op is Ops.AFTER else c.after(*resets) for c in outer.src))
+pm_add_queue_sig_resets = PatternMatcher([(UPat(Ops.LINEAR, name="outer"), add_queue_sig_resets)])
+
+# **************** Graph ****************
+
+class HCQ2Graph(GraphRunner):
+  def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()):
+    super().__init__(linear, input_uops)
+    self.dev = cast(HCQ2Compiled, Device[self.device])
+    self.hcq_ctx = HCQ2LowerCtx(name="hcq_graph")
+
+    self.input_addrs = Buffer("CPU", max(len(input_uops), 1), dtypes.uint64, preallocate=True)
+    self.input_addrs_uop = self.hcq_ctx.host_param(self.input_addrs)
+
+    self.linear = graph_rewrite(self.linear, pm_insert_deps, ctx=self, name="hcq: insert deps", walk=True)
+    self.linear, sizes = prep_runtime(self.hcq_ctx, self.linear)
+    for dev_name, sz in sizes.items():
+      buf = Buffer(dev_name, sz, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)
+      self.hcq_ctx.devs[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(buf, dev_name), UOp.const(dtypes.uint64, buf._buf.va_addr))
+
+    self.linear = graph_rewrite(self.linear, pm_replace_params, ctx=self, name="hcq: replace params", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_lower_kernargs + pm_lower_ops, ctx=self.hcq_ctx, name="hcq: lower ops")
+
+    # per-queue signal state — populated as a side-effect by pm_alloc_queue_sigs walking the lowered linear.
+    self.queue_sig_bufs:list[Buffer] = []
+    self.queue_sigs:dict[tuple[str, str], UOp] = {}
+    graph_rewrite(self.linear, pm_alloc_queue_sigs, ctx=self, name="hcq: alloc queue sigs", walk=True)
+
+    self.linear = graph_rewrite(self.linear, pm_lower_queue_deps, ctx=self, name="hcq: lower queue deps")
+    self.linear = graph_rewrite(self.linear, pm_split_into_queues, ctx=self.hcq_ctx, name="hcq: split into queues")
+    self.linear = graph_rewrite(self.linear, pm_add_barriers, ctx=self.hcq_ctx, name="hcq: add barriers", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_optimize_queue_deps, ctx=self, name="hcq: optimize queue deps", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_drop_dead_stores, ctx=self, name="hcq: drop dead stores")
+    self.linear = graph_rewrite(self.linear, pm_add_signals, ctx=self.hcq_ctx, name="hcq: add signals", walk=True)
+    self.linear = graph_rewrite(self.linear, self.dev.pm_lower, ctx=self.hcq_ctx, name=f"hcq: encode cmdbuf {self.dev.device}", walk=True)
+    self.linear = graph_rewrite(self.linear, pm_add_queue_sig_resets, ctx=self, name="hcq: add queue sig resets", walk=True)
+    self.host_call = build_host_program(self.hcq_ctx, self.linear, None, self.dev)
+
+    self.host_rt, self.host_globals = get_runtime("CPU", self.host_call.src[0]), self.host_call.src[0].arg.globals
+
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None:
+    addrs = self.input_addrs.as_memoryview(force_zero_copy=True).cast('Q')
+    for i, u in enumerate(input_uops):
+      buf = next(b for b in u.buffer.bufs if b.device == self.dev.device) if isinstance(u.buffer, MultiBuffer) else u.buffer
+      addrs[i] = buf._buf.va_addr
+    self.host_rt(*[self.hcq_ctx.inputs[i].get_buf("CPU") for i in self.host_globals], vals=self.host_call.src[0].arg.vals(var_vals), wait=True)
+    if wait:
+      st = time.perf_counter()
+      self.dev.synchronize()
+      return time.perf_counter() - st
+    return None
+
+  @staticmethod
+  def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:
+    all_devs = GraphRunner._all_devs(batch_devs, new_call)
+    return new_call.src[0].op in (Ops.PROGRAM, Ops.COPY) and len(all_devs) == 1 and isinstance(all_devs[0], HCQ2Compiled)

extra/hcq2/hcq2.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any, TYPE_CHECKING
-import struct, functools, time, itertools
+import struct, functools, time, collections
 from dataclasses import replace
 if TYPE_CHECKING: from tinygrad.engine.realize import ExecContext
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, wait_cond, mv_address, round_up, DEBUG
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites
 from tinygrad.dtype import dtypes
@@ -11,7 +11,7 @@ from dataclasses import dataclass, field
 from tinygrad.runtime.support.memory import BumpAllocator
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.engine.realize import pm_flatten_linear, to_program, track_stats
+from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params
 
 HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 
@@ -25,7 +25,8 @@ class HCQ2Compiled(Compiled):
                kernargs_size=(16 << 20), can_recover:bool=False, arch=None):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
-    super().__init__(device, allocator, compilers, runtime, None, arch=arch)
+    from extra.hcq2.graph.hcq import HCQ2Graph
+    super().__init__(device, allocator, compilers, lambda *a, **kw: None, HCQ2Graph, arch=arch)
 
     self.kernargs_size = kernargs_size
     self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(kernargs_size, wrap=True)
@@ -52,7 +53,9 @@ class HCQ2Compiled(Compiled):
     if not hasattr(self, 'iface'): return
     sig = self.timeline_signal._buf.cpu_view().mv.cast('Q')
     tl = self.timeline_value.as_memoryview(force_zero_copy=True).cast('Q')
-    wait_cond(lambda: sig[0] >= tl[0] - 1, timeout_ms=3000, msg=f"{sig[0]} < {tl[0] - 1}")
+    st = time.perf_counter()
+    while sig[0] < tl[0] - 1:
+      if time.perf_counter() - st > (timeout or 3000) / 1000: self.on_device_hang()
 
   def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
 
@@ -139,38 +142,36 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
 
 # **************** lower context ****************
 
+@dataclass
+class HCQ2DeviceCtx:
+  device:str                       # device name; resolve to instance via Device[device]
+  kernargs_host:UOp                # UOp whose .buffer is dev.kernargs_buf (BUFFER UOp in runtime, PARAM in graph)
+  kernargs_gpu:UOp                 # va_addr const of dev.kernargs_buf
+  kernargs_allocator:BumpAllocator = field(default_factory=lambda: BumpAllocator(2 << 20, wrap=False))
+
 @dataclass
 class HCQ2LowerCtx:
-  dev:HCQ2Compiled
   name:str
-
-  kernargs_host:UOp|None = None
-  kernargs_gpu:UOp|None = None
-  kernargs_allocator:BumpAllocator = field(default_factory=lambda: BumpAllocator(0x1000, wrap=False))
-
-  timestamps_gpu:UOp|None = None
-  next_timestamp:itertools.count = field(default_factory=itertools.count)
-
   inputs:list[Buffer] = field(default_factory=list)
   holds:list[UOp] = field(default_factory=list)
+  devs:dict[str, HCQ2DeviceCtx] = field(default_factory=dict)
 
   def host_param(self, buf:Buffer) -> UOp:
     if buf not in self.inputs: self.inputs.append(buf)
     return UOp.placeholder((buf.size,), buf.dtype, self.inputs.index(buf))
 
 class HCQEncoder:
-  def __init__(self, ctx:HCQ2LowerCtx): self.ctx, self.dev, self.blob, self.patches, self.deps = ctx, ctx.dev, b'', [], set()
+  def __init__(self, ctx:HCQ2LowerCtx, dev:HCQ2Compiled): self.ctx, self.dev, self.blob, self.patches, self.deps = ctx, dev, b'', [], []
 
   @property
-  def src(self) -> tuple[UOp, ...]: return tuple(self.patches + list(self.deps))
+  def src(self) -> tuple[UOp, ...]: return tuple(self.patches + dedup(self.deps))
 
   def get_dev_addr(self, uop:UOp) -> sint|UOp:
-    # unwrap transient AFTER on the value: deps flow into enc.deps separately, the outer wrapper never reaches the final graph
     while uop.op is Ops.AFTER:
-      self.deps.update(uop.src[1:])
+      self.deps.extend(uop.src[1:])
       uop = uop.src[0]
-    self.deps.add(uop)
-    return uop.buffer.get_buf(self.dev.device).va_addr if uop.op in (Ops.BUFFER, Ops.BUFFER_VIEW) else uop.ssimplify()
+    if isinstance(val:=uop.ssimplify(), UOp): self.deps.append(uop)
+    return uop.buffer.get_buf(self.dev.device).va_addr if uop.op in (Ops.BUFFER, Ops.BUFFER_VIEW) else val
 
   def append(self, *data, dtype=dtypes.uint32):
     for d in data:
@@ -186,59 +187,81 @@ class HCQEncoder:
 
 pm_prep_runtime = PatternMatcher([
   # device-specific lowering of the program
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),),
-        name="call", allow_any_len=True), lambda ctx,call,prg: call.replace(src=(ctx.dev.pm_lower.rewrite(prg, ctx),) + call.src[1:])),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(), UPat(), UPat(Ops.BINARY)), name="p"),), name="c", allow_any_len=True),
+    lambda ctx, c, p: c.replace(src=(Device[p.src[1].arg].pm_lower.rewrite(p, ctx),) + c.src[1:])),
 ])
 
-# **************** lower hcq ****************
+def calc_kernargs_sizes(ctx:dict[str,int], u:UOp) -> None:
+  d = u.src[0].buffer.device
+  ctx[d] = ctx.get(d, 0) + round_up(u.arg[0].kernargs_alloc_size, 16)
+pm_calc_kernargs_sizes = PatternMatcher([(UPat(Ops.PROGRAM, name="u"), calc_kernargs_sizes)])
+
+# **************** lower kernargs ****************
 
 def lower_kernargs(ctx:HCQ2LowerCtx, call:UOp, prg:UOp) -> UOp:
   data, info = prg.arg
+  # after amd_build_program, prg.src is (BUFFER_lib_gpu,); the buffer's device names the device
+  dctx = ctx.devs[prg.src[0].buffer.device]
 
-  enc = HCQEncoder(ctx)
+  enc = HCQEncoder(ctx, Device[dctx.device])
   for gi in info.globals: enc.append(enc.get_dev_addr(call.src[1+gi]), dtype=dtypes.uint64)
   for v in info.vars: enc.append(v, dtype=dtypes.uint32)
 
-  args_off = ctx.kernargs_allocator.alloc(data.kernargs_alloc_size, 16)
-  assert ctx.kernargs_host is not None and ctx.kernargs_gpu is not None
-  ctx.kernargs_host.buffer.view(len(enc.blob), dtypes.uint8, args_off).ensure_allocated().as_memoryview(force_zero_copy=True)[:] = enc.blob
+  args_off = dctx.kernargs_allocator.alloc(data.kernargs_alloc_size, 16)
+  dctx.kernargs_host.buffer.view(len(enc.blob), dtypes.uint8, args_off).ensure_allocated().as_memoryview(force_zero_copy=True)[:] = enc.blob
 
-  args_uop = (ctx.kernargs_gpu + args_off).after(ctx.kernargs_host.after(*tuple(p.replace(arg=p.arg+args_off) for p in enc.patches)))
+  args_uop = (dctx.kernargs_gpu + args_off).after(dctx.kernargs_host.after(*tuple(p.replace(arg=p.arg+args_off) for p in enc.patches)))
   return call.replace(src=(prg.replace(src=prg.src + (args_uop,), arg=(data, info)),) + call.src[1:])
 
+pm_lower_kernargs = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER),), name="prg"),), name="call", allow_any_len=True), lower_kernargs),
+])
+
+# **************** lower ops ****************
+
 def lower_program(ctx:HCQ2LowerCtx, call:UOp, prg:UOp) -> UOp:
-  sig, tl = UOp.from_buffer(ctx.dev.timeline_signal), ctx.host_param(ctx.dev.timeline_value)
-  return UOp(Ops.LINEAR, dtypes.void, (
-    sig.wait(tl[0] - 1),
-    UOp(Ops.BARRIER, dtypes.void),
-    UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(ctx.timestamps_gpu + next(ctx.next_timestamp) * 8,), arg="timestamp"),
-    prg,
-    UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(ctx.timestamps_gpu + next(ctx.next_timestamp) * 8,), arg="timestamp"),
-    sig.store(tl[0])))
+  q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(prg.src[0].buffer.device, "COMPUTE"))
+  return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag)
 
 def lower_copy(ctx:HCQ2LowerCtx, call:UOp, copy:UOp) -> UOp:
-  dst, src, dev = call.src[1], call.src[2], ctx.dev
-  devs = [dev, src_dev] if (src_dev:=Device[src.device]) is not dev else [dev]
-  sigs_tls = [(UOp.from_buffer(d.timeline_signal), ctx.host_param(d.timeline_value)) for d in devs]
-  return UOp(Ops.LINEAR, dtypes.void, (
-    *[s.wait(t[0] - 1) for s,t in sigs_tls],
-    UOp(Ops.BARRIER, dtypes.void),
-    UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(ctx.timestamps_gpu + next(ctx.next_timestamp) * 8,), arg="timestamp"),
-    UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),
-    UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(ctx.timestamps_gpu + next(ctx.next_timestamp) * 8,), arg="timestamp"),
-    *[s.store(t[0]) for s,t in sigs_tls]))
+  dst, src = call.src[1], call.src[2]
+  q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(dst.buffer.device, "COPY"))
+  return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag)
 
-# lower to hcq-specific commands
-pm_hcq_lower = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER),), name="prg"),), name="call", allow_any_len=True), lower_kernargs),
+pm_lower_ops = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER), UPat()), name="prg"),), name="call", allow_any_len=True), lower_program),
   (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), lower_copy),
 ])
 
+# **************** split into queues ****************
+
+def split_into_queues(ctx:HCQ2LowerCtx, outer:UOp) -> UOp:
+  groups:dict[tuple, list[UOp]] = collections.defaultdict(list)
+  for child in outer.src:
+    wrapper = child.src[0] if child.op is Ops.AFTER else child
+    for q in wrapper.src: groups[q.arg].extend(q.src)
+  return outer.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, tuple(cmds), arg=k) for k, cmds in groups.items()))
+pm_split_into_queues = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR, src=UPat(Ops.LINEAR)).or_after(), name="outer"), split_into_queues)])
+
+# **************** add signals (runtime) ****************
+
+def add_signals(ctx:HCQ2LowerCtx, outer:UOp) -> UOp:
+  def wrap(q:UOp) -> UOp:
+    (dev_name, qname), devs = q.arg, {q.arg[0]} | {u.buffer.device for u in q.toposort() if u.op in (Ops.BUFFER, Ops.BUFFER_VIEW)}
+    sigs_tls = [(UOp.from_buffer(Device[d].timeline_signal), ctx.host_param(Device[d].timeline_value)) for d in sorted(devs) if d.startswith("AMD")]
+    return q.replace(src=(*(s.wait(t[0]-1) for s,t in sigs_tls), *q.src, *(s.store(t[0]) for s,t in sigs_tls)), arg=qname)
+  return outer.replace(src=tuple(wrap(q) for q in outer.src))
+
+pm_add_barriers = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"),
+  lambda ctx, outer: outer.replace(src=tuple(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)) for q in outer.src)))])
+
+pm_add_signals = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"), add_signals)])
+
 # **************** build host program ****************
 
 def resolve_cmdbuf(ctx:HCQ2LowerCtx, blob:UOp) -> UOp:
   inner = blob.src[0] if blob.op is Ops.AFTER else blob
+  dev_name, qtype = inner.tag
 
   # prepare the cmdbuf and make it a param
   bb = Buffer("CPU", len(inner.arg)//4, dtypes.uint32, preallocate=True)
@@ -246,10 +269,10 @@ def resolve_cmdbuf(ctx:HCQ2LowerCtx, blob:UOp) -> UOp:
   bb_param = ctx.host_param(bb)
 
   submit_cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(bb_param.after(*(blob.src[1:] if blob.op is Ops.AFTER else ())),),
-                  arg=f"submit_{inner.tag.lower()}")
+                  arg=f"submit_{qtype.lower()}", tag=dev_name)
 
   # increment the timeline value
-  tl = ctx.host_param(ctx.dev.timeline_value)
+  tl = ctx.host_param(Device[dev_name].timeline_value)
   return tl.after(UOp(Ops.BARRIER, dtypes.void, src=(submit_cf,))).index(UOp.const(dtypes.int, 0), ptr=True).store(tl[0] + 1)
 
 def resolve_patches(ctx:HCQ2LowerCtx, buf:UOp) -> UOp|None:
@@ -289,65 +312,66 @@ pm_resolve_ref_buffers = PatternMatcher([(UPat((Ops.BUFFER, Ops.BUFFER_VIEW), na
 
 pm_callify = PatternMatcher([(UPat(Ops.SINK, name="sink"), hcq_callify)])
 
-def hcq_build_host_program(ctx:HCQ2LowerCtx, linear:UOp, ast:UOp) -> UOp:
+# **************** schedule ****************
+
+def prep_runtime(ctx:HCQ2LowerCtx, linear:UOp) -> tuple[UOp, dict[str,int]]:
+  linear = graph_rewrite(linear, pm_prep_runtime, ctx=ctx, name="hcq: prepare runtime")
+  graph_rewrite(linear, pm_calc_kernargs_sizes, ctx=(sizes:={}), enter_calls=True)
+  return linear, sizes
+
+def build_host_program(ctx:HCQ2LowerCtx, linear:UOp, ast:UOp, dev:HCQ2Compiled) -> UOp:
   sink = graph_rewrite(linear, pm_create_host_sink, ctx=ctx, name="hcq: create host sink", walk=True)
   sink = graph_rewrite(sink, pm_lower_cmdbufs, ctx=ctx, bottom_up=True, name="hcq: lower cmdbufs")
   sink = graph_rewrite(sink, pm_resolve_patches, ctx=ctx, bottom_up=True, name="hcq: resolve patches")
   sink = graph_rewrite(sink, pm_resolve_ref_buffers, ctx=ctx, bottom_up=True, name="hcq: resolve ref buffers")
-  sink = graph_rewrite(sink, ctx.dev.pm_lower, ctx=ctx, name="hcq: device lower", walk=True)
+  sink = graph_rewrite(sink, dev.pm_lower, ctx=ctx, name=f"hcq: device lower {dev.device}", walk=True)
   return graph_rewrite(sink, pm_callify, ctx=ctx, name="hcq: callify")
 
-# **************** schedule ****************
+@track_rewrites(name=lambda ctx,linear,ast,dev,**kw: f"hcq schedule {getattr(ast.arg, 'name', ast.op.name.lower())}")
+def hcq_schedule(ctx:HCQ2LowerCtx, linear:UOp, ast:UOp, dev:HCQ2Compiled) -> UOp:
+  linear, sizes = prep_runtime(ctx, linear)
+  for dev_name, sz in sizes.items():
+    off = dev.kernargs_offset_allocator.alloc(sz, 16)
+    ctx.devs[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(dev.kernargs_buf.view(sz, dtypes.uint8, off), dev_name),
+                                       UOp.const(dtypes.uint64, dev.kernargs_buf.get_buf(dev_name).va_addr + off))
+  linear = graph_rewrite(linear, pm_lower_kernargs + pm_lower_ops, ctx=ctx, name="hcq: lower ops")
+  linear = graph_rewrite(linear, pm_split_into_queues, ctx=ctx, name="hcq: split into queues")
+  linear = graph_rewrite(linear, pm_add_barriers, ctx=ctx, name="hcq: add barriers", walk=True)
+  linear = graph_rewrite(linear, pm_add_signals, ctx=ctx, name="hcq: add signals", walk=True)
+  linear = graph_rewrite(linear, dev.pm_lower, ctx=ctx, name=f"hcq: encode cmdbuf {dev.device}", walk=True)
+  return build_host_program(ctx, linear, ast, dev)
 
-@track_rewrites(name=lambda dev,ctx,linear,ast,**kw: f"hcq schedule {getattr(ast.arg, 'name', ast.op.name.lower())}")
-def hcq_schedule(dev:HCQ2Compiled, ctx:HCQ2LowerCtx, linear:UOp, ast:UOp) -> UOp:
-  linear = graph_rewrite(linear, pm_prep_runtime, ctx=ctx, name="hcq: prepare runtime")
-  linear = graph_rewrite(linear, pm_hcq_lower + pm_flatten_linear, ctx=ctx, name="hcq: lower to cmdbuf ops")
-  linear = UOp(Ops.LINEAR, dtypes.void, (graph_rewrite(linear, dev.pm_lower, ctx=ctx, name="hcq: encode cmdbuf ops"),))
-  return hcq_build_host_program(ctx, linear, ast)
+def ensure_accessible(ctx:HCQ2LowerCtx, call:UOp, copy:UOp) -> UOp|None:
+  src_buf = call.src[2].buffer # TODO: cleanup
+  dev = call.src[1].buffer.device
+  try: src_buf.get_buf(dev)
+  except Exception:
+    (cpubuf := Buffer("CPU", src_buf.nbytes, dtypes.uint8, preallocate=True)).copyin(src_buf.ensure_allocated().as_memoryview())
+    ctx.holds.append(buf_uop:=UOp.from_buffer(cpubuf, dev))
+    return call.replace(src=call.src[:2] + (buf_uop,) + call.src[3:])
+pm_ensure_bufs_accessible = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), ensure_accessible)])
 
-def _resolve_call(ctx:ExecContext, call:UOp, ast:UOp) -> UOp:
-  from tinygrad.engine.realize import resolve_params
-  return call.replace(src=(ast,) + tuple(resolve_params(call, ctx.input_uops)) + tuple(s for s in call.src[1:] if s.op is Ops.BIND))
-
-def _run_host_call(ctx:ExecContext, call:UOp, dev:HCQ2Compiled, host_call:UOp, bufs:list[Buffer], ts_buf:Buffer) -> float:
+def hcq_exec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   from tinygrad.engine.realize import run_linear
+
+  if ast.src[1].arg.split(":")[0] != "AMD": return None
+
+  # TODO: this mess should gone
+  resolved_call = call.replace(src=(ast,) + tuple(resolve_params(call, ctx.input_uops)) + tuple(s for s in call.src[1:] if s.op is Ops.BIND))
+  bufs = [cast(Buffer, resolved_call.src[1+gi].buffer) for gi in ast.arg.globals] if ast.op is Ops.PROGRAM \
+    else [cast(Buffer, resolved_call.src[i].buffer) for i in range(1, len(resolved_call.src))]
+  dev = cast(HCQ2Compiled, Device[bufs[0].device])
+  hcq_ctx = HCQ2LowerCtx(name="submit")
+  linear = graph_rewrite(UOp(Ops.LINEAR, dtypes.void, (resolved_call,)), pm_ensure_bufs_accessible, ctx=hcq_ctx)
+  host_call = hcq_schedule(hcq_ctx, linear, ast, dev)
   with track_stats(ctx, call, dev.device, bufs, ctx.var_vals) as tm:
+    st = time.perf_counter() if ctx.wait else 0.0
     run_linear(UOp(Ops.LINEAR, dtypes.void, (host_call,)), var_vals=ctx.var_vals, jit=True, update_stats=DEBUG>=3)
     if ctx.wait:
       dev.synchronize()
-      tss = ts_buf._buf.cpu_view().mv.cast('Q')
-      tm[0] = (tss[1] - tss[0]) / dev.timestamp_divider / 1e6
+      tm[0] = time.perf_counter() - st
   return tm[0] if tm[0] is not None else 0.0
 
-def hcq_exec_program(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  if ast.src[1].arg.split(":")[0] != "AMD": return None
-  dev, resolved_call = Device[ast.src[1].arg], _resolve_call(ctx, call, ast)
-  hcq_ctx = HCQ2LowerCtx(dev=dev, name="submit_program",
-                         kernargs_host=UOp.from_buffer(dev.kernargs_buf, dev.device),
-                         kernargs_gpu=UOp.const(dtypes.uint64, dev.kernargs_buf.get_buf(dev.device).va_addr),
-                         kernargs_allocator=dev.kernargs_offset_allocator, # allocator is passed and it will rotate kernargs
-                         timestamps_gpu=UOp.const(dtypes.uint64, dev.timestamps_buf.get_buf(dev.device).va_addr))
-  host_call = hcq_schedule(dev, hcq_ctx, UOp(Ops.LINEAR, dtypes.void, (resolved_call,), arg="COMPUTE"), ast)
-  prg_bufs = [cast(Buffer, resolved_call.src[1+gi].buffer) for gi in ast.arg.globals]
-  return _run_host_call(ctx, call, dev, host_call, prg_bufs, ts_buf=dev.timestamps_buf)
-
-def hcq_exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  if ast.src[1].arg.split(":")[0] != "AMD": return None
-  dev, resolved_call = Device[ast.src[1].arg], _resolve_call(ctx, call, ast)
-  hcq_ctx = HCQ2LowerCtx(name="submit_copy", dev=dev, timestamps_gpu=UOp.const(dtypes.uint64, dev.timestamps_buf.get_buf(dev.device).va_addr))
-  src_buf = resolved_call.src[2].buffer
-  try: src_buf.get_buf(dev.device)
-  except Exception:
-    (cpubuf := Buffer("CPU", src_buf.nbytes, dtypes.uint8, preallocate=True)).copyin(src_buf.ensure_allocated().as_memoryview())
-    hcq_ctx.holds.append(buf_uop:=UOp.from_buffer(cpubuf, dev.device))
-    resolved_call = resolved_call.replace(src=resolved_call.src[:2] + (buf_uop,) + resolved_call.src[3:])
-  host_call = hcq_schedule(dev, hcq_ctx, UOp(Ops.LINEAR, dtypes.void, (resolved_call,), arg="COPY"), ast)
-  bufs = [cast(Buffer, resolved_call.src[1].buffer), cast(Buffer, resolved_call.src[2].buffer)]
-  return _run_host_call(ctx, call, dev, host_call, bufs, ts_buf=dev.timestamps_buf)
-
 pm_hcq_exec = PatternMatcher([
-  # TODO: use upat device=?
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), hcq_exec_program),
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), hcq_exec_copy),
+  (UPat(Ops.CALL, src=(UPat({Ops.PROGRAM, Ops.COPY}, name="ast"),), name="call", allow_any_len=True), hcq_exec),
 ])

extra/hcq2/ops_amd2.py

@@ -16,10 +16,9 @@ from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
-from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
+from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 from tinygrad.runtime.support.usb import USB3
 from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
-from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
 from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
@@ -29,8 +28,8 @@ from tinygrad.engine.realize import get_runtime
 from tinygrad.uop.ops import Ops, UPat, PatternMatcher, graph_rewrite
 
 class AMDComputeQueue(HCQEncoder):
-  def __init__(self, ctx:HCQ2LowerCtx):
-    super().__init__(ctx)
+  def __init__(self, ctx:HCQ2LowerCtx, dev:AMDDevice):
+    super().__init__(ctx, dev)
     self.pm4, self.gc, self.nbio, self.soc = self.dev.pm4, self.dev.gc, self.dev.nbio, self.dev.soc
 
   def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, len(vals) - 1), *vals)
@@ -142,13 +141,16 @@ amd_inner_pm = PatternMatcher([
 ])
 
 def amd_lower_pm4(ctx, linear):
-  enc = AMDComputeQueue(ctx)
+  prg = next(s for s in linear.src if s.op is Ops.PROGRAM)
+  dev = Device[prg.src[1].arg]
+  enc = AMDComputeQueue(ctx, dev)
   graph_rewrite(linear, amd_inner_pm, ctx=enc, name="amd: encode")
-  return UOp(Ops.BINARY, dtypes.void, arg=enc.blob).rtag("COMPUTE").after(*enc.src)
+  return UOp(Ops.BINARY, dtypes.void, arg=enc.blob).rtag((dev.device, "COMPUTE")).after(*enc.src)
 
 def amd_submit_pm4(ctx, cf):
+  dev = Device[cf.tag]
   bb_param = cf.src[0]
-  q = ctx.dev.compute_queue
+  q = dev.compute_queue
   ring, wptr, doorbell, put_ptr = (ctx.host_param(b) for b in (q.ring, q.write_ptr, q.doorbell, q.put_value))
   size, ring_dwords = UOp.const(dtypes.uint32, bb_param.dtype.size), q.ring.size
 
@@ -164,8 +166,8 @@ def amd_submit_pm4(ctx, cf):
   return doorbell.after(flush)[0].store(next_put)
 
 class AMDCopyQueue(HCQEncoder):
-  def __init__(self, ctx:HCQ2LowerCtx, queue_idx=0):
-    super().__init__(ctx)
+  def __init__(self, ctx:HCQ2LowerCtx, dev:AMDDevice, queue_idx=0):
+    super().__init__(ctx, dev)
     self.sdma, self.queue_idx, self.max_copy_size = self.dev.sdma, queue_idx, self.dev.max_copy_size
 
   def copy(self, x):
@@ -192,9 +194,11 @@ class AMDCopyQueue(HCQEncoder):
            *data64_le(self.get_dev_addr(x.src[0])))
 
 def amd_lower_sdma(ctx, linear):
-  enc = AMDCopyQueue(ctx)
+  copy = next(s for s in linear.src if s.op is Ops.COPY)
+  dev = Device[copy.src[0].buffer.device]
+  enc = AMDCopyQueue(ctx, dev)
   graph_rewrite(linear, amd_inner_sdma_pm, ctx=enc, name="amd: encode sdma")
-  return UOp(Ops.BINARY, dtypes.void, arg=enc.blob).rtag("COPY").after(*enc.src)
+  return UOp(Ops.BINARY, dtypes.void, arg=enc.blob).rtag((dev.device, "COPY")).after(*enc.src)
 
 amd_inner_sdma_pm = PatternMatcher([
   (UPat(Ops.WAIT,  name="x"), lambda ctx, x: ctx.wait(x)),
@@ -205,8 +209,9 @@ amd_inner_sdma_pm = PatternMatcher([
 ])
 
 def amd_submit_sdma(ctx, cf):
+  dev = Device[cf.tag]
   bb_param = cf.src[0]
-  q = ctx.dev.sdma_queue(0)
+  q = dev.sdma_queue(0)
   ring, wptr, doorbell, put_ptr = (ctx.host_param(b) for b in (q.ring, q.write_ptr, q.doorbell, q.put_value))
   size_dw, ring_bytes = bb_param.dtype.size, q.ring.size * 4
 
@@ -237,23 +242,24 @@ class AMDProgramData:
 _amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,Buffer]] = {}
 
 def amd_build_program(ctx:HCQ2LowerCtx, prg:UOp) -> UOp:
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[4].arg, ctx.dev.device))) is None:
+  dev = Device[prg.src[1].arg]
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[4].arg, dev.device))) is None:
     image, sections, relocs = elf_loader(lib)
     rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
     for off, sym, typ, addent in relocs:
       assert typ == 5, f"unknown AMD reloc {typ}"  # R_AMDGPU_REL64
       image[off:off+8] = struct.pack('<q', sym - off + addent)
-    lib_gpu = Buffer(ctx.dev.device, round_up(image.nbytes, 0x1000), dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
-    ctx.dev.allocator._copyin(lib_gpu._buf, image)
-    ctx.dev.synchronize()
+    lib_gpu = Buffer(dev.device, round_up(image.nbytes, 0x1000), dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
+    dev.allocator._copyin(lib_gpu._buf, image)
+    dev.synchronize()
     desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytes(image[rodata:rodata+ctypes.sizeof(amdgpu_kd.llvm_amdhsa_kernel_descriptor_t)]))
-    if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (ctx.dev.iface.props['lds_size_in_kb']*1024)//512:
+    if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
       raise RuntimeError("Too many resources requested: group_segment_size")
-    ctx.dev._ensure_has_local_memory(desc.private_segment_fixed_size)
+    dev._ensure_has_local_memory(desc.private_segment_fixed_size)
     edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
     cached = _amd_program_cache[key] = (AMDProgramData(
       entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
-      rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if ctx.dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
+      rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
       rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
       wave32=bool(desc.kernel_code_properties & 0x400),
       kernargs_segment_size=desc.kernarg_size,
@@ -262,7 +268,7 @@ def amd_build_program(ctx:HCQ2LowerCtx, prg:UOp) -> UOp:
       enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER,
     ), lib_gpu)
   data, lib_gpu = cached
-  return prg.replace(src=(UOp.from_buffer(lib_gpu, ctx.dev.device),), arg=(data, prg.arg))
+  return prg.replace(src=(UOp.from_buffer(lib_gpu, dev.device),), arg=(data, prg.arg))
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
@@ -284,29 +290,6 @@ class AMDQueueDesc:
   put_value: Buffer   # uint64[1]
   params: tuple|None = None  # setup_ring params for recovery
 
-  @property
-  def ring_mv(self) -> MMIOInterface: return self.ring._buf.view.view(fmt='I')
-  @property
-  def rptr_mv(self) -> MMIOInterface: return self.read_ptr._buf.view.view(fmt='Q')
-  @property
-  def wptr_mv(self) -> MMIOInterface: return self.write_ptr._buf.view.view(fmt='Q')
-  @property
-  def doorbell_mv(self) -> MMIOInterface: return self.doorbell._buf.view.view(fmt='Q')
-  @property
-  def put(self) -> int: return self.put_value._buf.view.view(fmt='Q')[0]
-  @put.setter
-  def put(self, v:int): self.put_value._buf.view.view(fmt='Q')[0] = v
-
-  def signal_doorbell(self, dev, doorbell_value:int|None=None):
-    try:
-      self.wptr_mv[0] = self.put
-      System.memory_barrier()
-      if dev.is_am() and not dev.is_usb(): dev.iface.dev_impl.gmc.flush_hdp()
-      self.doorbell_mv[0] = self.put if doorbell_value is None else doorbell_value
-    except Exception as e:
-      dev.error_state = e
-      raise
-
 class PCIIface(PCIIfaceBase):
   def __init__(self, dev, dev_id):
     super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0)),), vram_bar=0,
@@ -348,22 +331,22 @@ class PCIIface(PCIIfaceBase):
         eop_buffer.va_addr, eop_buffer.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
 
     ext = lambda addr,n,dt: Buffer("CPU", n, dt, options=BufferSpec(external_ptr=addr), preallocate=True)
+    (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True))._buf.view.view(fmt='Q')[0] = 0
     return AMDQueueDesc(ring=ext(ring.va_addr, ring.size//4, dtypes.uint32),
       doorbell=ext(self.dev_impl.doorbell64.addr + doorbell_index*8, 1, dtypes.uint64),
       read_ptr=ext(gart.va_addr+rptr, 1, dtypes.uint64), write_ptr=ext(gart.va_addr+wptr, 1, dtypes.uint64),
-      put_value=Buffer("CPU", 1, dtypes.uint64, preallocate=True), params=rcvr_params)
+      put_value=put_value, params=rcvr_params)
 
   def _collect_interrupts(self, reset=False, drain_only=False):
-    devs:list[AMDDevice] = [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, AMDDevice) and d.is_am()]
-    for d in devs:
-      if drain_only: d.iface.dev_impl.ih.drain()
-      else: d.iface.dev_impl.ih.interrupt_handler()
+    d = self.dev
+    if drain_only: d.iface.dev_impl.ih.drain()
+    else: d.iface.dev_impl.ih.interrupt_handler()
 
-      if reset and d.iface.dev_impl.recover(force=d.error_state is not None):
-        d.compute_queue.put = d.compute_queue.rptr_mv[0] = d.compute_queue.wptr_mv[0] = 0
-        d.iface.dev_impl.gfx.setup_ring(*d.compute_queue.params)
-        d.timeline_signal.value = d.timeline_value - 1
-        d.error_state = None
+    if reset and d.iface.dev_impl.recover():
+      cq = d.compute_queue
+      for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
+      d.iface.dev_impl.gfx.setup_ring(*cq.params)
+      d.timeline_signal._buf.cpu_view().mv.cast('Q')[0] = d.timeline_value.as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
 
   def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):

extra/llama_kernels/cast_amax/__init__.py

@@ -53,8 +53,10 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
   inv_scale = (grad_amax_state_t.float() + 1e-8) / FP8_MAX
   new_grad_amax = scalar_amax(grad_amax_buf)
   store_effect = grad_amax_state_t.uop.store(new_grad_amax.uop)
-  # Stash fp8 companion + amax store for cdna_asm_gemm's bwd to attach to grad_a.
-  _grad_fp8_mailbox[grad_xw13.uop] = (grad_xw13_fp8.uop, inv_scale.uop, new_grad_amax.uop, store_effect)
+  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, inv_scale.uop)
   return (None, None, grad_xw13.uop, None, None)
 
 def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor) -> tuple[Tensor, Tensor, Tensor]:

test/backend/test_ops.py

@@ -1060,10 +1060,17 @@ class TestOps(unittest.TestCase):
     helper_test_op([()], torch.erf, Tensor.erf)
 
   def test_gelu(self):
-    helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="tanh"), Tensor.gelu)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="tanh"), lambda x: Tensor.gelu(x, approximate="tanh"))
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="none"), lambda x: Tensor.gelu(x, approximate="none"))
   def test_gelu_extreme(self):
-    helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="tanh"), Tensor.gelu, low=300, high=400)
-    helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="tanh"), Tensor.gelu, low=-400, high=-300)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="tanh"), lambda x: Tensor.gelu(x, approximate="tanh"),
+                   low=300, high=400)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="tanh"), lambda x: Tensor.gelu(x, approximate="tanh"),
+                   low=-400, high=-300)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="none"), lambda x: Tensor.gelu(x, approximate="none"),
+                   low=300, high=400)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="none"), lambda x: Tensor.gelu(x, approximate="none"),
+                   low=-400, high=-300)
   def test_quick_gelu(self):
     helper_test_op([(45,65)], lambda x: x * torch.sigmoid(1.702 * x), Tensor.quick_gelu)
     helper_test_op([()], lambda x: x * torch.sigmoid(1.702 * x), Tensor.quick_gelu)

test/external/fuzz_fast_idiv.py

@@ -14,7 +14,7 @@ if __name__ == "__main__":
       print(f"Progress: {i}")
     dt = random.choice(dtypes.ints + tuple(dt.vec(4) for dt in dtypes.ints))
     u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
-    d = random.randint(1, max(1, u.arg[2]))
+    d = random.randint(1, max(1, u.arg[2])*2)
     if d in powers_of_two: continue
     expr = fast_idiv(DEV.target(Device.DEFAULT), u, d)
     if expr is None: continue

test/helpers.py

@@ -76,7 +76,7 @@ def timeit(fxn:Callable[..., T], *args, **kwargs) -> tuple[T, float]:
   ret = fxn(*args, **kwargs)
   return ret, (time.perf_counter_ns()-st)*1e-6
 
-def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
+def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple[int, ...]=()):
   allocator = Device['PYTHON'].allocator
   bufs = []
   for buf_dt, data in inputs or []:
@@ -85,7 +85,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
   g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=())
   prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
   prog = PythonProgram("run", PythonCompiler().compile(prg.src[3].arg))
-  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
+  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 
 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:

test/null/test_uops.py

@@ -2,13 +2,13 @@
 import unittest
 import numpy as np
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import Timing, Context
+from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, ConstFloat  # noqa: F401
 from tinygrad.device import Device
 from tinygrad.uop.ops import Ops, UOp, UPat, exec_alu
 from tinygrad.uop.spec import spec_shared
 from tinygrad.uop.symbolic import sym
-from test.helpers import to_uops_list
+from test.helpers import eval_uop, to_uops_list
 
 class TestSafeCast(unittest.TestCase):
   def test_cast_folds(self):
@@ -201,6 +201,7 @@ class TestFastIdiv(unittest.TestCase):
       self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
       self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
 
+  @Context(DISABLE_FAST_IDIV=0)
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
   def test_fast_idiv_and_mod(self):
     g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
@@ -220,6 +221,13 @@ class TestFastIdiv(unittest.TestCase):
     self.assertIn(Ops.SHR, ops)
     self.assertNotIn(Ops.CMOD, ops)
 
+  @Context(DISABLE_FAST_IDIV=0)
+  def test_fast_idiv_bounded_numerator_zero(self):
+    x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
+    for val in range(2):
+      self.assertEqual(eval_uop(x.alu(Ops.CDIV, x.const_like(3)), vals=(val,)), cdiv(val, 3))
+
+  @Context(DISABLE_FAST_IDIV=0)
   def test_fast_idiv_remove_powers_of_two(self):
     ridx = UOp.range(2**20, 0)
     uops = to_uops_list([ridx//(7*64)], ren=Device[Device.DEFAULT].renderer)

test/unit/test_gguf.py

@@ -115,33 +115,35 @@ class TestGGUF(unittest.TestCase):
     with self.assertRaises(ValueError):
       ggml_data_to_tensor(Tensor.empty(512, dtype=dtypes.uint8), 256, 1337)
 
-  def test_multi_part_load(self):
-    def build(n_total, part_no, tensors):
-      # [header] [kv_data] [tensor_infos] [padding] [tensor_data_blob]
-      buf = bytearray()
-      # Header: magic "GGUF" + version=3 + n_tensors + n_kv=2
-      buf += struct.pack("<4siqq", b"GGUF", 3, len(tensors), 2)
-      # KV entries: [key_len: uint64][key bytes][type: int32][value]
-      for k, v in [("split.count", n_total), ("split.no", part_no)]:
-        kb = k.encode()
-        buf += struct.pack("<Q", len(kb)) + kb + struct.pack("<i", 4) + struct.pack("<I", v)
-      data_off = 0
-      # Tensor infos: [name_len][name][ndims][dims reversed][qtype][offset_into_data_blob]
-      for name, dims, qtype, data in tensors:
-        nb = name.encode()
-        buf += struct.pack("<Q", len(nb)) + nb + struct.pack("<I", len(dims))
-        for d in reversed(dims): buf += struct.pack("<Q", d)
-        buf += struct.pack("<i", qtype) + struct.pack("<Q", data_off)
-        data_off += len(data)
-      buf += b"\x00" * ((32 - len(buf) % 32) % 32)
-      for _, _, _, data in tensors: buf += data
-      return bytes(buf)
+  @staticmethod
+  def _build_gguf(tensors, kvs):
+    # [header] [kv_data] [tensor_infos] [padding] [tensor_data_blob]
+    buf = bytearray()
+    # Header: magic "GGUF" + version=3 + n_tensors + n_kv
+    buf += struct.pack("<4siqq", b"GGUF", 3, len(tensors), len(kvs))
+    # KV entries: [key_len: uint64][key bytes][type: int32][value]
+    for k, v in kvs:
+      kb = k.encode()
+      if isinstance(v, str): buf += struct.pack("<Q", len(kb)) + kb + struct.pack("<i", 8) + struct.pack("<Q", len(v)) + v.encode()
+      else: buf += struct.pack("<Q", len(kb)) + kb + struct.pack("<i", 4) + struct.pack("<I", v)
+    data_off = 0
+    # Tensor infos: [name_len][name][ndims][dims reversed][qtype][offset_into_data_blob]
+    for name, dims, qtype, data in tensors:
+      nb = name.encode()
+      buf += struct.pack("<Q", len(nb)) + nb + struct.pack("<I", len(dims))
+      for d in reversed(dims): buf += struct.pack("<Q", d)
+      buf += struct.pack("<i", qtype) + struct.pack("<Q", data_off)
+      data_off += len(data)
+    buf += b"\x00" * ((32 - len(buf) % 32) % 32)
+    for _, _, _, data in tensors: buf += data
+    return bytes(buf)
 
+  def test_multi_part_load(self):
     with tempfile.TemporaryDirectory() as d:
       d = pathlib.Path(d)
       a, b = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32), np.array([5.0, 6.0], dtype=np.float32)
-      (d / "test-00001-of-00002.gguf").write_bytes(build(2, 0, [("a", (4,), 0, a.tobytes())]))
-      (d / "test-00002-of-00002.gguf").write_bytes(build(2, 1, [("b", (2,), 0, b.tobytes())]))
+      (d / "test-00001-of-00002.gguf").write_bytes(self._build_gguf([("a", (4,), 0, a.tobytes())], [("split.count", 2), ("split.no", 0)]))
+      (d / "test-00002-of-00002.gguf").write_bytes(self._build_gguf([("b", (2,), 0, b.tobytes())], [("split.count", 2), ("split.no", 1)]))
       kv, ts = gguf_load(d / "test-00001-of-00002.gguf")
       self.assertEqual(kv["split.count"], 2)
       np.testing.assert_equal(ts["a"].numpy(), a)

test/unit/test_gradient.py

@@ -69,6 +69,21 @@ class TestTensorGradient(unittest.TestCase):
     np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0])
     self.assertIs(x.grad, old_grad)
 
+  def test_gradient_through_clone(self):
+    src = Tensor([1.0, 2.0, 3.0, 4.0])
+    x = src.clone().requires_grad_(True)
+    (x * 2.0).sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0])
+    self.assertIsNone(src.grad)
+
+    src = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
+    x = src.clone().requires_grad_(True)
+    try:
+      (x * 2.0).sum().backward()
+    except RuntimeError:
+      # TODO: this crashes now
+      pass
+
   def test_gradient_through_chained_unrealized_setitem(self):
     g1 = Tensor.zeros(4).contiguous()
     g1[2] = Tensor(1.0)

test/unit/test_linalg.py

@@ -12,7 +12,6 @@ def reconstruction_helper(A:list[Tensor],B:Tensor, tolerance=1e-5):
   np.testing.assert_allclose(reconstructed_tensor.numpy(),B.numpy(),atol=tolerance,rtol=tolerance)
 
 class TestLinAlg(unittest.TestCase):
-  @unittest.skip("flaky on CI")
   def test_svd_general(self):
     sizes = [(2,2),(5,3),(3,5),(3,4,4),(2,2,2,2,3)]
     for size in sizes:

tinygrad/codegen/__init__.py

@@ -69,7 +69,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, pm_add_loads, name="** add loads (code)")
 
   # create image buffers
-  if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON"}:
+  if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}:
     sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True, ctx=ren.target.arch)
 
   # devectorize (TODO: does this need opts?)

tinygrad/gradient.py

@@ -14,6 +14,13 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
     return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
   if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
+def unbroadcast(ctx:UOp, shape:tuple|None) -> UOp:
+  if ctx._shape is None or shape is None or ctx.shape == shape: return ctx
+  if len(shape) > len(ctx.shape): raise RuntimeError(f"can't unbroadcast {ctx.shape} to {shape}")
+  aligned = (1,)*(len(ctx.shape)-len(shape)) + shape
+  axis = tuple(i for i,(s,n) in enumerate(zip(aligned, ctx.shape)) if s != n)
+  return ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, axis).cast(ctx.dtype).reshape(shape)
+
 def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]:
   """Remove unused PARAMs from body and return compacted (body, args)."""
   used = sorted({p.arg: p for p in body.toposort() if p.op is Ops.PARAM}.items())
@@ -66,9 +73,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(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n))
-     .cast(ctx.dtype), None)),
+  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (unbroadcast(ctx, ret.src[0]._shape), 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[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)),)),
@@ -114,6 +119,7 @@ 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
+      v = unbroadcast(v, k._shape)
       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/helpers.py

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

tinygrad/llm/gguf.py

@@ -12,6 +12,14 @@ def _ggml_iq_grid(device: str, grid: tuple[int, ...], grid_shape: tuple[int, int
   values = [float((w >> (8*i)) & 0xFF) for w in grid for i in range(grid_shape[1])]
   return Tensor(values, dtype=dtypes.float32, device=device).reshape(grid_shape)
 
+# native types {ggml_type: dtype}
+_GGML_NATIVE = {0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8, 25: dtypes.int16,
+                26: dtypes.int32, 27: dtypes.int64, 28: dtypes.float64, 30: dtypes.bfloat16}
+
+# quant types {ggml_type: (number of elements, number of bytes)}
+_GGML_QUANT = {2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
+               12:(256,144), 13:(256,176), 14:(256,210), 18:(256,98), 21:(256,110), 22:(256,82), 23:(256,136), 39:(32,17), 41:(128,18)}
+
 def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   """
   Converts ggml tensor data to a tinygrad tensor.
@@ -24,11 +32,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   """
   # https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356
 
-  # native types
-  if (dtype := {
-    0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8,
-    25: dtypes.int16, 26: dtypes.int32, 27: dtypes.int64, 28: dtypes.float64, 30: dtypes.bfloat16,
-  }.get(ggml_type)) is not None:
+  if (dtype := _GGML_NATIVE.get(ggml_type)) is not None:
     return t[:dtype.itemsize * n].contiguous().bitcast(dtype)
 
   def q_to_uint8(t: Tensor, b: int) -> Tensor:
@@ -36,12 +40,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
     shift_tensor, bitmask = Tensor.stack(*[ Tensor(2**(i*b), device=t.device, dtype=t.dtype) for i in range(8//b) ]), 0xff >> (8 - b)
     return t.unsqueeze(-1).expand((*t.shape,8//b)).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
 
-  # map to (number of elements, number of bytes)
-  if (nelements_nbytes := {
-    2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
-    12:(256,144), 13:(256,176), 14:(256,210), 18:(256,98), 21:(256,110), 22:(256,82), 23:(256,136), 39:(32,17),
-    41:(128,18)
-  }.get(ggml_type)) is not None:
+  if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None:
     from tinygrad.runtime.autogen import ggml_common as _ggml
     blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
     if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
@@ -132,6 +131,8 @@ readers: dict[int, Callable[[io.BufferedIOBase], Any]] = { 8: read_str, 9: read_
 read_uint32, read_int32, read_uint64, read_int64 = readers[4], readers[5], readers[10], readers[11]
 
 def _gguf_parse(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
+  # TODO: remove the need for copy to default device
+  tensor = tensor.to(None).realize()
   r = io.BufferedReader(TensorIO(tensor), 1_000_000)
   magic, version, n_tensors, n_kv = r.read(4), read_int32(r), read_int64(r), read_int64(r)
   if magic != b"GGUF" or version not in [2, 3]: raise ValueError("Invalid GGUF format!")
@@ -169,10 +170,8 @@ def gguf_load(fn: Tensor|str|pathlib.Path) -> tuple[dict, dict[str, Tensor]]:
 
   NOTE: The provided tensor must be on a device that supports execution.
   """
-  # TODO: remove the need for copy to default device
-  def load(p): return _gguf_parse(p if isinstance(p, Tensor) else Tensor(p).to(None).realize())
-  kv, sd = load(fn)
+  kv, sd = _gguf_parse(fn if isinstance(fn, Tensor) else Tensor(pathlib.Path(fn)))
   if kv.get('split.count', 1) <= 1: return kv, sd
   if isinstance(fn, Tensor): raise ValueError("multi-part GGUF requires a path argument (got Tensor)")
-  for pp in _gguf_split_paths(pathlib.Path(fn), kv)[1:]: sd.update(load(pp)[1])
+  for pp in _gguf_split_paths(pathlib.Path(fn), kv)[1:]: sd.update(_gguf_parse(Tensor(pp))[1])
   return kv, sd

tinygrad/llm/model.py

@@ -6,9 +6,9 @@ from tinygrad.llm.gguf import gguf_load
 from tinygrad.uop.ops import resolve
 
 @functools.cache
-def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
-  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
-  freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|None=None) -> Tensor:
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, device=device)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end, device=device).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
   return freqs.cos().cat(freqs.sin(), dim=-1).contiguous()
 
 class ExpertWeights:
@@ -27,9 +27,9 @@ def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
 
 def pairwise_topk(x: Tensor, k: int) -> tuple[Tensor, Tensor]:
   n = x.shape[-1]
-  vals = Tensor.arange(n).reshape(1,1,n).cast(x.dtype).expand(x.shape)
+  vals = Tensor.arange(n, device=x.device).reshape(1,1,n).cast(x.dtype).expand(x.shape)
   cmp = (x.unsqueeze(-1) > x.unsqueeze(-2)) | ((x.unsqueeze(-1) == x.unsqueeze(-2)) & \
-    (Tensor.arange(n).reshape(1,1,n,1) < Tensor.arange(n).reshape(1,1,1,n)))
+    (Tensor.arange(n, device=x.device).reshape(1,1,n,1) < Tensor.arange(n, device=x.device).reshape(1,1,1,n)))
   sel = Tensor.zeros_like(x).scatter(-1, cmp.sum(axis=-1).cast('int32'), vals)[:,:,n-k:].cast('int32')
   return x.gather(-1, sel), sel
 
@@ -186,7 +186,7 @@ class TransformerBlock(FFNBlock):
     if not hasattr(self, "cache_kv"):
       # TODO: how is the dtype of this determined?
       self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, device=x.device)
-      self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta)
+      self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class MLATransformerBlock(FFNBlock):
   def __init__(self, config:TransformerConfig):
@@ -232,7 +232,7 @@ class MLATransformerBlock(FFNBlock):
   def _init_state(self, x:Tensor):
     if not hasattr(self, "cache_k"):
       self.cache_k = Tensor.empty(x.shape[0], 1, self.config.max_context, self.config.kv_lora_rank + self.config.rope_dim, device=x.device)
-      self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta)
+      self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class GatedDeltaNetBlock(FFNBlock):
   def __init__(self, config:TransformerConfig, ssm:SSMConfig):

tinygrad/mixin/__init__.py

@@ -5,7 +5,7 @@ from tinygrad.mixin.elementwise import ElementwiseMixin
 from tinygrad.mixin.movement import MovementMixin
 from tinygrad.mixin.reduce import ReduceMixin
 from tinygrad.uop import Ops
-from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element
+from tinygrad.uop.ops import resolve, smax, smin, identity_element
 from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
 from tinygrad.helpers import all_int, argfix, ceildiv, flatten, flat_to_grouped, make_tuple, prod, resolve_pool_pads, round_up
 
@@ -306,11 +306,6 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
   def _broadcasted(self, y, reverse=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
     # ptr dtypes aren't in the promo lattice
     if x.dtype == y.dtype or any(isinstance(d, PtrDType) for d in (x.dtype, y.dtype)): return x, y
     return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype)
@@ -1472,30 +1467,30 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     #prepare round robin pairing: identity on first half, reversed on second half
     permute = type(self).arange(num//2, dtype=dtypes.int, device=self.device).cat(
                 type(self).arange(num//2, num, dtype=dtypes.int, device=self.device).flip(0))
-    cols = type(self).arange(num, dtype=dtypes.int, device=self.device)
+    cols, h = type(self).arange(num, dtype=dtypes.int, device=self.device), num // 2
+    eye_num = type(self).eye(num, dtype=self.dtype, device=self.device).expand(b_shape + (num, num))
     def one_round_jacobi(U, V, permute):
-      # permutation matrix P such that X @ P == X[..., permute]; X @ P.T applies the inverse permutation
+      # permutation matrix P with P[a,b] = (a == permute[b]); first 2h columns are paired-column selectors
       P = cols.unsqueeze(1).eq(permute.unsqueeze(0)).cast(U.dtype)
-      #pair all the columns
-      U_perm, V_perm = U @ P, V @ P
-      U_permuted, runoff_U = U_perm.split(num - 1, -1) if num % 2 == 1 else (U_perm, None)
-      V_permuted, runoff_V = V_perm.split(num - 1, -1) if num % 2 == 1 else (V_perm, None)
-      U_left, U_right = U_permuted.split(num//2, -1)
-      V_left, V_right = V_permuted.split(num//2, -1)
-      #compute the jacobi rotations for each pairing
-      gamma = (U_left * U_right).sum(-2).reshape(b_shape + (1, num//2))
-      alpha, beta = U_permuted.square().sum(-2).unsqueeze(-2).split(num//2, -1)
+      P_pair = P[..., :2*h]  # drops the runoff column for odd num
+      # extract paired columns to compute Jacobi rotation params
+      U_pair = U @ P_pair
+      U_left, U_right = U_pair.split(h, -1)
+      gamma = (U_left * U_right).sum(-2).reshape(b_shape + (1, h))
+      alpha, beta = U_pair.square().sum(-2).unsqueeze(-2).split(h, -1)
       rot = gamma.ne(0)
       tau = (beta - alpha) / (2 * rot.where(gamma, 1))
       t = tau.ne(0).where(tau.sign(), 1) / (tau.abs() + (1 + tau.square()).sqrt())
       t = rot.where(t, 0)
       c = 1 / (1 + t.square()).sqrt()
       s = c * t
-      #apply the rotations and unpermute via P.T
-      U_left, U_right = c * U_left - s * U_right, s * U_left + c * U_right
-      U = U_left.cat(U_right.cat(runoff_U, dim=-1) if num % 2 == 1 else U_right, dim=-1) @ P.transpose(-2, -1)
-      V_left, V_right = c * V_left - s * V_right, s * V_left + c * V_right
-      V = V_left.cat(V_right.cat(runoff_V, dim=-1) if num % 2 == 1 else V_right, dim=-1) @ P.transpose(-2, -1)
+      # build rotation matrix R: identity + sum over pairs of 2x2 rotation deltas at (i_k, j_k) positions
+      Mi, Mj = P_pair.transpose(-2, -1).split(h, -2)  # paired-column selectors, each shape (h, num)
+      Mi_a, Mi_b = Mi.unsqueeze(-1), Mi.unsqueeze(-2)
+      Mj_a, Mj_b = Mj.unsqueeze(-1), Mj.unsqueeze(-2)
+      cc, ss = (c - 1).reshape(b_shape + (h, 1, 1)), s.reshape(b_shape + (h, 1, 1))
+      R = eye_num + (cc * (Mi_a * Mi_b + Mj_a * Mj_b) + ss * (Mi_a * Mj_b - Mj_a * Mi_b)).sum(-3)
+      U, V = U @ R, V @ R
       #prepare the next round robin pairings
       if num % 2 == 1: permute = (permute - 1) % num
       else: permute = permute[0].reshape(1).cat(((permute[1:num] - 2) % (num - 1)) + 1)

tinygrad/mixin/elementwise.py

@@ -750,7 +750,7 @@ class ElementwiseMixin(DTypeMixin, CreationMixin):
     """
     return self * (self * 1.702).sigmoid()
 
-  def gelu(self) -> Self:
+  def gelu(self, approximate:str="tanh") -> Self:
     """
     Applies the Gaussian Error Linear Unit (GELU) function element-wise.
 
@@ -760,7 +760,12 @@ class ElementwiseMixin(DTypeMixin, CreationMixin):
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).gelu().numpy())
     ```
     """
-    return 0.5 * self * (1 + (math.sqrt(2 / math.pi) * (self + 0.044715 * self ** 3)).tanh())
+    if approximate == "tanh":
+      return 0.5 * self * (1 + (math.sqrt(2 / math.pi) * (self + 0.044715 * self ** 3)).tanh())
+    elif approximate == "none":
+      return self * 0.5 * (1.0 + (self / math.sqrt(2)).erf())
+    else:
+      raise RuntimeError(f"{approximate=} is not supported")
 
   def swish(self) -> Self:
     """

tinygrad/nn/onnx.py

@@ -617,7 +617,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def softmax_13(x:Tensor, axis:int=-1): return x.softmax(axis)
   Softmax = {OpSetId(Domain.ONNX, 1):softmax_1, OpSetId(Domain.ONNX, 13):softmax_13}
   def HardSigmoid(x:Tensor, alpha:float=0.2, beta:float=0.5): return (alpha*x + beta).clip(0, 1)
-  def Gelu(x:Tensor, approximate:str|None=None): return x.gelu() if approximate == "tanh" else 0.5 * x * (1 + (x/math.sqrt(2)).erf())
+  def Gelu(x:Tensor, approximate:str|None=None): return x.gelu(approximate="none" if approximate is None else approximate)
   def BiasGelu(x: Tensor, bias: Tensor, approximate: str | None = None) -> Tensor: return Gelu(x + bias, approximate)
   def FastGelu(x:Tensor, bias:Tensor|None=None): return (x + bias).gelu() if bias is not None else x.gelu() # this is tanh approximated
   def PRelu(X:Tensor, slope:Tensor): return (X > 0).where(X, X * slope)

tinygrad/runtime/support/compiler_cpu.py

@@ -6,8 +6,8 @@ from tinygrad.runtime.autogen import llvm
 
 class ClangJITCompiler(Compiler):
   def __init__(self, arch:list[str], cachekey="compile_clang_jit"):
+    assert len(arch) >= 2, f"invalid arch string: {','.join(arch)!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
     self.arch, cpu, *feats = arch
-    assert self.arch and cpu, f"invalid arch string: {arch!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
     match self.arch:
       case "x86_64": self.args = [f"-march={cpu}"] + [f"-mno{f}" if f.startswith("-") else f"-m{f}" for f in feats]
       # on arm march means "runs on this arch and superset" instead of "optimize for this arch". x86 march == arm mcpu
@@ -92,8 +92,8 @@ class LLVMCompiler(Compiler):
 
 class CPULLVMCompiler(LLVMCompiler):
   def __init__(self, arch:list[str], cache_key=None):
+    assert len(arch) >= 2, f"invalid arch string: {','.join(arch)!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
     self.arch, cpu, *feats = arch
-    assert self.arch and cpu, f"invalid arch string: {arch!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
     featstr = ','.join(f if f.startswith('-') else '+'+f for f in feats)
     if cpu == "native":
       cpu = ctypes.string_at(llvm.LLVMGetHostCPUName()).decode()

tinygrad/schedule/indexing.py

@@ -144,11 +144,18 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
     case _: raise RuntimeError(f"{op} is not a MovementOp")
   return rngs
 
+pm_do_broadcast = PatternMatcher([
+  (UPat(GroupOp.Broadcastable, name="x"), lambda x: x.replace(src=tuple(y._broadcast_to(x.shape) for y in x.src))),
+])
+
 @profile_matches
 def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   if debug: print("**************************")
   rctx = IndexingContext()
 
+  # run broadcasting
+  tsink = graph_rewrite(tsink, pm_do_broadcast, name="do broadcast")
+
   # get ops to realize
   graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize")
 

tinygrad/tensor.py

@@ -1146,12 +1146,13 @@ class Tensor(OpMixin):
     0x8000000000008002, 0x8000000000000080, 0x800a, 0x800000008000000a, 0x8000000080008081, 0x8000000000008080, 0x80000001, 0x8000000080008008)]
 
     rate, dsbyte = {"sha3_224": (144, 6), "sha3_256": (136, 6), "shake_128": (168, 31)}[cfg] if isinstance(cfg, str) else cfg
-    data, data_pad = self.bitcast(dtypes.uint8).reshape(prod(self.shape[:-1]), self.shape[-1]), rate - (self.shape[-1] * self.dtype.itemsize % rate)
+    data = self.bitcast(dtypes.uint8).reshape(prod(self.shape[:-1]), self.shape[-1])
+    data_pad = rate - data.shape[-1] % rate
     # pad batches then pad blocks
-    data = data.pad((None, (0, data_pad))).reshape(bs := data.shape[0], -1, rate).pad((None, None, (0, 200 - rate)))
+    data = data.pad((None, (0, data_pad))).reshape(bs := data.shape[0], -1, rate).pad_to(None, None, 200)
 
     # create pad mask
-    lbe = prod(data.shape[1:]) + rate - data_pad - 200
+    lbe = (data.shape[1] - 1) * 200 + rate - data_pad
     if data_pad == 1: mb = [(lbe, 0), (1, dsbyte ^ 0x80), (200 - rate, 0)]
     else: mb = [(lbe, 0), (1, dsbyte), (data_pad - 2, 0), (1, 0x80), (200 - rate, 0)]
     pad_mask = Tensor.cat(*(Tensor(v, dtype=dtypes.uint8, device=data.device).expand(l) for l, v in mb if l > 0)).unsqueeze(0)
@@ -1160,7 +1161,7 @@ class Tensor(OpMixin):
 
     state = Tensor.zeros(bs, 25, device=self.device, dtype=dtypes.uint64)
     for k in range(int(data.shape[1])):
-      state = state ^ data.shrink((None, (k, k+1), None)).squeeze(1)
+      state = state ^ data[:, k]
       for i in range(24): # f1600
         # θ step
         p = state.reshape(bs, 5, 5).transpose(2, 1)
@@ -1179,11 +1180,7 @@ class Tensor(OpMixin):
     assert self.dtype == dtypes.uint8, "only support uint8 tensors for hashing"
     assert self.ndim == 2, "only support batched 1d tensors"
     assert self.shape[1] == 1024 * 1024, "only support messages of 1mb"
-
-    blocks = self.shape[0] * self.shape[1] // 4096
-    data = self.reshape(blocks, 4096)
-    block_hashes = data.keccak("shake_128").reshape(self.shape[0], 4096)
-    return block_hashes.keccak("shake_128").reshape(self.shape[0], 16)
+    return self.reshape(-1, 4096).keccak("shake_128").reshape(self.shape[0], -1).keccak("shake_128")
 
   def hash(self) -> Tensor:
     """
@@ -1193,19 +1190,14 @@ class Tensor(OpMixin):
     print(t.data().hex())
     ```
     """
-
     data = self.flatten().bitcast(dtypes.uint8)
-    if (tsize := data.shape[0]) % 2**20 != 0: data = data.pad((0, 2**20 - tsize % 2**20))
-    base_chunks = ceildiv(data.shape[0], 2**20)
-    tree_depth = math.ceil(math.log(base_chunks, 65536)) if base_chunks > 1 else 0
-
-    level_chunks = base_chunks
-    for _ in range(tree_depth + 1):
-      data = data.reshape(level_chunks, 2**20)._hash_1mb().flatten()
-      if (tsize := data.shape[0]) % 2**20 != 0: data = data.pad((0, 2**20 - tsize % 2**20))
-      level_chunks = ceildiv(data.shape[0], 2**20)
-
-    return data[:16]
+    n = data.shape[0]
+    assert isinstance(n, int), "hash requires concrete shape"
+    chunks = ceildiv(n, 2**20)
+    while chunks > 1:
+      data = data.pad_to(chunks * 2**20).reshape(chunks, 2**20)._hash_1mb().flatten()
+      chunks = ceildiv(chunks, 65536)
+    return data.pad_to(2**20).unsqueeze(0)._hash_1mb().flatten()[:16]
 
   # ***** processing ops *****
 
@@ -1477,7 +1469,7 @@ class Tensor(OpMixin):
   def image_conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding=0, dtype=None) -> Tensor:
     dtsz = 2 if FLOAT16 else 4
 
-    (bs,_,iy,ix), (cout,cin,H,W) = self.shape, weight.shape
+    (bs,_,_,_), (cout,cin,H,W) = self.shape, weight.shape
     assert isinstance(cin, int) and isinstance(cout, int)
     x, w = self, weight.reshape(groups, (rcout := cout//groups), cin, H, W)
 
@@ -1513,7 +1505,6 @@ class Tensor(OpMixin):
     def ipad(t, i, amt):
       shape = (None,)*i + (amt,) + (None,)*(t.ndim-i-1)
       return Tensor(True, device=t.device).expand(t.shape).pad_to(shape).where(t.pad_to(shape), Invalid) if amt != t.shape[i] else t
-
     # align a dimension, use at to specify the dimension to pad in, defaults to first
     def pad_align(t, dim, at=None, force=False):
       # align to 64 pixels when height is real, otherwise 64 bytes is sufficient
@@ -1531,7 +1522,7 @@ class Tensor(OpMixin):
     else: x, w = x.contiguous(), w.contiguous()
 
     # undo alignment hacks
-    if bank_conflict: x, w = x[:, :, :ix, :, :cin // 4, :], w[:, :H, :cin // 4, ...]
+    if bank_conflict: x, w = x[:, :, :, :, :cin // 4, :], w[:, :, :cin // 4, ...]
     else: x, w = x[:, :, :ix, :], w[:, :H, ...]
 
     # expand out
@@ -1554,13 +1545,9 @@ class Tensor(OpMixin):
     # the conv!
     ret = (x*w).cast(dtypes.float32).sum((-4, -3, -2, -1), dtype=dtype)
 
-    if added_ox:
-      ret = ret.reshape(bs, oy, ox + added_ox, groups, rcout)[:, :, :ox, ...]
-
+    ret = ret.reshape(bs, oy, ox + added_ox, groups, rcout)[:, :, :ox, :, :]
     # undo hack for non multiples of 4 on C.rcout
-    if added_output_channels:
-      ret = ret.reshape(bs, oy, ox, groups, rcout)[:, :, :, :, :-added_output_channels]
-
+    if added_output_channels: ret = ret[:, :, :, :, :-added_output_channels]
     # NCHW output
     ret = ret.reshape(bs, oy, ox, groups * (rcout - added_output_channels)).permute(0,3,1,2)
     return ret if bias is None else ret.add(bias.reshape(1, -1, 1, 1))

tinygrad/uop/__init__.py

@@ -118,6 +118,9 @@ class GroupOp:
   # TODO: is BITCAST always Elementwise if it's shape changing?
   Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST})
 
+  # all ops that support shape broadcasting
+  Broadcastable = set.union(Elementwise, {Ops.CAST, Ops.GROUP, Ops.STORE})
+
   Defines = {Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}
 
   Irreducible = {Ops.CONST, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.RANGE}

tinygrad/uop/decompositions.py

@@ -286,6 +286,7 @@ def fast_idiv(target: Target, x: UOp, d: int, dont_cast=False) -> UOp|None:
   is_unsigned = x.vmin>=0 or x.dtype in dtypes.uints
   assert d>0, "Sign should have been taken out of divisor"
   vmin,vmax = max(x.vmin, x.dtype.min), min(x.vmax, x.dtype.max)
+  if vmin > -d and vmax < d: return x.const_like(0)
   m,s = magicgu(max(vmax, abs(vmin)), d)
   if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
     return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)

tinygrad/uop/ops.py

@@ -51,7 +51,10 @@ def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
   max_dim = max(len(s) for s in shapes)
   return tuple((1,)*(max_dim-len(s))+s for s in shapes)
 def _broadcast_shape(*shapes:tuple[sint, ...]) -> tuple[sint, ...]:
-  return tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*_align_left(*shapes)))
+  ret = tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*_align_left(*shapes)))
+  if not all(resolve(s == ns) or resolve(s == 1) for shape in _align_left(*shapes) for s,ns in zip(shape, ret)):
+    raise ValueError(f"shape mismatch: objects cannot be broadcast to a single shape {shapes}")
+  return ret
 
 def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop
 def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop
@@ -213,6 +216,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION:
         return None
 
+      # hacks for NOOP
+      case Ops.NOOP:
+        return self.src[0]._shape if len(self.src) >= 1 else None
+
       case Ops.GETTUPLE:
         # GETTUPLE extracts from a TUPLE (possibly through a FUNCTION)
         in_tuple = self.src[0].src[0] if self.src[0].op is Ops.FUNCTION else self.src[0]
@@ -258,7 +265,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.WMMA | Ops.SHAPED_WMMA: return self.src[2]._shape
 
       # passthrough ops
-      case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.PATCH | Ops.LOAD:
+      case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.PATCH | Ops.LOAD | \
+           Ops.COPY | Ops.ALLREDUCE:
         return self.src[0]._shape
       # REDUCE with empty axis is passthrough (lowered form)
       case Ops.REDUCE if len(self.arg[1]) == 0:
@@ -312,11 +320,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
           return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))
 
     # elementwise ops keep the shape the same. all inputs with shape must match
-    if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE, Ops.STORE}):
-      input_shapes = [x._shape for x in self.src if x._shape is not None]
-      if len(input_shapes) == 0: return None
-      if not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes} {[x.op for x in self.src]}")
-      return input_shapes[0]
+    if self.op in GroupOp.Broadcastable:
+      input_shapes = [x._shape for x in self.src]
+      assert len(self.src) > 0 and all(x is not None for x in input_shapes), f"None input shape not supported for {self.op}"
+      return _broadcast_shape(*input_shapes)
 
     # all Ops must be explicitly handled
     raise NotImplementedError(f"no shape handling for {self.op} with {self.dtype}")
@@ -476,10 +483,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return UOp(Ops.CONTRACT, dtype=self.dtype.vec(prod([x.vmax+1 for x in rngs])), src=(self,), arg=tuple((x.arg[0], x.vmax+1) for x in rngs))
   def alu(self, op, *src:UOp, **kwargs):
     all_srcs = (self, *src)
-    # broadcast shaped operands to a common shape (None and () are falsy, so only real shapes participate)
-    if (shapes := [s for x in all_srcs if (s:=x._shape)]) and not all_same(shapes):
-      out_shape = _broadcast_shape(*shapes)
-      all_srcs = tuple(x._broadcast_to(out_shape) if x._shape else x for x in all_srcs)
     out_dtype = all_srcs[-1].dtype
     if op in {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool
     return UOp(op, out_dtype, all_srcs, **kwargs)

tinygrad/viz/cli.py

@@ -68,7 +68,12 @@ def main(args) -> None:
     data = viz.get_render(viz_data, step["query"])
     if isinstance(data.get("value"), Iterator):
       for m in data["value"]:
-        if "uop" in m: print(emit(m["graph"] if print_graph else m["uop"]))
+        if print_graph and "graph" in m and not args.json:
+          for k,v in m["graph"].items():
+            print(f"[{k}] {' '.join((lines:=v['label'].splitlines())[:5])}{'...' if len(lines) > 5 else ''}"+(f" tag={v['tag']}" if v['tag'] else ''))
+            if v["src"]:
+              print("  src: "+", ".join([f"{i}->[{x}]" for i,x in v["src"][:5]])+(f", ... and {len(v['src'])-5} more" if len(v["src"]) > 5 else ""))
+        elif "uop" in m: print(emit(m["graph"] if print_graph else m["uop"]))
         if not reconstruct_matches: return None
         if m.get("diff"):
           loc = pathlib.Path(m["upat"][0][0])
@@ -194,8 +199,8 @@ def main(args) -> None:
           if DEBUG >= 3 and s["name"] == "View Base AST": print_step(s)
           if DEBUG >= 4 and s["name"] == "View Source": print_step(s)
           if DEBUG >= 5 or ls: print(emit(" "*s["depth"]+s["name"]+(f" - {s['match_count']}" if s.get('match_count', 0) else '')))
-          if DEBUG >= 6 or (DEBUG >= 5 and s["name"] == "View Kernel Graph"): print_step(s, print_graph=True)
-          if DEBUG >= 7 or s["name"] in args.src: print_step(s, reconstruct_matches=True)
+          if DEBUG >= 6 or (DEBUG >= 5 and s["name"] == "View Kernel Graph") or (s["name"] in args.src): print_step(s, print_graph=True)
+          if DEBUG >= 7: print_step(s, reconstruct_matches=True)
       elif DEBUG >= 3 and k.get("ext"): print(emit(k["ext"]))
     for k in (produce_top_kernels if args.t else produce_all_kernels)(): render_event(k)