Merge branch 'master' into default_late_allreduce

examples/mlperf/model_train.py

@@ -1344,7 +1344,6 @@ def train_llama3():
   vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
 
   model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
-
   params = get_parameters(model)
   # weights are all bfloat16 for now
   assert params and all(p.dtype == dtypes.bfloat16 for p in params)
@@ -1382,20 +1381,16 @@ def train_llama3():
 
     vocab_mask.shard_(device, axis=2).realize()
 
-  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
-  is_fake_offload = Device.DEFAULT == "NULL"
-  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
+  optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
   optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
                            eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
 
   # init grads
-  if is_offload_optim:
-    for p in optim.params:
-      p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=optim_device, requires_grad=False).contiguous().realize()
-  else:
-    for p in optim.params:
-      p.grad = p.zeros_like().contiguous().realize()
+  for p in optim.params:
+    p.grad = p.empty_like().realize()
   grads: list[Tensor] = [p.grad for p in optim.params]
+  for p in optim.params:
+    p.grad.assign(p.grad.zeros_like()).realize()
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
 
@@ -1421,9 +1416,8 @@ def train_llama3():
     loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     loss.backward()
     assert all(p.grad is g for p,g in zip(optim.params, grads))
-    loss_cpu = loss.flatten().float().to("CPU")
-    Tensor.realize(loss_cpu, *grads)
-    return loss_cpu
+    Tensor.realize(loss, *grads)
+    return loss.flatten().float().to("CPU")
 
   @TinyJit
   def optim_step():
@@ -1431,13 +1425,12 @@ def train_llama3():
     scheduler.step()
 
     for g in grads:
-      g.assign(g.zeros_like())
+      g.assign(g.zeros_like()).realize()
 
-    lr_cpu = optim.lr.float().to("CPU")
-    grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads)
+    lr = optim.lr
+    Tensor.realize(lr, *grads)
 
-    return lr_cpu, grad_norm_cpu
+    return lr.float().to("CPU"), grad_norm.float().to("CPU")
 
   @TinyJit
   @Tensor.train(False)

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

@@ -12,7 +12,6 @@ export ALL2ALL=${ALL2ALL:-1}
 export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-1} MP=${MP:-8}

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

@@ -12,7 +12,6 @@ export ALL2ALL=${ALL2ALL:-1}
 export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-1} MP=${MP:-8}

extra/torch_backend/backend.py

@@ -590,7 +590,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.repeat": lambda x,*repeats: Tensor.repeat(x,*repeats).contiguous(), # not a view
   "aten._softmax": lambda self,dim,half_to_float: self.softmax(dim),
   "aten._log_softmax": lambda self,dim,half_to_float: self.log_softmax(dim),
-  "aten.random_": lambda self: Tensor.randint(*self.shape, low=self.dtype.min, high=self.dtype.max, device=self.device, dtype=self.dtype),
+  "aten.random_": lambda self: Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype),
   "aten.random_.from": lambda self, from_, to: Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype),
   "aten.uniform_": lambda self, low=0, high=1: Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype),
   "aten.normal_": lambda self, mean=0, std=1: Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype),

extra/viz/cli.py

@@ -48,16 +48,11 @@ def decode_profile(data:bytes) -> dict:
         name, ref, key, st, dur, fmt = u("<IIIIfI")
         v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
     else:
-      v["linear"] = u("<B")[0]
       v["peak"] = u("<Q")[0]
       for _ in range(event_count):
-        if v["linear"]:
-          ts, value = u("<IQ")
-          v["events"].append({"event":"freq", "ts":ts, "value":value})
-        else:
-          alloc, ts, key = u("<BII")
-          if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
+        alloc, ts, key = u("<BII")
+        if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
   return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
 
 if __name__ == "__main__":

test/amd/test_sqttmap.py

@@ -2,10 +2,9 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG, OSX, getenv, temp
+from tinygrad.helpers import DEBUG, OSX
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
-from tinygrad.viz.serve import sqtt_timeline
 from test.amd.disasm import disasm
 
 import tinygrad
@@ -54,7 +53,7 @@ class TestSQTTMapBase(unittest.TestCase):
   def setUpClass(cls):
     if cls is TestSQTTMapBase: raise unittest.SkipTest("base class")
     cls.examples = {}
-    for pkl_path in ([Path(temp("profile.pkl", append_user=True))] if getenv("LOAD_PROFILE") else sorted((EXAMPLES_DIR/cls.target).glob("*.pkl"))):
+    for pkl_path in sorted((EXAMPLES_DIR/cls.target).glob("*.pkl")):
       with open(pkl_path, "rb") as f:
         data = pickle.load(f)
       sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
@@ -73,29 +72,6 @@ class TestSQTTMapBase(unittest.TestCase):
           passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target, pass_rocprof_err)
           if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")
 
-  def test_sqtt_timeline(self):
-    for name, (events, kern_events, target) in self.examples.items():
-      for event in events:
-        if (p:=kern_events.get(event.kern)) is None: continue
-        with self.subTest(example=name, kern=event.kern):
-          if not (timeline:=sqtt_timeline(event.blob, p.lib, target)): continue
-          frequency = [e.key for e in timeline if type(e).__name__ == "ProfilePointEvent" and e.name == "freq_hz"]
-          mean = sum(frequency) / len(frequency)
-          variance = sum((v - mean) ** 2 for v in frequency) / len(frequency)
-          self.assertGreater(mean, 0)
-          self.assertGreater(variance, 0)
-          if DEBUG >= 2: print(f"{name:20s} SE:{event.se} {mean/1e9:.2f} GHz mean, {variance/1e18:.2f} GHz^2 variance")
-          events = [e for e in timeline if type(e).__name__ == "ProfileRangeEvent"]
-          insts, execs = 0, 0
-          for e in events:
-            if "EXEC" in e.device:
-              if "ALT" not in e.name.display_name: execs += 1
-            elif "WAVE" in e.device:
-              # sopk/immediates don't get ALU/MEM EXEC
-              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE"}: insts += 1
-            else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
-          self.assertEqual(execs, insts)
-
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"
 
 class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"

test/backend/test_dtype.py

@@ -10,7 +10,7 @@ from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
-from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
 
@@ -101,14 +101,14 @@ class TestDType(unittest.TestCase):
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
   def test_uint_overflow(self):
     if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
-    v = self.DTYPE.max
+    v = dtypes.max(self.DTYPE)
     _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
     _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)
 
   def test_dtypes_DTYPES_DICT(self):
     self.assertIn("float", DTYPES_DICT)
     self.assertIn("float32", DTYPES_DICT)
-    self.assertEqual(len(DTYPES_DICT), 28)
+    self.assertEqual(len(DTYPES_DICT), 26)
     self.assertTrue(all(isinstance(value, DType) for value in DTYPES_DICT.values()))
     self.assertTrue(all(issubclass(_to_np_dtype(value), np.generic) for value in DTYPES_DICT.values() if _to_np_dtype(value) is not None))
 
@@ -143,8 +143,6 @@ def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
 class TestFp8s(unittest.TestCase):
   def test_fp8e4m3_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3).dtype == dtypes.fp8e4m3
   def test_fp8e5m2_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2).dtype == dtypes.fp8e5m2
-  def test_fp8e4m3fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3fnuz).dtype == dtypes.fp8e4m3fnuz
-  def test_fp8e5m2fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2fnuz).dtype == dtypes.fp8e5m2fnuz
 
 class TestFp8sConversions(unittest.TestCase):
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
@@ -171,30 +169,6 @@ class TestFp8sConversions(unittest.TestCase):
   def test_fp8e5m2_to_float(self, x):
     np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())
 
-  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3FNUZ_MAX, max_value=FP8E4M3FNUZ_MAX))
-  def test_float_to_fp8e4m3fnuz(self, x):
-    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
-
-  def test_float_to_fp8e4m3fnuz_extreme_values(self):
-    for x in [FP8E4M3FNUZ_MAX, FP8E4M3FNUZ_MAX*1.01, -FP8E4M3FNUZ_MAX, -FP8E4M3FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
-      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
-
-  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2FNUZ_MAX, max_value=FP8E5M2FNUZ_MAX))
-  def test_float_to_fp8e5m2fnuz(self, x):
-    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
-
-  def test_float_to_fp8e5m2fnuz_extreme_values(self):
-    for x in [FP8E5M2FNUZ_MAX, FP8E5M2FNUZ_MAX*1.01, -FP8E5M2FNUZ_MAX, -FP8E5M2FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
-      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
-
-  @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e4m3fnuz_to_float(self, x):
-    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fnuz).float().item())
-
-  @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e5m2fnuz_to_float(self, x):
-    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
-
 @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16(unittest.TestCase):
   def test_bf16_creation_numpy(self):
@@ -516,3 +490,4 @@ class TestOpsBFloat16(unittest.TestCase):
 
 if __name__ == '__main__':
   unittest.main()
+

test/backend/test_dtype_alu.py

@@ -52,8 +52,6 @@ class ht:
 ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0)  # filter subnormal bfloat16
 ht.fp8e4m3 = ht.uint8
 ht.fp8e5m2 = ht.uint8
-ht.fp8e4m3fnuz = ht.uint8
-ht.fp8e5m2fnuz = ht.uint8
 
 def universal_test(a, b, dtype, op):
   if not isinstance(op, tuple): op = (op, op)
@@ -69,8 +67,7 @@ def universal_test(a, b, dtype, op):
     if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
       fe, fm = dtypes.finfo(dtype)
       atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
-    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1),
-                        dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz:(5e-1, 5e-1)}.get(dtype, (1e-10, 1e-7))
+    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7))
     np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
   else: np.testing.assert_equal(tensor_value, numpy_value)
 
@@ -90,8 +87,7 @@ def universal_test_unary(a, dtype, op):
   else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
   if dtype in dtypes.floats:
     atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
-      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1),
-      dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz: (5e-1, 5e-1)}.get(dtype, (1e-6, 1e-5))
+      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5))
     np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
   else: np.testing.assert_equal(tensor_value, numpy_value)
 
@@ -159,26 +155,6 @@ class TestDTypeALU(unittest.TestCase):
   def test_emulated_fp8e5m2(self, a, b, op):
     universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
-  def test_fp8e4m3fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
-  def test_fp8e5m2fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
-  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
-  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
-  def test_emulated_fp8e4m3fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
-  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
-  def test_emulated_fp8e5m2fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
   @given(ht.float32, strat.sampled_from(unary_operations))
   def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
 
@@ -222,30 +198,6 @@ class TestDTypeALU(unittest.TestCase):
     if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
     universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
-  def test_fp8e4m3fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
-  def test_fp8e5m2fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
-  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
-  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
-  def test_emulated_fp8e4m3fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
-  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
-  def test_emulated_fp8e5m2fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
   @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
   def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)
 
@@ -366,7 +318,7 @@ class TestDTypeALU(unittest.TestCase):
 
   @unittest.expectedFailure
   def test_unsafe_cast_float_to_int_failure(self):
-    val = float(dtypes.int32.max - 1)
+    val = float(dtypes.max(dtypes.int32) - 1)
     t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
     t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
     np.testing.assert_equal(t1.item(), t2.item())

test/backend/test_ops.py

@@ -479,9 +479,9 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], 3.])
     helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
     helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
     helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
     helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], True], forward_only=True)
     helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], [True, True, False]], forward_only=True)
 
@@ -496,9 +496,9 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], 3.])
     helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
     helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
     helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
     helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], True], forward_only=True)
     helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], [True, True, False]], forward_only=True)
 

test/backend/test_quantize_onnx.py

@@ -204,13 +204,13 @@ class TestQuantizeOnnx(unittest.TestCase):
     W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
     tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
     out = (X.int().matmul(W.int())//1000)
-    if clip: out = out.clip(tg_dtype.min, tg_dtype.max)
+    if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
     out = out.cast(tg_dtype)
     opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
     sexec(out, opts, replace_src, run_count=1)
     tout = out.numpy()
     mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
-    if clip: mout = mout.clip(tg_dtype.min, tg_dtype.max)
+    if clip: mout = mout.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
     mout = mout.astype(xi)
     print(tout)
     print(mout)

test/backend/test_renderer_failures.py

@@ -62,7 +62,7 @@ class TestRendererFailures(unittest.TestCase):
 class TestCStyleFailures(unittest.TestCase):
   def test_inline_const_alu(self):
     # CPU doesn't use the max function
-    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.min(dtypes.int)+1))
     self.assertEqual(ret[0], 1)
 
   def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):

test/backend/test_schedule.py

@@ -811,26 +811,6 @@ class TestSchedule(unittest.TestCase):
       self.assertEqual(cnt1, 5)
       self.assertEqual(cnt2, 5)
 
-  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
-  def test_image_f16_residual_fusion(self):
-    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
-      def cnt():
-        inp = Tensor.empty((512,), dtype='float')
-        b1, b2 = Tensor.empty((512, 1024), dtype='float'), Tensor.empty((1024, 512), dtype='float')
-        c1, c2 = Tensor.empty((1024,), dtype='float'), Tensor.empty((512,), dtype='float')
-        rb = (((((inp @ b1) + c1).relu() @ b2) + c2).relu() + inp).relu()
-        b16, c16 = Tensor.empty((512, 16), dtype='float'), Tensor.empty((16,), dtype='float')
-        b32, c32 = Tensor.empty((512, 32), dtype='float'), Tensor.empty((32,), dtype='float')
-        sched = Tensor.schedule((rb @ b16 + c16).relu(), (rb @ b32 + c32).relu())
-        for si in sched: si.lower()
-        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
-
-      with Context(IMAGE=1): cnt1 = cnt()
-      with Context(IMAGE=2): cnt2 = cnt()
-
-      self.assertEqual(cnt1, 9)
-      self.assertEqual(cnt2, 9)
-
   @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
   @unittest.expectedFailure
   def test_image_conv_fusion(self):

test/null/test_dtype_spec.py

@@ -75,20 +75,20 @@ class TestHelpers(unittest.TestCase):
   def test_dtype_range(self):
     for dt in core_dtypes:
       if dtypes.is_float(dt):
-        np.testing.assert_equal(dt.min, -math.inf)
-        np.testing.assert_equal(dt.max, math.inf)
+        np.testing.assert_equal(dtypes.min(dt), -math.inf)
+        np.testing.assert_equal(dtypes.max(dt), math.inf)
         np.testing.assert_equal(dt.min, -math.inf)
         np.testing.assert_equal(dt.max, math.inf)
       elif dtypes.is_int(dt):
         info = np.iinfo(_to_np_dtype(dt))
-        np.testing.assert_equal(dt.min, info.min)
-        np.testing.assert_equal(dt.max, info.max)
+        np.testing.assert_equal(dtypes.min(dt), info.min)
+        np.testing.assert_equal(dtypes.max(dt), info.max)
         np.testing.assert_equal(dt.min, info.min)
         np.testing.assert_equal(dt.max, info.max)
       else:
         assert dt == dtypes.bool, dt
-        np.testing.assert_equal(dt.min, False)
-        np.testing.assert_equal(dt.max, True)
+        np.testing.assert_equal(dtypes.min(dt), False)
+        np.testing.assert_equal(dtypes.max(dt), True)
         np.testing.assert_equal(dt.min, False)
         np.testing.assert_equal(dt.max, True)
 

test/null/test_schedule.py

@@ -1189,11 +1189,5 @@ class TestBufferView(unittest.TestCase):
     b = a.shrink(((200, 800),)).shrink(((0, 300),)).reshape((30, 10)).shrink(((20, 25), (0, 10))).contiguous()
     run_schedule(check_schedule(b, 0))
 
-class TestInvalidTensor(unittest.TestCase):
-  def test_full_invalid_is_zero_kernels(self):
-    from tinygrad.dtype import Invalid
-    t = Tensor.full((4,), Invalid, dtype=dtypes.float)
-    check_schedule(t, 0)
-
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/null/test_uop_symbolic.py

@@ -708,18 +708,6 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((31*b+1)//18, 0, 172, "(((b*13)+1)//18+b)")
     self.helper_test_variable((19*b+3)//7, 0, 271, "(((b*5)+3)//7+(b*2))")
 
-  def test_gcd_with_remainder(self):
-    # gcd_with_remainder: factor GCD out of non-constant terms and denominator
-    a = Variable("a", 0, 2)
-    self.helper_test_variable((a*4)//6, 0, 1, "(a*2//3)")
-    self.helper_test_variable((a*4+1)//6, 0, 1, "(a*2//3)")
-    self.helper_test_variable((a*4+2)//6, 0, 1, "((a*2+1)//3)")
-    self.helper_test_variable((a*4+3)//6, 0, 1, "((a*2+1)//3)")
-    self.helper_test_variable((a*4)%6, 0, 4, "(a*2%3*2)")
-    self.helper_test_variable((a*4+1)%6, 1, 5, "(a*2%3*2+1)")
-    self.helper_test_variable((a*4+2)%6, 0, 4, "((a*2+1)%3*2)")
-    self.helper_test_variable((a*4+3)%6, 1, 5, "((a*2+1)%3*2+1)")
-
   def test_div_by_factor_tie_break(self):
     a = Variable("a", 0, 1)
     b = Variable("b", 0, 1)
@@ -759,17 +747,6 @@ class TestSymbolic(unittest.TestCase):
     # f=3, k=2, const=1: (a*3+b+1)%6 = (a%2)*3 + b + 1
     self.helper_test_variable((a*3+b+1)%6, 1, 5, "(b+a%2*3+1)")
 
-  def test_div_nest_by_factor_with_const(self):
-    # nest_by_factor IDIV: (160*a + 5*b + 4*c + K) // 60 should pick div=5 (clean) over div=4 (dirty)
-    a = Variable("a", 0, 2)
-    b = Variable("b", 0, 31)
-    c = Variable("c", 0, 1)
-    self.helper_test_variable((160*a + 5*b + 4*c) // 60, 0, 7, "(a*2+(b+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 1) // 60, 0, 8, "(a*2+(b+c+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 2) // 60, 0, 8, "(a*2+(b+c+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 3) // 60, 0, 8, "(a*2+(b+c+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 59) // 60, 0, 8, "(a*2+(b+c+a*8+11)//12)")
-
   def test_div_mod_recombine_after_nesting(self):
     # when nest_div_by_factor simplifies the div, the mod must also nest so recombine can fire
     gidx0 = Variable("gidx0", 0, 15)
@@ -786,17 +763,6 @@ class TestSymbolic(unittest.TestCase):
     # div nests: y//12 -> a//2, mod nests: y%12 -> (a%2)*6+b, recombine
     self.helper_test_variable((y//12)*12 + y%12, 0, 43, "(b+a*6)")
 
-  def test_div_mod_recombine_in_additive_sum(self):
-    x = Variable("x", 0, 31)
-    y = Variable("y", 0, 5)
-    # recombine should work inside larger additive sums, not just in the two special y+... tree shapes
-    self.helper_test_variable((x//8)*4 + y + (x//2)%4, 0, 20, "(y+x//2)")
-    self.helper_test_variable(y + (x//8)*4 + (x//2)%4, 0, 20, "(y+x//2)")
-
-  def test_div_mod_recompose_low_order_remainder(self):
-    x = Variable("x", 0, 127)
-    self.helper_test_variable((x//2)%4*2 + x%2, 0, 7, "(x%8)")
-
   def test_reshape_index_roundtrip(self):
     # simulate reshape index decompose then recompose — the core pattern this enables
     # (8,8) decomposed for (16,4): combined=r0*8+r1, div and mod by 4
@@ -952,12 +918,6 @@ class TestSymbolic(unittest.TestCase):
     self.assertIn((a.cast(dtypes.long)+b.cast(dtypes.long)).render(), "(long)((a+b))")
     self.assertIn((a.cast(dtypes.long)*b.cast(dtypes.long)).render(), "(long)((a*b))")
 
-  def test_nested_mod_negative_range(self):
-    # (x%(k*c))%c = x%c holds for cmod regardless of signs since sign(x%(k*c)) = sign(x)
-    x = Variable("x", 0, 1575)
-    self.helper_test_variable(((x + (-1064)) % 512) % 4, -3, 3, "((x+-1064)%4)")
-    self.helper_test_variable(((x + (-1064)) % 512) % 128, -127, 127, "((x+-1064)%128)")
-
 class TestSymbolicNumeric(unittest.TestCase):
   def helper_test_numeric(self, f):
     MIN, MAX = 0, 10

test/null/test_uop_vmin_vmax.py

@@ -64,8 +64,8 @@ class TestVminVmaxProperties(unittest.TestCase):
 
     # negative mask: x & -1 could be anything since -1 has all bits set
     uop = x & -1
-    self.assertEqual(uop.vmin, dtypes.int32.min)
-    self.assertEqual(uop.vmax, dtypes.int32.max)
+    self.assertEqual(uop.vmin, dtypes.min(dtypes.int32))
+    self.assertEqual(uop.vmax, dtypes.max(dtypes.int32))
 
   def test_vmin_vmax_multiplication_with_variable(self):
     # vmin and vmax for multiplication with a variable
@@ -136,8 +136,8 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertEqual(x_bool.vmin, False)
     self.assertEqual(x_bool.vmax, True)
     x_uint = x.cast(dtypes.uint)
-    self.assertEqual(x_uint.vmin, dtypes.uint.min)
-    self.assertEqual(x_uint.vmax, dtypes.uint.max)
+    self.assertEqual(x_uint.vmin, dtypes.min(dtypes.uint))
+    self.assertEqual(x_uint.vmax, dtypes.max(dtypes.uint))
 
   def test_vmin_vmax_invalid(self):
     i = UOp.invalid()

test/unit/test_assign.py

@@ -911,36 +911,5 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now, silently dropped
       self.assertEqual(c.tolist(), [[5,5],[5,5]])
 
-class TestPartialAssignToSharedBuffer(unittest.TestCase):
-  def test_five_slices(self):
-    big = Tensor.zeros(50).contiguous().realize()
-    views = [big[i*10:(i+1)*10].reshape(2, 5) for i in range(5)]
-    for v in views: v.assign(v + 1)
-    Tensor.realize(*views)
-    for v in views:
-      np.testing.assert_allclose(v.numpy(), np.ones((2, 5)))
-
-  def test_many_slices(self):
-    n_params = 10
-    big = Tensor.zeros(n_params * 12).contiguous().realize()
-    grads = [big[i*12:(i+1)*12].reshape(3, 4) for i in range(n_params)]
-    for g in grads: g.assign(g + 1)
-    Tensor.realize(*grads)
-    for g in grads:
-      np.testing.assert_allclose(g.numpy(), np.ones((3, 4)))
-
-  def test_mixed_shapes(self):
-    big = Tensor.zeros(100).contiguous().realize()
-    shapes = [(3, 4), (4, 6), (6, 4), (2, 5), (4, 3)]
-    pos, views = 0, []
-    for s in shapes:
-      n = s[0] * s[1]
-      views.append(big[pos:pos+n].reshape(*s))
-      pos += n
-    for v in views: v.assign(v + 1)
-    Tensor.realize(*views)
-    for v, s in zip(views, shapes):
-      np.testing.assert_allclose(v.numpy(), np.ones(s))
-
 if __name__ == "__main__":
   unittest.main()

test/unit/test_dtype_spec.py

@@ -17,15 +17,13 @@ dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and is_dtype_sup
 
 FP8E4M3_MAX = 448.0
 FP8E5M2_MAX = 57344.0
-FP8E4M3FNUZ_MAX = 240.0
-FP8E5M2FNUZ_MAX = 57344.0
 
 def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float=1e-7):
   if DEBUG >= 2: print(tensor.numpy())
   try:
     assert tensor.dtype == target_dtype
     np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2,
-      dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1, dtypes.fp8e4m3fnuz:1e-1, dtypes.fp8e5m2fnuz:5e-1}.get(target_dtype, tol_target_dtype))
+      dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1}.get(target_dtype, tol_target_dtype))
 
   except AssertionError as e:
     raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e

test/unit/test_invalid_tensor.py

@@ -1,126 +0,0 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad.dtype import Invalid, dtypes
-from tinygrad.engine.realize import run_schedule
-
-class TestInvalidTensor(unittest.TestCase):
-  def _invalid_test_helper(self, out, expected):
-    sched = out.schedule()
-    buf = out.uop.buffer
-    buf.allocate()
-    sentinel = memoryview(bytearray(b'\x42' * buf.nbytes))
-    buf.copyin(sentinel)
-    before = buf.as_memoryview().cast(out.dtype.fmt).tolist()
-    run_schedule(sched)
-    ret = buf.as_memoryview().cast(out.dtype.fmt).tolist()
-
-    for i,v in enumerate(expected): self.assertEqual(ret[i], before[i] if v is None else v)
-
-  def test_where_x_invalid(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid)
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_where_invalid_x(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Invalid, Tensor([1.0, 2.0, 3.0, 4.0]))
-    self._invalid_test_helper(out, [None, None, 3.0, 4.0])
-
-  def test_where_invalid_2d(self):
-    mask = Tensor.arange(6).reshape(2, 3) < 3
-    vals = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
-    out = mask.where(vals, Invalid)
-    self._invalid_test_helper(out, [1.0, 2.0, 3.0, None, None, None])
-
-  def test_where_invalid_int(self):
-    mask = Tensor.arange(3) < 2
-    out = mask.where(Tensor([10, 20, 30]), Invalid)
-    self._invalid_test_helper(out, [10, 20, None])
-
-  def test_where_invalid_add(self):
-    mask = Tensor.arange(3) < 2
-    mixed = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
-    out = mixed + Tensor([1.0, 2.0, 3.0])
-    self._invalid_test_helper(out, [11.0, 22.0, None])
-
-  def test_where_invalid_add_left(self):
-    mask = Tensor.arange(3) < 2
-    mixed = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
-    out = Tensor([1.0, 2.0, 3.0]) + mixed
-    self._invalid_test_helper(out, [11.0, 22.0, None])
-
-  def test_where_always_true(self):
-    mask = Tensor.arange(3) < 10
-    out = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
-    self._invalid_test_helper(out, [10.0, 20.0, 30.0])
-
-  def test_where_cast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).cast(dtypes.int)
-    self._invalid_test_helper(out, [1, 2, None, None])
-
-  def test_where_compare(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid) > 1
-    self._invalid_test_helper(out, [False, True, None, None])
-
-  def test_where_unary(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 4.0, 9.0, 16.0]), Invalid).sqrt()
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_where_where(self):
-    mask1 = Tensor.arange(4) < 2
-    mask2 = Tensor.arange(4) > 0
-    out = mask2.where(mask1.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid), Invalid)
-    self._invalid_test_helper(out, [None, 2.0, None, None])
-
-  def test_where_reduce_always_true(self):
-    mask = Tensor.arange(4) < 9
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).sum()
-    self._invalid_test_helper(out, [10.0])
-
-  def test_invalid_unary(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float).sqrt())
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_binary(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float) + 2)
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_binary_left(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), 2 + Tensor.full((4,), Invalid, dtype=dtypes.float))
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_reshape(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).reshape(2,2)
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_cast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int).cast(dtypes.float))
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_bitcast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int).bitcast(dtypes.float))
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_where_bitcast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int)).bitcast(dtypes.int)
-    self._invalid_test_helper(out, [0x3f800000, 0x40000000, None, None])
-
-  # tensor indexing uses reduce, so the entire result becomes invalid
-  @unittest.expectedFailure
-  def test_tensor_index(self):
-    idx = (Tensor.arange(4) < 2).where(Tensor([0, 1, 2, 3]), Invalid)
-    out = Tensor([1.0, 2.0, 3.0, 4.0])[idx]
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-if __name__ == '__main__':
-  unittest.main()

tinygrad/device.py

@@ -353,12 +353,11 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
     if device == "NV": return not CI and not NV_PTX and not NV_NAK
     if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not CPU_LVP
     return device in {"AMD", "CL", "PYTHON", "NULL"}
-  if dtype in dtypes.fp8_ocp:
+  if dtype in dtypes.fp8s:
     if device == "CUDA": return not CI and not CUDA_PTX
     if device == "NV": return not CI and not NV_PTX and not NV_NAK
-    if device == "AMD": return not CI and getattr(Device["AMD"], "target") == (9,5,0)
+    if device == "AMD": return not CI and getattr(Device["AMD"], "target") in {(9,4,2), (9,5,0)}
     return device in {"PYTHON", "NULL"}
-  if dtype in dtypes.fp8_fnuz: return device in {"PYTHON", "NULL"}
   if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
                                           dtypes.ushort, dtypes.float, dtypes.int32, dtypes.uint32, dtypes.half]
   # for CI GPU and OSX, cl_khr_fp16 isn't supported

tinygrad/dtype.py

@@ -30,7 +30,6 @@ class InvalidType:
   def __hash__(self): return id(self)
   def __repr__(self): return "Invalid"
   def __reduce__(self): return (InvalidType, ())  # unpickle returns the singleton
-  def __format__(self, spec): return "Invalid"
 
 Invalid = InvalidType()
 
@@ -79,14 +78,10 @@ class DType(metaclass=DTypeMetaClass):
     return PtrDType(self.priority, self.bitsize, self.name, self.fmt, self.count, None, self, addrspace, 1, size)
   def scalar(self) -> DType: return self._scalar if self._scalar is not None else self
   def nbytes(self) -> int: raise RuntimeError("only ptr types have nbytes")
-  @functools.cached_property
-  def min(self):
-    if dtypes.is_int(self): return 0 if dtypes.is_unsigned(self) else -2**(self.scalar().bitsize-1)
-    return -float("inf") if dtypes.is_float(self) else False
-  @functools.cached_property
-  def max(self):
-    if dtypes.is_int(self): return 2**(self.scalar().bitsize)-1+self.min
-    return float("inf") if dtypes.is_float(self) else True
+  @property
+  def min(self): return dtypes.min(self)
+  @property
+  def max(self): return dtypes.max(self)
 
 @dataclass(frozen=True, eq=False)
 class PtrDType(DType):
@@ -176,11 +171,21 @@ class dtypes:
     # int is the default. wrap floats in ConstFloat to distinguish -0.0 from 0.0 in cache
     return ConstFloat(float(val)) if dtypes.is_float(dtype) else bool(val) if dtypes.is_bool(dtype) else int(val)
   @staticmethod
+  @functools.cache
+  def min(dtype:DType):
+    if dtypes.is_int(dtype): return 0 if dtypes.is_unsigned(dtype) else -2**(dtype.scalar().bitsize-1)
+    return -float("inf") if dtypes.is_float(dtype) else False
+  @staticmethod
+  @functools.cache
+  def max(dtype:DType):
+    if dtypes.is_int(dtype): return 2**(dtype.scalar().bitsize)-1+dtypes.min(dtype)
+    return float("inf") if dtypes.is_float(dtype) else True
+  @staticmethod
   def finfo(dtype:DType) -> tuple[int, int]:
     """(exponent, mantissa)"""
     if not dtypes.is_float(dtype): raise ValueError(f"{dtype} is not a floating point type")
     return {dtypes.float16: (5, 10), dtypes.bfloat16: (8, 7), dtypes.float32: (8, 23), dtypes.float64: (11, 52),
-            dtypes.fp8e4m3: (4, 3), dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3fnuz: (4, 3), dtypes.fp8e5m2fnuz: (5, 2)}[dtype]
+            dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3: (4, 3)}[dtype]
   void: Final[DType] = DType.new(-1, 0, "void", None)
   index: Final[DType] = DType.new(-1, 800, "index", None)
   bool: Final[DType] = DType.new(0, 1, "bool", '?')
@@ -196,8 +201,6 @@ class dtypes:
   _uint256: Final[DType] = DType.new(8, 256, "uint256", None)
   fp8e4m3: Final[DType] = DType.new(9, 8, "float8_e4m3", None)
   fp8e5m2: Final[DType] = DType.new(10, 8, "float8_e5m2", None)
-  fp8e4m3fnuz: Final[DType] = DType.new(9, 8, "float8_e4m3fnuz", None)
-  fp8e5m2fnuz: Final[DType] = DType.new(10, 8, "float8_e5m2fnuz", None)
   float16: Final[DType] = DType.new(11, 16, "half", 'e')
   # bfloat16 has higher priority than float16, so least_upper_dtype(dtypes.int64, dtypes.uint64) = dtypes.float16
   bfloat16: Final[DType] = DType.new(12, 16, "__bf16", None)
@@ -218,9 +221,7 @@ class dtypes:
   default_float: ClassVar[DType] = float32
   default_int: ClassVar[DType] = int32
 
-  fp8_ocp = (fp8e4m3, fp8e5m2)
-  fp8_fnuz = (fp8e4m3fnuz, fp8e5m2fnuz)
-  fp8s = fp8_ocp + fp8_fnuz
+  fp8s = (fp8e4m3, fp8e5m2)
   floats = fp8s + (float16, bfloat16, float32, float64)
   int8s = (uint8, int8)
   int16s = (uint16, int16)
@@ -242,9 +243,8 @@ def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType)
 # we don't support weak type and complex type
 promo_lattice = { dtypes.bool: [dtypes.int8, dtypes.uint8], dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
   dtypes.int64: [dtypes.uint64], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32],
-  dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.fp8e4m3, dtypes.fp8e5m2, dtypes.fp8e4m3fnuz, dtypes.fp8e5m2fnuz],
-  dtypes.fp8e4m3: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e5m2: [dtypes.float16, dtypes.bfloat16],
-  dtypes.fp8e4m3fnuz: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e5m2fnuz: [dtypes.float16, dtypes.bfloat16],
+  dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.fp8e4m3, dtypes.fp8e5m2],
+  dtypes.fp8e5m2: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e4m3: [dtypes.float16, dtypes.bfloat16],
   dtypes.float16: [dtypes.float32], dtypes.bfloat16: [dtypes.float32], dtypes.float32: [dtypes.float64], }
 
 @functools.cache
@@ -299,14 +299,10 @@ def float_to_bf16(x):
 _fp8_cfg = {
   dtypes.fp8e4m3: (7, 4, 0x7, 0x3F50000000000000, 0x407D000000000000, 0x7E, 0x3F90000000000000),
   dtypes.fp8e5m2: (15, 3, 0x3, 0x3EE0000000000000, 0x40EE000000000000-1, 0x7B, 0x3F10000000000000),
-  dtypes.fp8e4m3fnuz: (8, 4, 0x7, 0x3F40000000000000, 0x406F000000000000-1, 0x7F, 0x3F80000000000000),
-  dtypes.fp8e5m2fnuz: (16, 3, 0x3, 0x3ED0000000000000, 0x40EE000000000000-1, 0x7F, 0x3F00000000000000),
 }
 
 def float_to_fp8(x: float, dtype: DType) -> int:
   assert dtype in dtypes.fp8s, "Only for fp8s"
-  if dtype in dtypes.fp8_fnuz and not math.isfinite(x): return 0x80
-  if dtype in dtypes.fp8_fnuz and x == 0.0: return 0x00
   # e4m3 don't support inf, return 0x7f(+NaN) and 0xff(-NaN) to match jax
   # NaN is unordered, can't compare with zero, use math.copysign to get sign
   if dtype == dtypes.fp8e4m3 and not math.isfinite(x): return 0x7f if math.copysign(1, x) > 0 else 0xff
@@ -326,17 +322,16 @@ def float_to_fp8(x: float, dtype: DType) -> int:
     res, half = mantissa >> shift, half_ulp << shift
     round_bits = (xbits | (1 << 52)) & ((half << 1) - 1)
     if round_bits > half or (round_bits == half and res & 1): res += 1
-  return 0 if dtype in dtypes.fp8_fnuz and res == 0 else int(res | sign)  # fnuz has no negative zero
+  return int(res | sign)
 
 def fp8_to_float(x: int, dtype: DType) -> float:
   assert dtype in dtypes.fp8s, "Only for fp8s"
-  if dtype in dtypes.fp8_fnuz and x == 0x80: return math.nan
   if (x & 0x7F) == 0: return -0.0 if x & 0x80 else 0.0
   bias, sig_bits, *_ = _fp8_cfg[dtype]
   mant_bits, exp_bits = sig_bits - 1, 8 - sig_bits
   exp_max, mant_max = (1 << exp_bits) - 1, (1 << mant_bits) - 1
   sign, exp, mantissa = (x >> 7) & 1, (x >> mant_bits) & exp_max, x & mant_max
-  if dtype not in dtypes.fp8_fnuz and exp == exp_max:
+  if exp == exp_max:
     if dtype == dtypes.fp8e5m2: return math.copysign(math.nan if mantissa else math.inf, -1 if sign else 1)
     if mantissa == mant_max: return math.nan
   val = (mantissa / (mant_max + 1)) * 2 ** (1 - bias) if exp == 0 else (1 + mantissa / (mant_max + 1)) * 2 ** (exp - bias)

tinygrad/nn/onnx.py

@@ -497,7 +497,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     o_ = [((i - 1) // s + 1) for i,s in zip(i_, s_)]
     return _onnx_pads_to_tiny_pads(_auto_pad([(o-1)*s+k-i for o,i,k,s in zip(o_, i_, k_, s_)], auto_pad))
 
-  def _clamp_cast(x:Tensor, dtype:DType): return x.clamp(dtype.min, dtype.max).cast(dtype)
+  def _clamp_cast(x:Tensor, dtype:DType): return x.clamp(dtypes.min(dtype), dtypes.max(dtype)).cast(dtype)
 
   def _prepare_quantize(x:Tensor, scale:Tensor, zero_point:Tensor|int, axis=1, block_size=0):
     if axis < 0: axis += x.ndim
@@ -1209,7 +1209,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   def DynamicQuantizeLinear(x: Tensor):
     # only support uint8
-    qmin, qmax = dtypes.uint8.min, dtypes.uint8.max
+    qmin, qmax = dtypes.min(dtypes.uint8), dtypes.max(dtypes.uint8)
     scale = (x.max().maximum(0) + ((-x).max()).maximum(0)) / (qmax - qmin)
     zero_point = _clamp_cast((qmin - x.min() / scale).round(), dtypes.uint8)
     y = _clamp_cast((x / scale).round() + zero_point, dtypes.uint8)

tinygrad/renderer/amd/sqtt.py

@@ -106,7 +106,7 @@ class InstOpRDNA4(Enum):
   SMEM = 0x1
   JUMP = 0x3
   JUMP_NO = 0x4
-  CALL = 0x5
+  JUMP_UNCOND = 0x5
   MESSAGE = 0x9
   VALU_TRANS = 0xb
   VALU_B2 = 0xd
@@ -187,16 +187,19 @@ class TS_DELTA_SHORT(PacketType):
 class TS_DELTA_OR_MARK(PacketType):
   encoding = bits[6:0] == 0b0000001
   delta = bits[47:12]
-  pl = bits[8:8]
-  rt = bits[9:9]
+  bit8 = bits[8:8]
+  bit9 = bits[9:9]
   @property
-  def is_marker(self) -> bool: return bool(self.rt and not self.pl)
+  def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
 
-class TS_DELTA_OR_MARK_RDNA4(TS_DELTA_OR_MARK):
+class TS_DELTA_OR_MARK_RDNA4(PacketType):  # Layout 4: 48->64 bits
+  encoding = bits[6:0] == 0b0000001
   delta = bits[63:12]
-  rt = bits[7:7]
-  pl = bits[8:8]
-  tl = bits[9:9]
+  bit7 = bits[7:7]
+  bit8 = bits[8:8]
+  bit9 = bits[9:9]
+  @property
+  def is_marker(self) -> bool: return bool((self.bit9 and not self.bit8) or self.bit7)
 
 class TS_DELTA_S5_W2(PacketType):
   encoding = bits[4:0] == 0b11100

tinygrad/runtime/ops_nv.py

@@ -718,8 +718,6 @@ class NVDevice(HCQCompiled[NVSignal]):
       if coloc_size > self.vid_coloc_buf.size: self.vid_coloc_buf, _ = self._realloc(self.vid_coloc_buf, coloc_size, force=True)
       if filter_size > self.vid_filter_buf.size: self.vid_filter_buf, _ = self._realloc(self.vid_filter_buf, filter_size, force=True)
 
-  def hw_copy_queues(self): return super().hw_copy_queues() + ([("NVDEC:0", NVVideoQueue)] if hasattr(self, 'vid_gpfifo') else [])
-
   def invalidate_caches(self):
     if self.is_nvd(): self.iface.rm_control(self.subdevice, nv_gpu.NV2080_CTRL_CMD_INTERNAL_BUS_FLUSH_WITH_SYSMEMBAR, None)
     else:

tinygrad/schedule/indexing.py

@@ -26,12 +26,14 @@ def realize_assign_src(ctx:dict[UOp, None], buf:UOp, x:UOp):
   if buf.base in x.backward_slice_with_self: ctx[x] = None
 
 pm_generate_realize_map = PatternMatcher([
+  # always realize SINK src
+  (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
   # always realize
-  (UPat({Ops.COPY, Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.ASSIGN}, name="tr"), realize),
   # realize srcs of these
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # sometimes realize src of assign
-  (UPat(Ops.STORE, src=(UPat.var("buf"), UPat.var("x"))), realize_assign_src),
+  (UPat(Ops.ASSIGN, src=(UPat.var("buf"), UPat.var("x"))), realize_assign_src),
 ])
 
 @dataclass(frozen=True)

tinygrad/schedule/rangeify.py

@@ -61,7 +61,7 @@ def normalize_assign_target_chain(assign:UOp, target:UOp, src:UOp):
   root_target = target
   while root_target.op is Ops.ASSIGN: root_target = root_target.src[0]
   # when RHS depends on the previous assign result, break with contiguous
-  #if target in src.toposort(): src = src.contiguous()
+  if target in src.toposort(): src = src.contiguous()
   return assign.replace(src=(root_target, src))
 
 def split_reduceop(reduce:UOp, x:UOp):
@@ -160,7 +160,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.ASSIGN, src=(UPat(Ops.ASSIGN, name="target"), UPat(name="src")), allow_any_len=True, name="assign"), normalize_assign_target_chain),
 
   # make source contiguous if it has hazardous movement ops on the dest buffer
-  (UPat(Ops.STORE, src=(UPat.var("target"), UPat.var("src")), name="assign"), fix_assign_hazard),
+  (UPat(Ops.ASSIGN, src=(UPat.var("target"), UPat.var("src")), name="assign"), fix_assign_hazard),
 
   # ** size 0 **
 
@@ -210,54 +210,56 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   # if it's user contiguous, we never remove it
   if src.op in ALWAYS_RUN_OPS or not buf.arg.removable: return None
 
-  # *** here is where we compute the cost ***
-  # if we return None, the bufferize is kept
+  # we don't want to bufferize threefry, also causes problems because not all platforms support long
+  if src.op is not Ops.THREEFRY:
+    # *** here is where we compute the cost ***
+    # if we return None, the bufferize is kept
 
-  accessed_buffers: list[UOp] = []
-  indexes: list[UOp] = []
-  reduces: list[UOp] = []
-  def red_gate(x:UOp):
-    if (x.op is Ops.BUFFERIZE and x.arg.addrspace == AddrSpace.GLOBAL) or x.op is Ops.MSTACK:
-      accessed_buffers.append(x)
-      return False
-    if x.op is Ops.PARAM:
-      accessed_buffers.append(x)
-    if x.op is Ops.INDEX:
-      indexes.append(x)
-    if x.op is Ops.REDUCE: reduces.append(x)
-    return True
-  src.toposort(gate=red_gate)
-  del red_gate
-  accessed_buffers = dedup(accessed_buffers)
+    accessed_buffers: list[UOp] = []
+    indexes: list[UOp] = []
+    reduces: list[UOp] = []
+    def red_gate(x:UOp):
+      if (x.op is Ops.BUFFERIZE and x.arg.addrspace == AddrSpace.GLOBAL) or x.op is Ops.MSTACK:
+        accessed_buffers.append(x)
+        return False
+      if x.op is Ops.PARAM:
+        accessed_buffers.append(x)
+      if x.op is Ops.INDEX:
+        indexes.append(x)
+      if x.op is Ops.REDUCE: reduces.append(x)
+      return True
+    src.toposort(gate=red_gate)
+    del red_gate
+    accessed_buffers = dedup(accessed_buffers)
 
-  # if this is generated from multiple buffers, don't remove this buffer
-  if len(accessed_buffers) > 3 and not (PCONTIG > 2): return None
+    # if this is generated from multiple buffers, don't remove this buffer
+    if len(accessed_buffers) > 3 and not (PCONTIG > 2): return None
 
-  # if any reduces access a buffer, don't remove this buffer
-  buffer_in_reduce = False
-  def buf_gate(x:UOp):
-    nonlocal buffer_in_reduce
-    if x.op in {Ops.PARAM, Ops.BUFFERIZE}: buffer_in_reduce = True
-    return not buffer_in_reduce
-  UOp.sink(*[x.src[0] for x in reduces]).toposort(gate=buf_gate)
-  del buf_gate
-  if buffer_in_reduce:
-    if PCONTIG > 2:
-      out_in_ratio = (prod(buf.shape)+1) / (sum([x.size for x in accessed_buffers])+1)
-      if out_in_ratio < 10: return None
-      # here we have to check the indexes, we might do a partial contig here
-      local_indexes = [x for x in indexes if x.src[0].op is Ops.BUFFERIZE and x.src[0].arg.addrspace == AddrSpace.LOCAL]
-      exclude_ranges = UOp.group(*[UOp.group(*x.src[1:]) for x in local_indexes]).ranges
-      subs = [(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST]
-      # if it's bufferized or a reduce, it's pcontig
-      is_pcontig, is_subs = partition(subs, lambda x: x[0] in exclude_ranges or any([r.arg[-1] == AxisType.REDUCE for r in x[1].ranges]))
-      if not len(is_subs):
+    # if any reduces access a buffer, don't remove this buffer
+    buffer_in_reduce = False
+    def buf_gate(x:UOp):
+      nonlocal buffer_in_reduce
+      if x.op in {Ops.PARAM, Ops.BUFFERIZE}: buffer_in_reduce = True
+      return not buffer_in_reduce
+    UOp.sink(*[x.src[0] for x in reduces]).toposort(gate=buf_gate)
+    del buf_gate
+    if buffer_in_reduce:
+      if PCONTIG > 2:
+        out_in_ratio = (prod(buf.shape)+1) / (sum([x.size for x in accessed_buffers])+1)
+        if out_in_ratio < 10: return None
+        # here we have to check the indexes, we might do a partial contig here
+        local_indexes = [x for x in indexes if x.src[0].op is Ops.BUFFERIZE and x.src[0].arg.addrspace == AddrSpace.LOCAL]
+        exclude_ranges = UOp.group(*[UOp.group(*x.src[1:]) for x in local_indexes]).ranges
+        subs = [(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST]
+        # if it's bufferized or a reduce, it's pcontig
+        is_pcontig, is_subs = partition(subs, lambda x: x[0] in exclude_ranges or any([r.arg[-1] == AxisType.REDUCE for r in x[1].ranges]))
+        if not len(is_subs):
+          return None
+        if len(is_pcontig):
+          ret = src.substitute(dict(is_subs), extra_pm=pm_gate_substitute)
+          return ret.bufferize(*[x[0] for x in is_pcontig], arg=BufferizeOpts(None, AddrSpace.LOCAL)).index(*[x[1] for x in is_pcontig])
+      else:
         return None
-      if len(is_pcontig):
-        ret = src.substitute(dict(is_subs), extra_pm=pm_gate_substitute)
-        return ret.bufferize(*[x[0] for x in is_pcontig], arg=BufferizeOpts(None, AddrSpace.LOCAL)).index(*[x[1] for x in is_pcontig])
-    else:
-      return None
 
   # if it makes it here, the bufferize is removed
   # this is the ranges replaced
@@ -357,15 +359,10 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
     assign_target, assign_src = assign.src[0], assign.src[1]
     assert assign_target.op is Ops.INDEX, f"{assign_target.op} is not index"
     while assign_src.op is Ops.NOOP: assign_src = assign_src.src[0]
-
-    store_target = assign_target
-    if assign.arg and assign_target.src[0].op is Ops.BUFFERIZE and assign_target.src[0].src[0].op is Ops.INDEX:
-      # BUFFERIZE(INDEX(...)); store through the underlying global index instead.
-      store_target = assign_target.src[0].src[0]
-
-    end_rngs = sorted(dedup(tuple(store_target.ranges) + tuple(rngs)), key=lambda x: x.arg)
-    ret = store_target.buf_uop.base
-    if assign_src is not store_target: ret = ret.after(store_target.replace(dtype=sdtype).store(assign_src).end(*end_rngs))
+    # skip self-assign from same-device copy, otherwise create the store
+    # in assign, this is the buffer size, not the bufferize size
+    if assign_src is assign_target: ret = assign_target.src[0]
+    else: ret = assign_target.src[0].after(assign_target.replace(dtype=sdtype).store(assign_src).end(*rngs))
     for op, marg in reversed(assign.arg or ()): ret = ret._mop(op, marg)
     return ret
 

tinygrad/tensor.py

@@ -5,7 +5,7 @@ from contextlib import ContextDecorator
 from typing import Any, Callable, ClassVar, Sequence, cast, get_args, Literal, SupportsIndex, ParamSpec, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
-from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst, Invalid, InvalidType
+from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten
 from tinygrad.helpers import IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
 from tinygrad.helpers import suppress_finalizing, disable_gc
@@ -113,7 +113,7 @@ class Tensor(OpMixin):
   __slots__ = "uop", "requires_grad", "grad"
   training: ClassVar[bool] = False
 
-  def __init__(self, data:ConstType|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.Path|None,
+  def __init__(self, data:PyConst|bytes|list|tuple|UOp|'numpy.ndarray'|pathlib.Path|None,
                device:str|tuple|list|None=None, dtype:DTypeLike|None=None, requires_grad:bool|None=None, _force_unique:bool=False):
     if device is None and isinstance(data, pathlib.Path): device = f"DISK:{data.resolve()}"  # keep it on the disk if device is None
     _dtype:DType|None = to_dtype(dtype) if dtype is not None else None
@@ -141,9 +141,6 @@ class Tensor(OpMixin):
       data = Tensor(0, device=_device, dtype=_dtype or dtypes.default_float, requires_grad=requires_grad).uop
     elif isinstance(data, get_args(PyConst)):
       data = (UOp.unique_const if _force_unique or requires_grad else UOp.const)(_dtype or dtypes.from_py(data), data, _device)
-    elif isinstance(data, InvalidType):
-      assert _dtype is not None
-      data = UOp.const(_dtype, data, _device)
     elif isinstance(data, bytes): data = _frompy(data, dtypes.uint8 if _dtype is None else _dtype)
     elif isinstance(data, (list, tuple)):
       if _dtype is None:
@@ -1064,7 +1061,7 @@ class Tensor(OpMixin):
     for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient, materialize_grads=True)):
       assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
       if t.grad is None: t.grad = g
-      else: t.grad.assign(t.grad + g.to(t.grad.device))
+      else: t.grad.assign(t.grad + g)
     return self
 
   # ***** movement low level ops *****
@@ -2109,7 +2106,7 @@ class Tensor(OpMixin):
     x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None))
     x_cummax, _ = x.cummax(-1)
     mask = Tensor.ones(last_dim_size, last_dim_size, requires_grad=False, device=self.device).tril()
-    ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax
+    ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), dtypes.min(self.dtype)).exp().sum(-1).log() + x_cummax
     return ret.transpose(-1, axis)
 
   def argmax(self, axis=None, keepdim=False) -> Tensor:
@@ -2306,12 +2303,12 @@ class Tensor(OpMixin):
     axis = tuple(range(-len(k_ := make_tuple(kernel_size, 2)), 0))
     pads = self._resolve_pool_pads(padding, len(k_))
     if ceil_mode: pads = self._apply_ceil_mode(pads, k_, stride if stride is not None else k_, dilation)
-    pooled = self.pad(pads, value=self.dtype.min)._pool(k_, stride if stride is not None else k_, dilation)
+    pooled = self.pad(pads, value=dtypes.min(self.dtype))._pool(k_, stride if stride is not None else k_, dilation)
     if not return_indices: return pooled.max(axis)
     spatial_sz = int(math.prod(spatial_shape := self.shape[-len(k_):]))
     idx = Tensor.arange(spatial_sz,0,-1, requires_grad=False, device=self.device).reshape(spatial_shape)
     m = pooled == pooled.max(axis, keepdim=True)
-    idx = m * idx.pad(pads, value=idx.dtype.min)._pool(k_, stride if stride is not None else k_, dilation)
+    idx = m * idx.pad(pads, value=dtypes.min(idx.dtype))._pool(k_, stride if stride is not None else k_, dilation)
     return pooled.max(axis), spatial_sz - idx.max(axis)
 
   def max_unpool2d(self, indices:Tensor, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0, output_size=None):
@@ -2752,8 +2749,8 @@ class Tensor(OpMixin):
     def _inv_mask(a:Tensor|PyConst, b:Tensor|PyConst) -> Tensor: return mask.any(-1).logical_not().where(a, b)
     if reduce == "sum": return mask.where(src, 0).sum(-1).add(self if include_self else _inv_mask(self, 0))
     if reduce == "prod": return mask.where(src, 1).prod(-1).mul(self if include_self else _inv_mask(self, 1))
-    if reduce == "amax": return mask.where(src, m := src.dtype.min).max(-1).maximum(self if include_self else _inv_mask(self, m))
-    if reduce == "amin": return mask.where(src, m := src.dtype.max).min(-1).minimum(self if include_self else _inv_mask(self, m))
+    if reduce == "amax": return mask.where(src, m := dtypes.min(src.dtype)).max(-1).maximum(self if include_self else _inv_mask(self, m))
+    if reduce == "amin": return mask.where(src, m := dtypes.max(src.dtype)).min(-1).minimum(self if include_self else _inv_mask(self, m))
     if reduce == "mean":
       count = mask.where(1, 0).sum(-1).add(1 if include_self else _inv_mask(1, 0))
       return mask.where(src, 0).sum(-1).add(self if include_self else _inv_mask(self, 0)).div(count)
@@ -2782,7 +2779,7 @@ class Tensor(OpMixin):
     # pad to power of 2
     n_stages = (orig_len-1).bit_length()
     pads = tuple((0, 2**n_stages - orig_len) if i == dim else None for i in range(x.ndim))
-    x = x.pad(pads, value=x.dtype.min if descending else x.dtype.max).unflatten(dim, (2,)*n_stages)
+    x = x.pad(pads, value=dtypes.min(x.dtype) if descending else dtypes.max(x.dtype)).unflatten(dim, (2,)*n_stages)
     # https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort1.svg
     for stage in range(1, n_stages+1):
       if stage != n_stages:
@@ -2947,8 +2944,7 @@ class Tensor(OpMixin):
     if not isinstance(y, Tensor):
       # make y a Tensor
       assert isinstance(y, (*get_args(ConstType), UOp)), f"{type(y)=}, {y=}"
-      if y is Invalid or isinstance(x.dtype, ImageDType) or dtypes.is_float(x.dtype) or (dtypes.is_int(x.dtype) and isinstance(y, int)):
-        y_dtype = x.dtype
+      if isinstance(x.dtype, ImageDType) or dtypes.is_float(x.dtype) or (dtypes.is_int(x.dtype) and isinstance(y, int)): y_dtype = x.dtype
       elif not isinstance(y, UOp): y_dtype = dtypes.from_py(y)
       if isinstance(y, UOp): y = Tensor.from_uop(y, device=x.device)
       else: y = Tensor(dtypes.as_const(y, y_dtype), x.device, y_dtype, requires_grad=False)
@@ -3655,15 +3651,11 @@ class Tensor(OpMixin):
 
     # contiguous creates the image, and early realize static weights (TODO: test for the static weight)
     if IMAGE == 1:
-      # pad with Invalid
-      def _invalid_pad_to(t, shape):
-        if all(p is None or p == s for p,s in zip(shape, t.shape)): return t
-        return Tensor(True, device=t.device).expand(t.shape).pad_to(shape).where(t.pad_to(shape), Invalid)
       # hacks for pitch alignment
       assert isinstance(ix, int) and isinstance(H, int)
       ALIGN = 64 // dtsz
-      x = _invalid_pad_to(x, (None, None, round_up(ix, ALIGN // math.gcd(groups * cin, ALIGN)), None))
-      w = _invalid_pad_to(w, (None, round_up(H, ALIGN // math.gcd(W * cin * 4, ALIGN))) + (None,) * (w.ndim - 2))
+      x = x.pad_to(None, None, round_up(ix, ALIGN // math.gcd(groups * cin, ALIGN)), None)
+      w = w.pad_to((None, round_up(H, ALIGN // math.gcd(W * cin * 4, ALIGN))) + (None,) * (w.ndim - 2))
 
       if FLOAT16: x, w = x.cast(dtypes.half).contiguous().cast(dtypes.float), w.cast(dtypes.half).contiguous().cast(dtypes.float)
       else: x, w = x.contiguous(), w.contiguous()

tinygrad/uop/__init__.py

@@ -82,8 +82,7 @@ class Ops(FastEnum):
   # ** 6 -- ops that don't exist in programs **
 
   # tensor graph ops
-  UNIQUE = auto(); DEVICE = auto() #; ASSIGN = auto()
-  ASSIGN = AFTER  # ASSIGN is AFTER now (remove it)
+  UNIQUE = auto(); DEVICE = auto(); ASSIGN = auto()
 
   # local unique
   LUNIQUE = auto()

tinygrad/uop/decompositions.py

@@ -14,7 +14,7 @@ def _lazy_map_numbers(x:UOp, inf:UOp, _inf:UOp, nan:UOp, ratio:UOp):
 
 # *** helper functions for bit manipulation ***
 def mantissa_bits(d:DType) -> int: return dtypes.finfo(d.scalar())[1]
-def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d.scalar())[0] - 1)) - (0 if d.scalar() in dtypes.fp8_fnuz else 1)
+def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d.scalar())[0] - 1)) - 1
 def exponent_mask(d:DType) -> int: return (1 << dtypes.finfo(d.scalar())[0]) - 1
 
 # **** utils ****
@@ -287,7 +287,7 @@ def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None:
   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)
   m,s = magicgu(max(vmax, abs(vmin)), d)
-  if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
+  if m*vmin >= dtypes.min(x.dtype) and m*vmax <= dtypes.max(x.dtype):
     return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
   # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
   if (largest_factor_of_two_in_d := (d & -d)) > 1:
@@ -295,7 +295,7 @@ def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None:
   if dont_cast: return None
   # promo_lattice needs to return an unsigned type if the type is unsigned
   if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and is_dtype_supported(next_dtype, device):
-    if m*vmin >= next_dtype.min and m*vmax <= next_dtype.max:
+    if m*vmin >= dtypes.min(next_dtype) and m*vmax <= dtypes.max(next_dtype):
       return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0)
   return None
 
@@ -388,10 +388,6 @@ def f2f(v, fr:DType, to:DType):
     sign, nosign = shl((v & shl(1, fs-1)).cast(f2f_dt[to]), ts - fs), (v & (shl(1, fs-1) - 1)).cast(f2f_dt[to])
     exp, norm = shr(nosign, fm), shl(nosign, tm - fm) + shl(tb - fb, tm)
     nan = shl(nosign, tm - fm) | shl((shl(1, te) - 1), tm)
-    if fr in dtypes.fp8_fnuz:
-      fnuz_nan = sign.ne(0) & nosign.eq(0)
-      qnan = shl(shl(1, te) - 1, tm) | shl(1, tm - 1)
-      return fnuz_nan.where(qnan, sign | exp.eq(0).where(0, norm)).bitcast(to)
     # fp8e4m3 has only one nan
     is_nan = (nosign.eq(shl(1, fm + fe) - 1) if fr == dtypes.fp8e4m3 else exp.eq(shl(1, fe) - 1))
     return (sign | exp.eq(0).where(0, is_nan.where(nan, norm))).bitcast(to)
@@ -403,14 +399,12 @@ def f2f(v, fr:DType, to:DType):
     nan_mantissa = (shl(1, tm) - 1) if to == dtypes.fp8e4m3 else (shr(nosign, fm - tm) & (shl(1, tm) - 1))
     nan = (sign | nan_mantissa | shl(shl(1, te) - 1, tm)).cast(f2f_dt[to])
     is_nan = (shr(v, fm) & (shl(1, fe) - 1)).eq(shl(1, fe) - 1)
-    if to in dtypes.fp8_fnuz: return is_nan.where(shl(1, ts - 1), underflow.where(0, sign.cast(f2f_dt[to]) | norm))
     return is_nan.where(nan, sign.cast(f2f_dt[to]) | underflow.where(0, norm))
   else: raise NotImplementedError(f"unsupported decomp {fr} -> {to}")
 
 def f2f_clamp(val:UOp, dt:DType) -> UOp:
   e, m = dtypes.finfo(dt)
-  if dt in dtypes.fp8_fnuz: max_exp, max_man = (1 << e) - 1, (1 << m) - 1
-  else: max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1)
+  max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1)
   mx = val.const_like(2.0**(max_exp - exponent_bias(dt)) * (1.0 + max_man / (1 << m)))
   sat = mx if dt in dtypes.fp8s else val.const_like(float('inf'))
   # FIXME: CMPLT of nan is undefined

tinygrad/uop/divandmod.py

@@ -1,4 +1,4 @@
-import functools, itertools, math
+import functools, itertools
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import cdiv, cmod, CORRECT_DIVMOD_FOLDING, unwrap
@@ -12,8 +12,8 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
   x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
   assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
   if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
-  if y_min*y_max > 0 and (qv:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
-    return x - qv*y if d.op is Ops.MOD else d.const_like(qv)
+  if y_min*y_max > 0 and (q:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
+    return x - q*y if d.op is Ops.MOD else d.const_like(q)
 
   # split uops for the rest of the processing
   x_peeled, const = x.pop_const()
@@ -22,11 +22,11 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
   # ** Constant Denominator Rules **
   # these rules strictly require y to be a scalar constant > 0
   if y.op is Ops.CONST and (c := y.arg) > 0:
-    # nested_div_mod: (x%(k*c))//c -> (x//c)%k, and (x%(k*c))%c -> x%c
-    if x.op is Ops.MOD and (k := x.src[1].divides(c)) is not None:
-      return x.src[0] // y % k if d.op is Ops.IDIV else x.src[0] % y
+    # canonicalize_mod_div: (x%(d*k))//d -> (x//d)%k, puts nested div/mod in div-first canonical form for recombine
+    if d.op is Ops.IDIV and x.op is Ops.MOD and x.src[1].op is Ops.CONST and x.vmin >= 0 and x.src[1].arg % c == 0:
+      return x.src[0] // y % x.ufix(x.src[1].arg // c)
 
-    # remove_nested_mod in sum: (a%4 + b)%2 -> (a+b)%2, requires non-negative sums
+    # remove_nested_mod: remove nested mod in case the inner mod is a multiple of the outer mod, example: (a%4 + b)%2 -> (a+b)%2
     if d.op is Ops.MOD and x.vmin >= 0:
       new_xs, changed = [], False
       for u in uops_no_const:
@@ -42,25 +42,28 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
 
     # fold_binary_numerator: fold if expression has one non-constant term that takes on two values
     if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1:
-      y1 = (cmod if d.op is Ops.MOD else cdiv)(factors[0]*v.vmin+const, c)
-      y2 = (cmod if d.op is Ops.MOD else cdiv)(factors[0]*v.vmax+const, c)
+      y1 = cmod(factors[0]*v.vmin+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmin+const, c)
+      y2 = cmod(factors[0]*v.vmax+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmax+const, c)
       return (y2-y1)*(v-v.vmin) + y1
 
     # fold_divmod_congruence: fold if a is congruent to an expression whose range is between 0 and c
     if not (x.vmin<0 and correct_divmod_folding):
       # when f%c == c//2, abs(r) == abs(r-c) is a tie, try both signs since either may fit in one period
-      rem_choices = [(r, r-c) if (r:=f%c)*2 == c else (min(r, r-c, key=abs),) for f in factors]
+      rem_choices = [((r:=f%c), r-c) if (r:=f%c)*2 == c else (min(r, r-c, key=abs),) for f in factors]
       for rems in itertools.product(*rem_choices):
         if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c:
           if d.op is Ops.MOD: return rem - rem.vmin//c*c
-          return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + const//c + rem.vmin//c
+          return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + (const-const%c+rem.vmin//c*c)//c
 
     # gcd_with_remainder: factor out common gcd from numerator
-    if x.vmin >= 0 and (g:=math.gcd(*factors, c)) > 1:
-      new_x = unwrap(x_peeled.divides(g)).simplify() + (const//g)%(c//g)
-      if new_x.vmin >= 0:
-        if d.op is Ops.MOD: return new_x % (c//g) * g + const%g
-        return new_x // (c//g) + const//c
+    # Note: this rule uses uops_no_const to exclude the additive constant from the GCD calculation
+    if x.vmin >= 0:
+      gcd = UOp.gcd(*uops_no_const, y).simplify()
+      if gcd.op is Ops.CONST and gcd.arg > 1:
+        new_x = unwrap(x_peeled.divide_exact(gcd)).simplify() + (const%c)//gcd.arg
+        if new_x.vmin >= 0:
+          ret = new_x.alu(d.op, x.ufix(c//gcd.arg))
+          return ret*gcd + const%gcd.arg if d.op is Ops.MOD else ret+const//c
 
     # nest_by_factor: x//c -> (x//f)//(c//f), x%c -> (x//f%(c//f))*f + b where b=x%f
     if x.vmin >= 0:
@@ -80,7 +83,7 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
   # ** Variable Denominator / Fallback Rules **
   # These rules apply to variables OR constants that failed the checks above.
   # Reconstruct all uops including const for these checks.
-  all_uops = list(x.split_uop(Ops.ADD))
+  all_uops = uops_no_const + ([x.const_like(const)] if const != 0 else [])
 
   # divide_by_gcd: x//y -> (x//gcd)//(y//gcd)
   gcd = UOp.gcd(*all_uops, y).simplify()

tinygrad/uop/ops.py

@@ -29,7 +29,7 @@ axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisTy
 range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.COPY: 2, Ops.BUFFER_VIEW: 1}
 
 # https://en.wikipedia.org/wiki/Identity_element
-def identity_element(op:Ops, dt:DType) -> PyConst: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dt.min}[op], dt)
+def identity_element(op:Ops, dt:DType) -> PyConst: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
 
 # With True as the default, this matches the old symbolic behavior
 def resolve(x:UOp|bool, default:bool=True):
@@ -301,10 +301,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
             raise ValueError(f"invalid type for axis: {axis_arg}")
           return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))
 
-    if self.op is Ops.STORE: return self.src[1]._shape
+    if self.op is Ops.ASSIGN: return self.src[1]._shape
 
     # 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}):
+    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}")
@@ -447,7 +447,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), (self, UOp.const(self.dtype, src) if not isinstance(src, UOp) else src), **kwargs)
   def end(self, *src:UOp): return UOp(Ops.END, src=(self,)+src) if len(src) else self
   def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs) if len(src) else self
-  def assign(self, x:UOp): return self.after(self.store(x))
+  def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
   def contract(self, *rngs:UOp):
     assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast"
@@ -843,8 +843,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if self.op is Ops.GEP: return self.src[0]._min_max
     # TODO: CAST to bool/unsigned is not monotone, still some case can be simplified
     if self.op is Ops.CAST and self.dtype in dtypes.floats+dtypes.sints+(dtypes.index,):
-      return max(self.dtype.min, self.src[0].vmin), min(self.src[0].vmax, self.dtype.max)
-    return self.dtype.min, self.dtype.max
+      return max(dtypes.min(self.dtype), self.src[0].vmin), min(self.src[0].vmax, dtypes.max(self.dtype))
+    return dtypes.min(self.dtype), dtypes.max(self.dtype)
 
   @functools.cached_property
   def _sym_fxn(self):
@@ -1051,12 +1051,12 @@ class UPat(OpMixin):
   def gep(self, i:int|None=None, **kwargs): return UPat(Ops.GEP, None, (self,), (i,) if i is not None else None, **kwargs)
   def load(self, *src:UPat, **kwargs): return UPat(Ops.LOAD, src=(self,)+src, **kwargs)
   def store(self, *src:UPat, **kwargs): return UPat(Ops.STORE, self.match_dtype, (self,)+src, **kwargs)
+  def assign(self, x:UPat, **kwargs): return UPat(Ops.ASSIGN, self.match_dtype, (self,x), **kwargs)
   def reduce(self, *src:UPat, **kwargs): return UPat(Ops.REDUCE, self.match_dtype, src=(self,)+src, **kwargs)
   def broadcast(self, **kwargs): return UPat(Ops.VECTORIZE, self.match_dtype, src=self, **kwargs)
   def contiguous(self, *args, **kwargs): return UPat(Ops.CONTIGUOUS, dtype=self.match_dtype, src=(self,)+args, **kwargs)
   def after(self, *src:UPat, **kwargs): return UPat(Ops.AFTER, self.match_dtype, (self,)+src, **kwargs)
   def end(self, *src:UPat, **kwargs): return UPat(Ops.END, self.match_dtype, (self,)+src, **kwargs)
-  def assign(self, x:UPat, **kwargs): return self.after(self.store(x), **kwargs)
 
   def const_like(self, b:ConstLike): return UPat.const(self.match_dtype, cast(ConstType, b))
   def alu(self, op:Ops, *src:UPat):

tinygrad/uop/spec.py

@@ -96,9 +96,6 @@ _tensor_spec = PatternMatcher([
   # ASSIGN has a target and a value. It can also optionally depend on other assigns
   (UPat(Ops.ASSIGN, name="x"), lambda x: len(x.src) >= 2 and all(s.op is Ops.ASSIGN for s in x.src[2:])),
 
-  # STORE in tensor graph: store a value into a target
-  (UPat(Ops.STORE, dtypes.void, (UPat(), UPat())), lambda: True),
-
   # MSELECT chooses one of the multi buffers
   (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
 
@@ -270,6 +267,8 @@ full_spec = PatternMatcher([
   # linearizer: outputs + intermediate KERNELs
   (UPat(Ops.CALL, dtype=dtypes.void), lambda: True),
 
+  # Invalid must have type Index
+  (UPat(Ops.CONST, arg=Invalid, name="x"), lambda x: x.dtype.scalar() == dtypes.index),
   # where on index in rhs position is fine
   (UPat(Ops.WHERE, dtype=dtypes.index, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.index))), lambda: True),
   # allow index dtype on a restricted set of UOps

tinygrad/uop/symbolic.py

@@ -25,48 +25,16 @@ def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
 invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
 invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
 
-def fold_add_divmod_recombine(x:UOp) -> UOp|None:
-  terms = list(x.split_uop(Ops.ADD))
-  for i,u in enumerate(terms):
-    if u.op is Ops.MOD and u.src[1].op is Ops.CONST: base, div, mul = u.src[0], u.src[1].arg, 1
-    elif u.op is Ops.MUL and u.src[1].op is Ops.CONST and (m:=u.src[0]).op is Ops.MOD and m.src[1].op is Ops.CONST:
-      base, div, mul = m.src[0], m.src[1].arg, u.src[1].arg
-    else: continue
-    for j,v in enumerate(terms):
-      if i == j: continue
-      if v.op is not Ops.MUL or v.src[1].op is not Ops.CONST or v.src[1].arg != div*mul: continue
-      q, exact = v.src[0], False
-      # (base%div)*mul + (base//div)*(div*mul) -> base*mul
-      if q.op is Ops.IDIV and q.src[1].op is Ops.CONST and q.src[1].arg == div: exact = q.src[0] is base
-      # ((base//d)%div)*mul + (base//(d*div))*(div*mul) -> (base//d)*mul
-      if not exact and base.op is Ops.IDIV and base.src[1].op is Ops.CONST:
-        exact = q.op is Ops.IDIV and q.src[1].op is Ops.CONST and q.src[0] is base.src[0] and q.src[1].arg == base.src[1].arg*div
-      if exact: return functools.reduce(operator.add, (t for k,t in enumerate(terms) if k not in (i,j)), base*mul)
-      # ((base//div)%d)*div + base%div -> base%(div*d)
-      if mul == 1 and div > 0 and q.op is Ops.MOD and q.src[1].op is Ops.CONST and (d:=q.src[1].arg) > 0 and q.src[0].op is Ops.IDIV:
-        if q.src[0].src[0] is base and q.src[0].src[1].op is Ops.CONST and q.src[0].src[1].arg == div:
-          return functools.reduce(operator.add, (t for k,t in enumerate(terms) if k not in (i,j)), base % (div*d))
-  return None
-
 # this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0
 propagate_invalid = PatternMatcher([
   # propagate invalid, push it past children
-  (invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype) if i.dtype is dtypes.index else None),
-  (UPat(GroupOp.Unary, src=(invalid_gate,), name="alu"), lambda cond,x,alu,i: cond.where(x.alu(alu.op), i)),
-  (UPat(GroupOp.Binary-GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i)),
-  (UPat(GroupOp.Binary-GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond.where(y.alu(alu.op,x), i)),
+  (invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype)),
+  *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i))
+    for op in GroupOp.Binary-GroupOp.Comparison),
   # TODO: when can this happen? and is it always safe to just drop invalid?
-  (UPat(GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i:
-     x.alu(alu.op,y) if i.dtype is dtypes.index else cond.where(x.alu(alu.op,y), i.cast(dtypes.bool))),
-  (UPat(GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i:
-     y.alu(alu.op,x) if i.dtype is dtypes.index else cond.where(y.alu(alu.op,x), i.cast(dtypes.bool))),
+  *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: x.alu(alu.op,y)) for op in GroupOp.Comparison),
   # alu with invalid -> invalid
-  (UPat(GroupOp.Unary, src=(invalid_pat,)), lambda i: i),
-  (UPat(GroupOp.Binary-GroupOp.Comparison, src=[invalid_pat, UPat()]), lambda i: i),
-  # normalize where(cond, Invalid, val) -> where(~cond, val, Invalid)
-  (UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not().where(val, i) if val.arg != Invalid else i),
-  (UPat(Ops.BITCAST, src=(invalid_pat,), name="bc"), lambda bc,i: i.cast(bc.dtype)),
-  (UPat(Ops.BITCAST, src=(invalid_gate,), name="bc"), lambda bc,cond,x,i: cond.where(x.bitcast(bc.dtype), i.bitcast(bc.dtype))),
+  *((invalid_pat.alu(op, UPat(dtype=dtypes.index)), lambda i: i) for op in GroupOp.Binary-GroupOp.Comparison),
 ])
 
 symbolic_simple = propagate_invalid + PatternMatcher([
@@ -78,8 +46,24 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat.var("x") // 1, lambda x: x),   # x//1 -> x
   (UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
   ((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base),  # (x%y)%y = -> x%y (rewritten with base for speed)
-  # variations of (x%c)+(x//c)*c = x
-  (UPat(Ops.ADD, dtype=dtypes.index, name="x"), fold_add_divmod_recombine),
+  # variations of (x%c)+(x//c)*c = x   TODO: add sorting to remove some variations
+  (UPat.var("x")%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"), lambda x,c: x), # (x%c)+(x//c)*c = x
+  ((UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c"),
+    lambda x,a,b,c: x//a if a.arg*c.arg==b.arg else None), # ((x//a)%c)+(x//a*c)*c = x//a. Note if a = 1 it degenerates to the one above
+  ((UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c1")*UPat.cvar("c2")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c3"),
+    lambda x,a,b,c1,c2,c3: x//a*c2 if c1.arg>0 and a.arg*c1.arg==b.arg and c1.arg*c2.arg==c3.arg else None),
+  ((UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3")+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"),
+    lambda x,c1,c2,c3: x*c2 if c1.arg*c2.arg==c3.arg else None), # (x%c1)*c2+(x//c1)*c3 = x*c2 if c1*c2==c3
+  ((UPat.var("y")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"))+UPat.var("x")%UPat.cvar("c"), lambda y,x,c: y+x),
+  ((UPat.var("y")+UPat.var("x")%UPat.cvar("c"))+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"), lambda y,x,c: y+x),
+  ((UPat.var("y")+(UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3"))+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"),
+    lambda y,x,c1,c2,c3: y+x*c2 if c1.arg*c2.arg==c3.arg else None),
+  ((UPat.var("y")+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"))+(UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3"),
+    lambda y,x,c1,c2,c3: y+x*c2 if c1.arg*c2.arg==c3.arg else None),
+  ((UPat.var("y")+(UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c1")*UPat.cvar("c2"))+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c3"),
+    lambda y,x,a,b,c1,c2,c3: y+x//a*c2 if c1.arg>0 and a.arg*c1.arg==b.arg and c1.arg*c2.arg==c3.arg else None),
+  ((UPat.var("y")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c3"))+(UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c1")*UPat.cvar("c2"),
+    lambda y,x,a,b,c1,c2,c3: y+x//a*c2 if c1.arg>0 and a.arg*c1.arg==b.arg and c1.arg*c2.arg==c3.arg else None),
   (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c", vec=False), lambda x,c: x if c.arg else c),
   (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c", vec=False), lambda x,c: c if c.arg else x),
   (UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x),
@@ -138,12 +122,6 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat.cvar("gate", vec=False).where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
   # a.where(b.where(c, d), d) -> (a & b).where(c, d)
   (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
-  # where over invalid -> invalid
-  (invalid_gate.where(UPat.var("a"), UPat.var("b")), lambda a,b,cond,x,i: i.cast(a.dtype)),
-  (invalid_pat.where(UPat.var("a"), UPat.var("b")), lambda a,b,i: i.cast(a.dtype)),
-  # reduce with invalid -> invalid
-  (UPat(Ops.REDUCE, src=(invalid_gate,), allow_any_len=True, name="r"), lambda r,cond,x,i: i.cast(r.dtype)),
-  (UPat(Ops.REDUCE, src=(invalid_pat,), allow_any_len=True, name="r"), lambda r,i: i.cast(r.dtype)),
 ])
 
 # ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ********
@@ -434,10 +412,6 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
   # fold gated LOAD/STORE
   (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"),
     lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0
-  (UPat(Ops.STORE, src=(UPat(), invalid_pat), allow_any_len=True), lambda i: UOp(Ops.NOOP)),
-  # store of where with invalid -> gated store
-  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="index"), UPat.var("cond").where(UPat.var("val"), invalid_pat)), allow_any_len=True, name="store"),
-   lambda index, cond, val, store, i: UOp.store(index.src[0].index(cond.where(index.src[1], UOp.invalid())), val, *store.src[2:])),
   ((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()),  # 1/(x^c) -> (1/x)^c
   ((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()),
   ((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)

tinygrad/viz/js/index.js

@@ -268,7 +268,7 @@ function setFocus(key) {
   const html = d3.select(".info").html("");
   if (eventType === EventTypes.EXEC) {
     const [n, _, ...rest] = e.arg.tooltipText.split("\n");
-    html.append(() => tabulate([["Name", colored(e.arg.label)], ["Duration", formatTime(e.width)], ["Start Time", formatTime(e.x)]]));
+    html.append(() => tabulate([["Name", colored(e.label)], ["Duration", formatTime(e.width)], ["Start Time", formatTime(e.x)]]));
     let group = html.append("div").classed("args", true);
     for (const r of rest) group.append("p").text(r);
     group = html.append("div").classed("args", true);
@@ -322,12 +322,11 @@ function setFocus(key) {
 }
 
 const EventTypes = { EXEC:0, BUF:1 };
-const GraphConfig = [{ pcolor:"#c9a8ff", unit:"B", fillColor:"#2B1B72"}, { pcolor:"#4fa3cc", unit:"Hz", fillColor:"#4fa3cc"}];
 
-async function renderProfiler(path, opts) {
+async function renderProfiler(path, unit, opts) {
   displaySelection("#profiler");
   // support non realtime x axis units
-  formatTime = opts.unit === "ms" ? formatMicroseconds : formatCycles;
+  formatTime = unit === "realtime" ? formatMicroseconds : formatCycles;
   if (data?.path !== path) { data = {tracks:new Map(), axes:{}, path, first:null, pcToShape:new Map()}; focusedDevice = null; focusedShape = null; }
   setFocus(focusedShape);
   // layout once!
@@ -379,12 +378,16 @@ async function renderProfiler(path, opts) {
       for (let j=0; j<eventsLen; j++) {
         const e = {name:strings[u32()], ref:optional(u32()), key:optional(u32()), st:u32(), dur:f32(), info:strings[u32()] || null};
         // find a free level to put the event
-        let depth = levels.findIndex(levelEt => e.st >= levelEt);
-        const et = e.st+Math.trunc(e.dur);
-        if (depth === -1) {
-          depth = levels.length;
-          levels.push(et);
-        } else levels[depth] = et;
+        let depth = 0;
+        if (opts.levelKey != null) { depth = opts.levelKey(e); levels[depth] = 0; }
+        else {
+          depth = levels.findIndex(levelEt => e.st >= levelEt);
+          const et = e.st+Math.trunc(e.dur);
+          if (depth === -1) {
+            depth = levels.length;
+            levels.push(et);
+          } else levels[depth] = et;
+        }
         if (depth === 0 || opts.colorByName) colorKey = e.name.split(" ")[0];
         if (!colorMap.has(colorKey)) {
           const color = typeof colors === "function" ? colors(colorKey)
@@ -426,15 +429,13 @@ async function renderProfiler(path, opts) {
       }
       div.style("height", levelHeight*levels.length+padding+"px").style("pointerEvents", "none");
     } else {
-      const linear = u8(), peak = u64();
-      const config = GraphConfig[linear];
+      const peak = u64();
       const timestamps = [], valueMap = new Map();
       // start by unpacking the raw events
       const memEvents = [];
       let x = 0, y = 0, shapeIdx = 0;
       const allocs = new Map();
       for (let j=0; j<eventsLen; j++) {
-        if (linear) { const ts = u32(), value = u64(); timestamps.push(ts); valueMap.set(ts, value); continue; }
         const alloc = u8(), ts = u32(), key = u32();
         if (alloc) {
           const dtype = strings[u32()], sz = u64(), nbytes = dtypeSize[dtype]*sz;
@@ -452,11 +453,11 @@ async function renderProfiler(path, opts) {
         }
       }
       timestamps.push(dur);
-      const height = linear ? (baseHeight-padding)*(opts.heightScale ?? 1)*2 : heightScale(peak);
+      const height = heightScale(peak);
       const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
       // generic polygon merger
       const base0 = yscale(0);
-      const sum = {x:[], y0:[], y1:[], fillColor:config.fillColor};
+      const sum = {x:[], y0:[], y1:[], fillColor:"#2B1B72"};
       for (let i=0; i<timestamps.length-1; i++) {
         const yv = yscale(valueMap.get(timestamps[i]));
         sum.x.push(timestamps[i], timestamps[i+1]); sum.y1.push(yv, yv); sum.y0.push(base0, base0);
@@ -495,10 +496,9 @@ async function renderProfiler(path, opts) {
         return bufShapes;
       };
       if (timestamps.length > 0) data.first = data.first == null ? timestamps[0] : Math.min(data.first, timestamps[0]);
-      data.tracks.set(k, { shapes:[sum], eventType, linear, visible, offsetY, pcolor:config.pcolor, height, peak, scaleFactor:maxheight*4/height,
-                           get views() { return [[sum], linear ? null : buildBufShapes()]; }, valueMap, rowBorderColor });
+      data.tracks.set(k, { shapes:[sum], eventType, visible, offsetY, pcolor:"#c9a8ff", height, peak, scaleFactor:maxheight*4/height,
+                           get views() { return [[sum], buildBufShapes()]; }, valueMap, rowBorderColor });
       div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
-        if (linear) return;
         const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
         let offset = 0;
         for (const [tid, track] of data.tracks) {
@@ -506,7 +506,7 @@ async function renderProfiler(path, opts) {
           if (tid === newFocus) { track.shapes = track.views[1]; offset += rescaleTrack(track, tid, track.scaleFactor); }
           else if (tid === focusedDevice) { track.shapes = track.views[0]; offset += rescaleTrack(track, tid, 1/track.scaleFactor); }
         }
-        data.axes.y = newFocus != null ? { domain:[0, (t=data.tracks.get(newFocus)).peak], range:[t.offsetY+t.height, t.offsetY], fmt:config.unit } : null;
+        data.axes.y = newFocus != null ? { domain:[0, (t=data.tracks.get(newFocus)).peak], range:[t.offsetY+t.height, t.offsetY], fmt:"B" } : null;
         toggleCls(document.getElementById(focusedDevice), document.getElementById(newFocus), "expanded");
         focusedDevice = newFocus;
         return resize();
@@ -545,27 +545,26 @@ async function renderProfiler(path, opts) {
     const visibleYStart = profilerEl.scrollTop-canvasTop + rect(profilerEl).top, visibleYEnd = visibleYStart+profilerEl.clientHeight;
     ctx.textBaseline = "middle";
     // draw shapes
-    for (const [k, { shapes, eventType, linear, visible, offsetY, valueMap, pcolor, scolor, rowBorderColor }] of data.tracks) {
+    for (const [k, { shapes, eventType, visible, offsetY, valueMap, pcolor, scolor, rowBorderColor }] of data.tracks) {
       visible.length = 0;
       const trackHeight = rect(document.getElementById(k)).height;
       if (offsetY+trackHeight < visibleYStart || offsetY > visibleYEnd) continue;
       const addBorder = scolor != null ? (w) => { if (w > 10) { ctx.strokeStyle = scolor; ctx.stroke(); } } : null;
-      const config = GraphConfig[linear];
       for (const e of shapes) {
         if (eventType === EventTypes.BUF) { // generic polygon
           if (e.x[0]>et || e.x.at(-1)<st) continue;
           ctx.beginPath();
           const x = e.x.map(xscale);
-          ctx.moveTo(x[0], offsetY+e.y1[0]);
+          ctx.moveTo(x[0], offsetY+e.y0[0]);
           for (let i=1; i<x.length; i++) {
-            ctx.lineTo(x[i], offsetY+e.y1[i]);
+            ctx.lineTo(x[i], offsetY+e.y0[i]);
             let arg = e.arg;
-            if (arg == null && valueMap != null) arg = {tooltipText: formatUnit(valueMap.get(e.x[i-1]), config.unit)}
+            if (arg == null && valueMap != null) arg = {tooltipText: `Total: ${formatUnit(valueMap.get(e.x[i-1]), 'B')}`}
             visible.push({ x0:x[i-1], x1:x[i], y0:offsetY+e.y1[i-1], y1:offsetY+e.y0[i], arg });
           }
-          if (linear) { ctx.strokeStyle = e.fillColor; ctx.lineWidth = 2; ctx.stroke(); ctx.lineWidth = 1; }
-          // walk the path back and fill the complete shape
-          else { for (let i=x.length-1; i>=0; i--) ctx.lineTo(x[i], offsetY+e.y0[i]); ctx.closePath(); ctx.fillStyle = e.fillColor; ctx.fill(); }
+          for (let i=x.length-1; i>=0; i--) ctx.lineTo(x[i], offsetY+e.y1[i]);
+          ctx.closePath();
+          ctx.fillStyle = e.fillColor; ctx.fill();
         } else { // contiguous rect
           if (e.x>et || e.x+e.width<st) continue;
           const x = xscale(e.x);
@@ -598,7 +597,7 @@ async function renderProfiler(path, opts) {
       const labelX = x+ctx.lineWidth+2;
       if (labelX <= lastLabelEnd) continue;
 
-      const label = formatTime(tick, et-st <= 1e3);
+      const label = formatTime(tick, et-st <= 1e3 ? true : false);
       ctx.textBaseline = "top";
       ctx.fillText(label, labelX, tickSize);
       lastLabelEnd = labelX + ctx.measureText(label).width + 4;
@@ -881,7 +880,7 @@ async function main() {
     if (url.pathname+url.search !== ckey) e.close();
     else if (e.readyState === EventSource.OPEN) activeSrc = e;
   }
-  if (ctx.name === "Profiler") return renderProfiler("/get_profile", {unit:"ms", width:"132px"});
+  if (ctx.name === "Profiler") return renderProfiler("/get_profile", "realtime", { width:"132px" });
   if (workerUrl == null) await initWorker();
   if (ckey in cache) {
     ret = cache[ckey];
@@ -899,8 +898,8 @@ async function main() {
     }
     // timeline with cycles on the x axis
     if (ret instanceof ArrayBuffer) {
-      const pkts = step.name.includes("PKTS");
-      return renderProfiler(ckey, {unit:"clk", heightScale:0.5, hideLabels:true, colorByName:pkts});
+      opts = {heightScale:0.5, hideLabels:true, levelKey:step.name.includes("PKTS") ? (e) => parseInt(e.name.split(" ")[1].split(":")[1]) : null, colorByName:ckey.includes("pkts")};
+      return renderProfiler(ckey, "clk", opts);
     }
     metadata.replaceChildren(...((ret.metadata ?? []).map((m) => {
       return tabulate(m.map((e) => [e.label.trim(), typeof e.value === "string" ? e.value : formatUnit(e.value)]));

tinygrad/viz/serve.py

@@ -239,16 +239,13 @@ def encode_mem_free(key:int, ts:int, execs:list[ProfilePointEvent], scache:dict)
     ei_encoding.append((e.key, enum_str(e.arg["name"], scache), num, mode))
   return struct.pack("<BIII", 0, ts, key, len(ei_encoding))+b"".join(struct.pack("<IIIB", *t) for t in ei_encoding)
 
-def graph_layout(k:str, dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
-                 scache:dict[str, int]) -> tuple[str, bytes|None]:
-  if k.startswith("LINE:"):
-    xy = [(rel_ts(e.ts, start_ts), e.key) for st,_,_,e in dev_events if isinstance(e, ProfilePointEvent)]
-    peaks.append(peak:=max([y for _,y in xy]))
-    return k.replace("LINE:", ""), struct.pack("<BIBQ", 1, len(xy), 1, peak)+b"".join(struct.pack("<IQ", x, y) for x,y in xy)
+def mem_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
+               scache:dict[str, int]) -> bytes|None:
   peak, mem = 0, 0
   temp:dict[int, int] = {}
   events:list[bytes] = []
   buf_ei:dict[int, list[ProfilePointEvent]] = {}
+
   for st,_,_,e in dev_events:
     if not isinstance(e, ProfilePointEvent): continue
     if e.name == "alloc":
@@ -265,7 +262,7 @@ def graph_layout(k:str, dev_events:list[tuple[int, int, float, DevEvent]], start
       mem -= temp.pop(e.key)
   for t in temp: events.append(encode_mem_free(t, rel_ts(end_ts, start_ts), buf_ei.pop(t, []), scache))
   peaks.append(peak)
-  return f"{k} Memory", struct.pack("<BIBQ", 1, len(events), 0, peak)+b"".join(events) if events else None
+  return struct.pack("<BIQ", 1, len(events), peak)+b"".join(events) if events else None
 
 # by default, VIZ does not start when there is an error
 # use this to instead display the traceback to the user
@@ -275,7 +272,7 @@ def soft_err(fn:Callable):
   except Exception: fn({"src":traceback.format_exc()})
 
 def row_tuple(row:str) -> tuple[tuple[int, int], ...]:
-  return ((0, 0),) if "Clock" in row else tuple((ord(ss[0][0]), int(ss[1])) if len(ss:=x.split(":"))>1 else (999,999) for x in row.split())
+  return tuple((ord(ss[0][0]), int(ss[1])) if len(ss:=x.split(":"))>1 else (999,999) for x in row.split())
 
 # *** Performance counters
 
@@ -339,29 +336,21 @@ def load_amd_counters(ctxs:list[dict], profile:list[ProfileEvent]) -> None:
 
 def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
   from tinygrad.renderer.amd.sqtt import map_insts, InstructionInfo, PacketType, INST, InstOp, VALUINST, IMMEDIATE, IMMEDIATE_MASK, VMEMEXEC, ALUEXEC
-  from tinygrad.renderer.amd.sqtt import INST_RDNA4, InstOpRDNA4, TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4
+  from tinygrad.renderer.amd.sqtt import INST_RDNA4, InstOpRDNA4
   ret:list[ProfileEvent] = []
-  row_ends:dict[str, Decimal] = {}
-  NS_PER_TICK = 10  # 100MHz
-  prev_pair:tuple[int, int]|None = None # (shader, realtime)
-  def add(name:str, p:PacketType, width=1, op:str|None=None, wave:int|None=None, info:InstructionInfo|None=None) -> None:
-    row = f"WAVE:{wave}" if (wave:=getattr(p, "wave", wave)) is not None else f"{p.__class__.__name__}:0 {name}"
-    ret.append(e:=ProfileRangeEvent(row, TracingKey(op or name, ret=f"PC:{info.pc}" if info else None), Decimal(p._time), Decimal(p._time+width)))
-    if (et:=row_ends.get(row)) is not None and e.st < et: raise RuntimeError(f"packet {p} overlaps another packet in {row}.")
-    row_ends[row] = unwrap(e.en)
+  rows:dict[str, None] = {}
+  trace:dict[str, set[int]] = {}
+  def add(name:str, p:PacketType, idx=0, width=1, op_name=None, wave=None, info:InstructionInfo|None=None) -> None:
+    if hasattr(p, "wave"): wave = p.wave
+    rows.setdefault(r:=(f"WAVE:{wave}" if wave is not None else f"{p.__class__.__name__}:0 {name}"))
+    key = TracingKey(f"{op_name if op_name is not None else name} OP:{idx}", ret=f"PC:{info.pc}" if info is not None else None)
+    ret.append(ProfileRangeEvent(r, key, Decimal(p._time), Decimal(p._time+width)))
   for p, info in map_insts(data, lib, target):
     if len(ret) > getenv("MAX_SQTT_PKTS", 50_000): break
-    if isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4)) and p.is_marker:
-      pair = (p._time, p.delta)
-      if prev_pair is None: prev_pair = pair
-      elif ret:
-        (s0, r0), (s1, r1) = prev_pair, pair
-        freq_hz = (s1 - s0) * 1_000_000_000 // ((r1 - r0) * NS_PER_TICK)
-        ret.append(ProfilePointEvent("LINE:Shader Clock", "freq_hz", freq_hz, ts=Decimal(p._time)))
-        prev_pair = pair
     if isinstance(p, (INST, INST_RDNA4)):
-      name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) else f"0x{p.op:02x}"
-      add(name, p, width=10 if "BARRIER" in name else 1, info=info)
+      op_name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) else f"0x{p.op:02x}"
+      name, width = (op_name, 10 if "BARRIER" in op_name else 1)
+      add(name, p, width=width, idx=int("OTHER" in name), info=info)
     if isinstance(p, (VALUINST, IMMEDIATE)): add(p.__class__.__name__, p, info=info)
     if isinstance(p, IMMEDIATE_MASK): add("IMMEDIATE", p, wave=unwrap(info).wave, info=info)
     if isinstance(p, (VMEMEXEC, ALUEXEC)):
@@ -370,9 +359,11 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
         add("VALU", p)
         add("SALU", p)
       else:
-        add(name.replace("_ALT", ""), p, op=name)
+        add(name.replace("_ALT", ""), p, op_name=name)
+      if p._time in trace.setdefault(name, set()): raise AssertionError(f"packets overlap in shared resource! {name}")
+      trace[name].add(p._time)
   pc_map = {addr:str(inst) for addr,inst in amd_decode(lib, target).items()}
-  return [ProfilePointEvent(r, "JSON", "pcMap", pc_map, ts=Decimal(0)) for r in row_ends]+ret
+  return [ProfilePointEvent(r, "JSON", "pcMap", pc_map, ts=Decimal(0)) for r in rows]+ret
 
 # ** SQTT OCC only unpacks wave start, end time and SIMD location
 
@@ -444,7 +435,7 @@ def get_profile(profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_
   for k,v in dev_events.items():
     v.sort(key=lambda e:e[0])
     layout[k] = timeline_layout(v, start_ts, scache)
-    layout.update([graph_layout(k, v, start_ts, unwrap(end_ts), peaks, dtype_size, scache)])
+    layout[f"{k} Memory"] = mem_layout(v, start_ts, unwrap(end_ts), peaks, dtype_size, scache)
   sorted_layout = sorted([k for k,v in layout.items() if v is not None], key=sort_fn)
   ret = [b"".join([struct.pack("<B", len(k)), k.encode(), unwrap(layout[k])]) for k in sorted_layout]
   index = json.dumps({"strings":list(scache), "dtypeSize":dtype_size, "markers":[{"ts":rel_ts(e.ts, start_ts), **e.arg} for e in markers],