now that needs rangeify 2 also

* fix bad range merges

* fix rng

* fix uop gc

.github/workflows/test.yml

@@ -259,14 +259,16 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-12
-        pydeps: "pillow"
+        pydeps: "pillow numpy ftfy regex"
         deps: testing_unit
     - name: Run unit tests
       run: python -m pytest -n=auto test/unit/ --durations=20
     - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 test/test_multitensor.py TestMultiTensor.test_data_parallel_resnet_train_step
+      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
     - name: Run SDXL on NULL backend
       run: MAX_BUFFER_SIZE=0 NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
+    - name: Run Clip tests for SD MLPerf on NULL backend
+      run: MAX_BUFFER_SIZE=0 NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
     # TODO: support fake weights
     #- name: Run LLaMA 7B on 4 fake devices
     #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -376,7 +378,7 @@ jobs:
         run: |
           ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot model with rangeify
-        run: FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+        run: RANGEIFY=1 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot alt model correctness (float32)
         run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot fastvits model correctness (float32)
@@ -525,23 +527,21 @@ jobs:
         llvm: "true"
     - name: Test CPU=1 RANGEIFY=1
       # TODO: add more passing tests here
-      # test_instancenorm_3d is very slow
+      # rangeify diamond cycle gives the wrong answer
       run: |
         CPU=1 CPU_LLVM=0 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
-          -k "not test_instancenorm_3d" \
+          -k "not test_assign_diamond_cycle" \
           test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_symbolic_ops.py test/test_symbolic_jit.py test/test_tensor_variable.py \
           test/test_outerworld_range.py test/test_randomness.py test/test_nn.py test/test_arange.py test/test_tensor.py test/test_optim.py \
-          test/test_setitem.py
-    - name: Test const folding
-      run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestReduceOpsConstFolding and not TestMultiConstFolding"
-    - name: Test multitensor
+          test/test_setitem.py test/test_assign.py test/test_multitensor.py
+    - name: Test CPU=1 CPU_LLVM=1 RANGEIFY=1
       run: |
-        CPU=1 RANGEIFY=1 python3 test/test_multitensor.py TestMultiTensor.test_matmul_shard_1_1 TestMultiTensor.test_simple_add_W TestMultiTensor.test_simple_reduce \
-        TestMultiTensor.test_elementwise_dtype TestMultiTensor.test_shard_no_recompile TestHandleData.test_copied_to_device TestMultiRamUsage
-        CPU=1 RANGEIFY=1 python3 -m pytest test/test_multitensor.py::TestMultiAssign -k 'not (multi_assign_piece_noncontig or multi_assign_var_offset)'
-        CPU=1 RANGEIFY=1 python3 -m pytest -n=auto test/test_multitensor.py::TestMultiTensor test/unit/test_allreduce.py -k 'not const_folding'
-    - name: Test CPU=1 RANGEIFY=2
-      run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
+        CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_edgecases.py
+    - name: Test const folding
+      run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestReduceOpsConstFolding"
+    # RANGEIFY=2 isn't supported
+    #- name: Test CPU=1 RANGEIFY=2
+    #  run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
     # slow (and still wrong on beautiful_mnist)
     #- name: Test LLVM RANGEIFY=1 (slow tests)
     #  run: CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
@@ -584,7 +584,7 @@ jobs:
         key: metal
         deps: testing
     - name: some unit tests
-      run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py --durations=20
+      run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py test/unit/test_linalg.py --durations=20
     - name: Test METAL=1 RANGEIFY=1
       run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py --durations=20
     - name: Run process replay tests

examples/gpt2.py

@@ -26,8 +26,8 @@ class Attention:
       start_pos = start_pos.val
 
     if HALF: x = x.half()
-    xqkv = self.c_attn(x)
-    xq, xk, xv = [xqkv.shrink((None, None, (i*self.dim, (i+1)*self.dim))).reshape(None, None, self.n_heads, self.head_dim) for i in range(3)]
+    xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
+    xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
     bsz, seqlen, _, _ = xq.shape
 
     # create kv cache
@@ -35,11 +35,11 @@ class Attention:
       self.cache_kv = Tensor.zeros(2, bsz, MAX_CONTEXT, self.n_heads, self.head_dim, dtype=x.dtype).contiguous().realize()
 
     # update the cache
-    self.cache_kv.shrink((None, None,(start_pos,start_pos+seqlen),None,None)).assign(Tensor.stack(xk, xv)).realize()
+    self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
 
     if start_pos > 0:
-      keys = self.cache_kv[0].shrink((None, (0, start_pos+seqlen), None, None))
-      values = self.cache_kv[1].shrink((None, (0, start_pos+seqlen), None, None))
+      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
+      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
     else:
       keys = xk
       values = xv
@@ -64,7 +64,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
     h = x + self.attn(self.ln_1(x), start_pos, mask).float()
-    return (h + self.mlp(self.ln_2(h)))
+    return (h + self.mlp(self.ln_2(h))).contiguous()
 
 class Transformer:
   def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):

examples/hlb_cifar10.py

@@ -229,7 +229,8 @@ def train_cifar():
     if getenv("RANDOM_CROP", 1):
       X = random_crop(X, crop_size=32)
     if getenv("RANDOM_FLIP", 1):
-      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
+      # NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
+      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
     X, Y = X[perms], Y[perms]
     return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])
 

examples/mlperf/initializers.py

@@ -17,6 +17,10 @@ def he_normal(*shape, a: float = 0.00, **kwargs) -> Tensor:
   std = math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:])) / 0.87962566103423978
   return std * rand_truncn(*shape, **kwargs)
 
+# Stable Diffusion v2 training uses default torch gelu, which doesn't use tanh approximation
+def gelu_erf(x:Tensor) -> Tensor:
+  return 0.5 * x * (1.0 + (x / 1.4142135623730951).erf())
+
 class Conv2dHeNormal(nn.Conv2d):
   def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
     super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)

extra/models/clip.py

@@ -9,6 +9,9 @@ from PIL import Image
 import numpy as np
 import re, gzip
 
+# Allow for monkeypatching for mlperf.
+gelu = Tensor.gelu
+
 @lru_cache()
 def default_bpe():
   # Clip tokenizer, taken from https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py (MIT license)
@@ -53,8 +56,8 @@ class Tokenizer:
     cs = [chr(n) for n in cs]
     return dict(zip(bs, cs))
   class ClipTokenizer:
-    def __init__(self):
-      self.byte_encoder = Tokenizer.bytes_to_unicode()
+    def __init__(self, version=None):
+      self.byte_encoder, self.version = Tokenizer.bytes_to_unicode(), version
       merges = gzip.open(default_bpe()).read().decode("utf-8").split('\n')
       merges = merges[1:49152-256-2+1]
       merges = [tuple(merge.split()) for merge in merges]
@@ -62,11 +65,17 @@ class Tokenizer:
       vocab = vocab + [v+'</w>' for v in vocab]
       for merge in merges:
         vocab.append(''.join(merge))
-      vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+      if self.version == "sd_mlperf_v5_0":
+        import regex
+        vocab.extend(['<start_of_text>', '<end_of_text>'])
+        self.cache = {'<start_of_text>': '<start_of_text>', '<end_of_text>': '<end_of_text>'}
+        self.pat = regex.compile(r"""<start_of_text>|<end_of_text>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", regex.IGNORECASE)
+      else:
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)
       self.encoder = dict(zip(vocab, range(len(vocab))))
       self.bpe_ranks = dict(zip(merges, range(len(merges))))
-      self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
-      self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)
 
     def bpe(self, token):
       if token in self.cache:
@@ -110,8 +119,17 @@ class Tokenizer:
 
     def encode(self, text:str, pad_with_zeros:bool=False) -> List[int]:
       bpe_tokens: List[int] = []
-      text = Tokenizer.whitespace_clean(text.strip()).lower()
-      for token in re.findall(self.pat, text):
+      if self.version == "sd_mlperf_v5_0":
+        import regex, ftfy, html
+        text = ftfy.fix_text(text)
+        text = html.unescape(html.unescape(text)).strip()
+        text = Tokenizer.whitespace_clean(text).lower()
+        re_module = regex
+      else:
+        text = Tokenizer.whitespace_clean(text.strip()).lower()
+        re_module = re
+
+      for token in re_module.findall(self.pat, text):
         token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
         bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
       # Truncation, keeping two slots for start and end tokens.
@@ -252,10 +270,8 @@ class Open:
       q,k,v = [y.reshape(T, B*self.n_heads, self.d_head).transpose(0, 1).reshape(B, self.n_heads, T, self.d_head) for y in proj.chunk(3)]
 
       attn_output = Tensor.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
-      attn_output = attn_output.permute(2, 0, 1, 3).reshape(T*B, C)
-
+      attn_output = attn_output.permute(2, 0, 1, 3).reshape(T, B, C)
       attn_output = self.out_proj(attn_output)
-      attn_output = attn_output.reshape(T, B, C)
 
       return attn_output
 
@@ -263,9 +279,10 @@ class Open:
     def __init__(self, dims, hidden_dims):
       self.c_fc   = Linear(dims, hidden_dims)
       self.c_proj = Linear(hidden_dims, dims)
+      self.gelu = gelu
 
     def __call__(self, x:Tensor) -> Tensor:
-      return x.sequential([self.c_fc, Tensor.gelu, self.c_proj])
+      return x.sequential([self.c_fc, self.gelu, self.c_proj])
 
   # https://github.com/mlfoundations/open_clip/blob/58e4e39aaabc6040839b0d2a7e8bf20979e4558a/src/open_clip/transformer.py#L210
   class ResidualAttentionBlock:
@@ -350,15 +367,15 @@ class Open:
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L396
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L498
 class FrozenOpenClipEmbedder(Embedder):
-  def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False):
-    self.tokenizer = Tokenizer.ClipTokenizer()
+  def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False, clip_tokenizer_version=None):
+    self.tokenizer = Tokenizer.ClipTokenizer(version=clip_tokenizer_version)
     self.model = Open.ClipTextTransformer(dims, n_heads, layers)
     self.return_pooled = return_pooled
     self.input_key = "txt"
     self.ln_penultimate = ln_penultimate
 
   def tokenize(self, text:str, device:Optional[str]=None) -> Tensor:
-    return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int64, device=device).reshape(1,-1)
+    return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int32, device=device).reshape(1,-1)
 
   def text_transformer_forward(self, x:Tensor, attn_mask:Optional[Tensor]=None):
     for r in self.model.transformer.resblocks:
@@ -449,7 +466,7 @@ class OpenClipEncoder:
     x = x + self.positional_embedding
     x = self.transformer(x, attn_mask=self.attn_mask)
     x = self.ln_final(x)
-    x = x[:, tokens.argmax(axis=-1)]
+    x = x[Tensor.arange(x.shape[0], device=x.device), tokens.argmax(axis=-1)]
     x = x @ self.text_projection
     return x
 

extra/remu/src/thread.rs

@@ -930,7 +930,7 @@ impl<'a> Thread<'a> {
 
             let op = ((instr >> 16) & 0x3ff) as u32;
             match op {
-                764 | 765 | 288 | 289 | 290 | 766 | 768 | 769 => {
+                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
                     let vdst = (instr & 0xff) as usize;
                     let sdst = ((instr >> 8) & 0x7f) as usize;
                     let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
@@ -944,6 +944,16 @@ impl<'a> Thread<'a> {
                     assert_eq!(clmp, 0);
 
                     let vcc = match op {
+                        767 => {
+                            let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
+                            let (mul_result, overflow_mul) = (s0 as i64).overflowing_mul(s1 as i64);
+                            let (ret, overflow_add) = mul_result.overflowing_add(s2 as i64);
+                            let overflowed = overflow_mul || overflow_add;
+                            if self.exec.read() {
+                                self.vec_reg.write64(vdst, ret as u64);
+                            }
+                            overflowed
+                        },
                         766 => {
                             let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
                             let (mul_result, overflow_mul) = (s0 as u64).overflowing_mul(s1 as u64);

test/device/test_null.py

@@ -0,0 +1,13 @@
+import unittest
+from tinygrad import dtypes, Device
+from tinygrad.device import is_dtype_supported
+
+@unittest.skipUnless(Device.DEFAULT=="NULL", "Don't run when testing non-NULL backends")
+class TestNULLSupportsDTypes(unittest.TestCase):
+  def test_null_supports_ints_floats_bool(self):
+    dts = dtypes.ints + dtypes.floats + (dtypes.bool,)
+    not_supported = [dt for dt in dts if not is_dtype_supported(dt, "NULL")]
+    self.assertFalse(not_supported, msg=f"expected these dtypes to be supported by NULL: {not_supported}")
+
+if __name__ == "__main__":
+  unittest.main()

test/external/external_uop_gc.py

@@ -63,6 +63,7 @@ if __name__ == "__main__":
     views_to_valid_uop.cache_clear()
 
     new_uops = uops_allocated()
+    print_uops()
     gc.collect()
     new_uops_gc = uops_allocated()
     print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")

test/external/mlperf_stable_diffusion/external_test_models.py

@@ -0,0 +1,53 @@
+import unittest
+from tinygrad import Tensor, dtypes, Device
+from tinygrad.nn.state import get_parameters
+from extra.models import clip
+from examples.mlperf.initializers import gelu_erf
+Device.DEFAULT="NULL"
+GPUS = [f"NULL:{i}" for i in range(8)]
+
+clip_params = {"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True, "clip_tokenizer_version": "sd_mlperf_v5_0"}
+def get_cond_stage_model(GPUS:list[str]|None=None) -> clip.FrozenOpenClipEmbedder:
+  clip.gelu = gelu_erf
+  model = clip.FrozenOpenClipEmbedder(**clip_params)
+  if GPUS and len(GPUS) > 1:
+    for p in get_parameters(model): p.to_(GPUS)
+  return model
+def get_tokens(BS:int) -> Tensor: return Tensor([0] * 77 * BS, dtype=dtypes.int32).reshape(-1, 77)
+
+class TestOpenClip(unittest.TestCase):
+  def test_tokenizer(self):
+    prompt = "Beautiful is better than ugly.\nExplicit is better than implicit.\nSimple is better than complex.\nComplex is better than complicated."
+    model = get_cond_stage_model()
+    tokens = model.tokenizer.encode(prompt, pad_with_zeros=True)
+    expected = [49406, 1215, 533, 1539, 1126, 8159, 269, 33228, 533, 1539, 1126, 15269, 585, 269, 4129, 533, 1539, 1126, 6324, 269, 6324, 533,
+                1539, 1126, 16621, 269, 49407] + [0]*50
+    self.assertEqual(tokens, expected)
+
+  def test_clip_gelu_init(self):
+    for resblock in get_cond_stage_model().model.transformer.resblocks:
+      self.assertEqual(resblock.mlp.gelu, gelu_erf)
+
+  def test_multigpu_clip_embed(self):
+    BS = 304
+    model = get_cond_stage_model(GPUS)
+    tokens = get_tokens(BS)
+    embeds = model.embed_tokens(tokens.shard(GPUS, axis=0)).realize()
+    self.assertEqual(embeds.shape, (BS, 77, 1024))
+    self.assertEqual(embeds.dtype, dtypes.float32)
+
+  def test_multigpu_clip_score(self):
+    BS = 240
+    vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
+    text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
+    clip.gelu = gelu_erf
+    clip_encoder = clip.OpenClipEncoder(1024, text_cfg, vision_cfg)
+    for p in get_parameters(clip_encoder): p.to_(GPUS)
+    tokens = get_tokens(BS)
+    imgs = Tensor.zeros(BS,3,224,224).contiguous()
+    scores = clip_encoder.get_clip_score(tokens.shard(GPUS, axis=0), imgs.shard(GPUS, axis=0)).realize()
+    self.assertEqual(scores.shape, (BS,))
+    self.assertEqual(scores.dtype, dtypes.float32)
+
+if __name__=="__main__":
+  unittest.main()

test/models/test_real_world.py

@@ -114,6 +114,16 @@ class TestRealWorld(unittest.TestCase):
 
       helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 93)
 
+  @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
+  def test_forward_cifar(self):
+    BS = 32
+    # with training batchnorm still though
+    with Tensor.train():
+      model = SpeedyResNet(Tensor.ones((12,3,2,2)))
+      @TinyJit
+      def run(X): return model(X)
+      helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, (1.0/48)*BS, 126)
+
   @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
   def test_train_cifar(self):
     with Tensor.train():

test/test_assign.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python
 import unittest
+import contextlib
 import numpy as np
 from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import temp
+from tinygrad.helpers import temp, RANGEIFY
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -254,6 +255,8 @@ class TestAssign(unittest.TestCase):
     b.assign(a.contiguous()).realize()
     assert GlobalCounters.kernel_count - kc == 2
 
+  # passing in RANGEIFY=1, RANGEIFY=0 asserts permuted assigns it can't fuse
+  def assert_permuted_assign(self): return self.assertRaisesRegex(RuntimeError, "contiguous") if not RANGEIFY else contextlib.nullcontext()
   def test_permuted_assignment(self):
     a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
     b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
@@ -277,7 +280,7 @@ class TestAssign(unittest.TestCase):
     #GlobalCounters.cache = []
     ba1 = a.uop.base.realized # noqa: F841
     bb1 = b.uop.base.realized # noqa: F841
-    with self.assertRaisesRegex(RuntimeError, "contiguous"):
+    with self.assert_permuted_assign():
       a.assign(a.permute(1,0) + b)   # this should not work!
       a.realize()
       ba2 = a.uop.base.realized # noqa: F841
@@ -285,6 +288,22 @@ class TestAssign(unittest.TestCase):
       #assert ba1 == ba2 and ba1 != bb1
       np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
 
+  @unittest.skipUnless(RANGEIFY, "only correct in rangeify")
+  def test_post_permuted_assignment_alt(self):
+    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    new_a = (a.T+b).numpy()
+    a.assign(a.T+b)
+    np.testing.assert_allclose(a.numpy(), new_a)
+
+  def test_post_reshape_assignment_fine(self):
+    a = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
+    b = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
+    rhs = a.reshape(-1).reshape(N, N)
+    new_a = (rhs+b).numpy()
+    a.assign(rhs+b)  # self-assign with reshape view is fine
+    np.testing.assert_allclose(a.numpy(), new_a)
+
   @unittest.skip("multi output not supported anymore")
   def test_simple_assignment_multioutput(self):
     a = Tensor.randn(32, 32).realize()
@@ -309,8 +328,8 @@ class TestAssign(unittest.TestCase):
   def test_permuted_assignment_correct(self):
     a = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
     b = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
-    # TODO: scheduler limitation, should NOT raise AssertionError from numpy.
-    with self.assertRaisesRegex(RuntimeError, "contiguous"):
+    # TODO: swizzler.py limitation, should NOT raise AssertionError from numpy.
+    with self.assert_permuted_assign():
       a = a.permute(1, 0)
       new_val = a + b
       a.assign(new_val)
@@ -319,10 +338,11 @@ class TestAssign(unittest.TestCase):
   def test_permuted_reduceop_child_dual_use(self):
     a = Tensor.randn(32, 32, 32).realize()
     b = Tensor.full((32, 32), 1.).contiguous().realize()
-    with self.assertRaisesRegex(RuntimeError, "contiguous"):
+    with self.assert_permuted_assign():
       r = a.sum(axis=1)
       b.assign(r + b.permute(1, 0))
       b.realize()
+      np.testing.assert_allclose(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32)).transpose(1, 0), atol=1e-6, rtol=1e-3)
 
   @unittest.skip("multi output not supported anymore")
   def test_permuted_reduceop_multioutput_dual_use(self):
@@ -359,15 +379,17 @@ class TestAssign(unittest.TestCase):
     a.assign(a + b)
     kc = GlobalCounters.kernel_count
     a.realize()
-    assert GlobalCounters.kernel_count - kc == 1
+    # rangeify makes two kernels
+    assert GlobalCounters.kernel_count - kc == (2 if RANGEIFY else 1)
     np.testing.assert_equal(a.numpy(), np.ones((4, 4))+np.pad(np.ones((4, 4))[:, 0:2], ((0, 0), (0, 2)), constant_values=2))
 
   def test_permuted_assignment_masked_view_not_contiguous(self):
     a = Tensor.ones(4, 4).contiguous().realize()
-    with self.assertRaisesRegex(RuntimeError, "contiguous"):
+    with self.assert_permuted_assign():
       b = a.shrink((None, (0, 2))).pad((None, (0, 2)), value=2).permute(1, 0)
       a.assign(a + b)
       a.realize()
+      self.assertListEqual(a.tolist(), [[2.,2.,2.,2.],[2.,2.,2.,2.],[3.,3.,3.,3.], [3.,3.,3.,3.]])
 
   # TODO: is there a way to sneak in a permute such that it returns the wrong answer?
 

test/test_const_folding.py

@@ -245,7 +245,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
         t = Tensor.ones(16, dtype=dt).reshape(4, 4)
         assert t.sum().dtype == t.contiguous().sum().dtype
 
-@unittest.skipIf(not_support_multi_device(), "no multi")
+@unittest.skipIf(not_support_multi_device() or RANGEIFY, "no multi, RANGEIFY doesn't support multi const folding")
 class TestMultiConstFolding(unittest.TestCase):
   def test_multi_const_folding_literal(self):
     ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))

test/test_jit.py

@@ -833,5 +833,20 @@ class TestJitGraphSplit(unittest.TestCase):
       multigraph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
       hcqgraph=[self.ji_graph(4)])
 
+class TestJitRandom(unittest.TestCase):
+  def test_jit_rangeify(self):
+    tst = {0:[], 1:[]}
+    for r in [0,1]:
+      Tensor.manual_seed(1337)
+      with Context(RANGEIFY=r):
+        _ = Tensor.randint(4, high=3)
+        # this second one makes the behavior different
+        _ = Tensor.randint(4, high=3)
+        @TinyJit
+        def f(): return Tensor.randint(20, high=5)
+        for _ in range(5): tst[r].append(f().tolist())
+    for i, (t0, t1) in enumerate(zip(tst[0], tst[1])):
+      self.assertListEqual(t0, t1, msg=f"mismatch at list {i}")
+
 if __name__ == '__main__':
   unittest.main()

test/test_linearizer.py

@@ -123,6 +123,7 @@ class TestLinearizer(unittest.TestCase):
     assert num_loads <= 4, "more load uops than needed"
     assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
 
+  @unittest.skip("this is handled at higher level now")
   def test_upcast_cse(self):
     # when upcasting, within a subtree, there may be common expressions.
 

test/test_linearizer_dumb.py

@@ -15,18 +15,18 @@ class TestLinearizerFailure(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
   def test_failure_beam_mnist(self):
     c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(4014080), arg=0, src=())
-    c1 = UOp.range(UOp.const(dtypes.int, 512), 0, AxisType.GLOBAL)
-    c2 = UOp.range(UOp.const(dtypes.int, 784), 1, AxisType.GLOBAL)
-    c3 = UOp.range(UOp.const(dtypes.int, 10), 3, AxisType.GLOBAL)
+    c1 = UOp.range(UOp.const(dtypes.index, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.index, 784), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.index, 10), 3, AxisType.GLOBAL)
     c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
-    c5 = c4.index(c1, UOp.const(dtypes.bool, True)).load()
-    c6 = UOp.range(UOp.const(dtypes.int, 6000), 1004, AxisType.REDUCE)
-    c7 = UOp.range(UOp.const(dtypes.int, 3750), 2006, AxisType.REDUCE)
-    c8 = UOp.range(UOp.const(dtypes.int, 16), 2007, AxisType.GROUP_REDUCE)
+    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True))).load()
+    c6 = UOp.range(UOp.const(dtypes.index, 6000), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(dtypes.index, 3750), 2006, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(dtypes.index, 16), 2007, AxisType.GROUP_REDUCE)
     c9 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(47040000), arg=2, src=())
-    c10 = c9.index((((c3*UOp.const(dtypes.int, 4704000))+c2)+(c6*UOp.const(dtypes.int, 784))), UOp.const(dtypes.bool, True)).load()
-    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.int, 6000))+c6)+((c7*UOp.const(dtypes.int, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.int, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
-    c12 = c0.index((((c1*UOp.const(dtypes.int, 7840))+(c2*UOp.const(dtypes.int, 10)))+c3), UOp.const(dtypes.bool, True)).store(c11, c1, c2, c3)
+    c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True))).load()
+    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
+    c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11, c1, c2, c3)
     ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
     _ = get_program(ast, Device["METAL"].renderer)
 

test/test_multitensor.py

@@ -54,6 +54,17 @@ class TestMultiTensor(unittest.TestCase):
       assert lb.shape == (128,)
     (X + X).realize()
 
+  def _test_shard_op(self, op, out, n=4):
+    t = Tensor.ones(n).contiguous().realize().shard(devices_2, 0)
+    r = op(t).realize()
+    assert t.uop.is_realized, "shard didn't realize"
+    self.assertEqual(r.tolist(), out)
+  def test_shard_reshape(self): self._test_shard_op(lambda t:t.reshape(2, 2), [[1.,1.],[1.,1.]])
+  def test_shard_elementwise(self): self._test_shard_op(lambda t:(t+t).reshape(2, 2), [[2.,2.],[2.,2.]])
+  def test_shard_reduce(self):
+    self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=1), [3.,3.], n=6)
+    self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=0), [2.,2.,2.], n=6)
+
   def test_shard_not_multiple(self):
     X = Tensor.ones(256).contiguous().realize()
     with self.assertRaises(RuntimeError):
@@ -782,6 +793,7 @@ class TestMultiTensor(unittest.TestCase):
     t = Tensor.rand(16, 16).shard(devices_2, axis=0)
     np.testing.assert_allclose(t.numpy(), t.clone().numpy())
 
+  @unittest.skipIf(RANGEIFY, "RANGEIFY doesn't support multi const folding")
   def test_multi_const_folding(self):
     with Context(TRACK_MATCH_STATS=0):
       a = Tensor.arange(3).realize()

test/test_ops.py

@@ -3164,6 +3164,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf))
     helper_test_op([(32,10)], lambda x: x.masked_fill((x<0.1).detach(), -math.inf))
 
+  @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "AMD" and RANGEIFY, "very slow on MOCKGPU because reduce does not fold")
   def test_masked_select(self):
     helper_test_op([(32, 10)], lambda x: x.masked_select(x>0.5), lambda x: x.masked_select(x>0.5), forward_only=True)
     helper_test_op([(32, 10)], lambda x: x.masked_select(torch.tensor(True)), lambda x: x.masked_select(Tensor(True)), forward_only=True)

test/test_pickle.py

@@ -2,7 +2,7 @@ import unittest, pickle, types
 import numpy as np
 from tinygrad import Tensor, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context
-from tinygrad.uop.ops import PatternMatcher, UPat, UOp, Ops
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp
 
 class TestPickle(unittest.TestCase):
   def test_pickle_code_object(self):
@@ -45,10 +45,9 @@ class TestPickle(unittest.TestCase):
     t_values = t.numpy()
     del t # free buffers
     print("** post pickle")
-    init = GlobalCounters.kernel_count
     t2:Tensor = pickle.loads(st)
+    assert t2.uop.is_realized
     np.testing.assert_equal(t_values, t2.numpy())
-    self.assertEqual(GlobalCounters.kernel_count-init, 0)
 
   def test_pickle_realized_tensor_alt2(self):
     print("** init")
@@ -70,14 +69,14 @@ class TestPickle(unittest.TestCase):
   def test_pickle_buffer_uop(self):
     t = Tensor.arange(4).realize()
     a = t.uop
-    assert a.op is Ops.BUFFER
-    self.assertIsNotNone(buffer:=a.realized)
+    assert a.is_realized
+    self.assertIsNotNone(buffer:=a.base.realized)
     s = pickle.dumps(a)
     # free buffers
     del a
     del buffer
     a2:UOp = pickle.loads(s)
-    self.assertListEqual(a2.realized.as_buffer().cast("I").tolist(), [0, 1, 2, 3])
+    self.assertListEqual(a2.base.realized.as_buffer().cast("I").tolist(), [0, 1, 2, 3])
 
   def test_pickle_unrealized_tensor(self):
     t = Tensor.ones(10, 10)

test/test_rangeify.py

@@ -15,11 +15,36 @@ class TestRangeifyAssign(unittest.TestCase):
     print(lst)
     print(lst2)
     print(lst3)
+    self.assertListEqual(lst, lst3)
+    self.assertListEqual(lst2, B.permute(1, 0).tolist())
 
 N = 256
 
+class TestRangeifyOpt(unittest.TestCase):
+  def test_randperm(self):
+    Tensor.randperm(10000).realize()
+
+  def test_one_getitem(self):
+    X = Tensor.empty(10000)
+    sel = Tensor.arange(1000).contiguous().realize()
+    Xsel = X[sel]
+    Tensor.realize(Xsel)
+
+  def test_two_getitem(self):
+    # this is splitting on the child even when it really shouldn't
+    X = Tensor.empty(10000)
+    Y = Tensor.empty(10000)
+    sel = Tensor.arange(1000).contiguous().realize()
+    Xsel, Ysel = X[sel], Y[sel]
+    Tensor.realize(Xsel, Ysel)
+
 @unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
 class TestRangeify(unittest.TestCase):
+  def test_groupnorm(self):
+    # ranges 1 and 3 are merging
+    x = nn.GroupNorm(32, 128)
+    x(Tensor.empty(1, 128, 64, 64)).realize()
+
   def test_expand_children(self):
     A = Tensor.empty(N, N).sum(axis=1)
     ba = A.expand(N, N)
@@ -57,6 +82,14 @@ class TestRangeify(unittest.TestCase):
     C = Tensor.empty(N, N)
     (((A@B).exp()@C).exp()).realize()
 
+  def test_double_gemm_exp_child(self):
+    A = Tensor.empty(N, N)
+    B = Tensor.empty(N, N)
+    C = Tensor.empty(N, N)
+    # A@B is used with exp, and also on the sum. this is two kernels now, is this right?
+    ret = A@B
+    ((ret.exp()@C)+ret).realize()
+
   def test_double_gemm_relu(self):
     A = Tensor.empty(N, N)
     B = Tensor.empty(N, N)
@@ -95,6 +128,11 @@ class TestRangeify(unittest.TestCase):
     w1 = Tensor.empty(8, 4, 3, 3)
     x.conv2d(w1).realize()
 
+  def test_conv2d_elu(self):
+    x = Tensor.empty(1, 4, 32, 32)
+    w1 = Tensor.empty(8, 4, 3, 3)
+    x.conv2d(w1).elu().realize()
+
   def test_conv2d_t(self):
     x = Tensor.empty(1, 4, 32, 32)
     w1 = Tensor.empty(8, 4, 3, 3)
@@ -106,6 +144,13 @@ class TestRangeify(unittest.TestCase):
     w2 = Tensor.empty(12, 8, 3, 3)
     x.conv2d(w1).conv2d(w2).realize()
 
+  def test_xception_conv2d(self):
+    # NOTE: this fusion is bad, it's recomputing the inner many times
+    x = Tensor.empty(1, 4, 32, 32)
+    w1 = Tensor.empty(8, 4, 1, 1)
+    w2 = Tensor.empty(8, 1, 3, 3)
+    x.conv2d(w1).conv2d(w2, groups=8).realize()
+
   def test_conv_maxpool_contig(self): self.test_conv_maxpool(True)
   def test_conv_maxpool(self, contig=False):
     GlobalCounters.reset()

test/test_schedule.py

@@ -43,6 +43,7 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
   return sched
 
 def expect_rangeify_fails(fxn): return (unittest.expectedFailure if RANGEIFY else (lambda f:f))(fxn)
+def expect_nonrangeify_fails(fxn): return (unittest.expectedFailure if not RANGEIFY else (lambda f:f))(fxn)
 
 def _realize_weights(m):
   for p in nn.state.get_parameters(m): p.realize()
@@ -358,7 +359,7 @@ class TestSchedule(unittest.TestCase):
     out1 = r1 + y
     schedule = check_schedule([out0, out1], 2 if RANGEIFY else 4)
     reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
-    assert len(reduceops) == 2
+    assert len(reduceops) == (3 if RANGEIFY else 2)
 
   def test_div_collapse_buffer(self):
     a = Tensor.full((4,), 4.0).contiguous().realize()
@@ -697,7 +698,6 @@ class TestSchedule(unittest.TestCase):
     c = (a.sum(2).contiguous() + b).contiguous()
     check_schedule(c, 2)
 
-  @expect_rangeify_fails
   def test_kernelize(self):
     a = Tensor.empty(10)
     b = Tensor.empty(10)
@@ -705,20 +705,20 @@ class TestSchedule(unittest.TestCase):
     d = c+2
     check_schedule(d, 2)
 
-  @expect_rangeify_fails
   def test_kernelize_view(self):
     a = Tensor.empty(4,1)
     b = a*2
     c = b.kernelize()+Tensor.empty(4,4)
     check_schedule(c, 2)
 
-  @expect_rangeify_fails
   def test_kernelize_diamond(self):
     a = Tensor([0]).realize()
     prev_a = (a+1).contiguous()
     a.assign(Tensor([2]))
     a.kernelize(prev_a)
-    assert prev_a.uop in a.uop.src, "contiguous usage must run before assign"
+    # RANGEIFY doesn't apply the post diamond graph, it's fine since we can always apply the fixup on each kernelize call
+    if not RANGEIFY:
+      assert prev_a.uop in a.uop.src, "contiguous usage must run before assign"
     self.assertEqual((prev_a+a*3).item(), 1+2*3)
 
   @expect_rangeify_fails
@@ -733,7 +733,6 @@ class TestSchedule(unittest.TestCase):
     self.assertEqual(b.buffer.numpy(), [12])
 
   # unlike schedule, kernelize can be called multiple times on a Tensor
-  @expect_rangeify_fails
   def test_double_kerenlize(self):
     a = Tensor.empty(10)
     b = Tensor.empty(10)
@@ -742,7 +741,6 @@ class TestSchedule(unittest.TestCase):
     e = c.kernelize()+d.kernelize()
     check_schedule(e, 3)
 
-  @expect_rangeify_fails
   def test_kernelize_bw(self):
     a = Tensor.full((3,), 2.0, requires_grad=True).contiguous()
     b = Tensor.full((3,), 3.0, requires_grad=True).contiguous()
@@ -753,7 +751,6 @@ class TestSchedule(unittest.TestCase):
     self.assertEqual(z.item(), 18.0)
     self.assertEqual(z.grad.item(), 1.0)
 
-  @expect_rangeify_fails
   def test_kernelize_bw_view(self):
     a = Tensor.full((3,1), 2.0, requires_grad=True).contiguous()
     b = Tensor.full((3,1), 3.0, requires_grad=True).contiguous()
@@ -1661,7 +1658,6 @@ class TestSchedule(unittest.TestCase):
     check_schedule(constv, 1)
 
   @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
-  @expect_rangeify_fails
   def test_image_matmul(self):
     with Context(IMAGE=2):
       x = Tensor.randn((9, 9)).realize()
@@ -1903,6 +1899,18 @@ class TestSchedule(unittest.TestCase):
       # NOTE: this is a bug on non rangeify
       np.testing.assert_equal(tst.numpy(), a.numpy())
 
+  def test_setitem_sched(self, transpose=False):
+    a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize()
+    a2 = a.T if transpose else a
+    expected = (a+a2).tolist()
+    a.assign(a+a2)
+    kcount = len(sched:=a.schedule())
+    run_schedule(sched)
+    self.assertListEqual(a.tolist(), expected)
+    self.assertEqual(kcount, 2 if transpose else 1)
+  @unittest.skipUnless(RANGEIFY>0, "this asserts on non rangeify")
+  def test_setitem_permuted_sched(self): self.test_setitem_sched(transpose=True)
+
   def test_sparse_categorical_crossentropy_simple(self):
     X = Tensor([[0, 2, 3], [1, 2, 3]]).realize()
     Y = Tensor([1, 2]).realize()
@@ -2268,7 +2276,6 @@ class TestCopyFolding(unittest.TestCase):
     b.realize()
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
-  @expect_rangeify_fails
   def test_permute_on_disk(self):
     with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer())
     a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
@@ -2276,6 +2283,14 @@ class TestCopyFolding(unittest.TestCase):
     b.realize()
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
+  @expect_nonrangeify_fails
+  def test_permute_on_disk_contiguous(self):
+    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer())
+    a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
+    b = a.reshape(2, 2).permute(1, 0).contiguous().to("CPU")
+    b.realize()
+    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
+
   def test_permute_after_shrink(self):
     a = Tensor.arange(5)
     b = a.shrink(((0, 4),)).reshape(2, 2).permute(1, 0).to("CPU")
@@ -2284,7 +2299,7 @@ class TestCopyFolding(unittest.TestCase):
 
   # NOTE: disk permute must come after COPY
   # TODO: this is wrong because of the permute
-  @unittest.expectedFailure
+  @expect_nonrangeify_fails
   def test_permute_after_shrink_on_disk(self):
     with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).realize().uop.base.buffer.as_buffer())
     a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
@@ -2427,6 +2442,7 @@ class TestUOpBecome(unittest.TestCase):
   # sometimes we prefer to perform an op before movement ops, in this case we should stack the mops on top of the new buffer
 
   # NOTE: this expand is not reordered because there's before it to fuse
+  @expect_rangeify_fails
   def test_reorder_expand(self):
     a = Tensor.empty(4, 1)
     b = a.expand(4, 4).reciprocal()

test/test_softmax_fusion.py

@@ -30,7 +30,7 @@ def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Te
 def run_one_schedule_item(out): lower_schedule_item(get_single_element(out.schedule())).run()
 
 class TestFuse(unittest.TestCase):
-  def _test_fuse(self, fxn, *args, atol=1e-7, allow_multiple=False, **kwargs):
+  def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs):
     GlobalCounters.reset()
     out_single = fxn(*args, **kwargs).fuse()
     if not allow_multiple: run_one_schedule_item(out_single)
@@ -44,6 +44,7 @@ class TestFuse(unittest.TestCase):
     a = Tensor.rand(50,50).realize()
     self._test_fuse(lambda a: a / a.mean(axis=1), a)
 
+  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
   def test_fuse_argmax(self):
     a = Tensor.rand(50,50).realize()
     self._test_fuse(lambda a: a.argmax(axis=-1), a)

test/test_uop_graph.py

@@ -581,12 +581,20 @@ class TestUOpGraph(unittest.TestCase):
       glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
       gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 42),), "gidx0")
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<8),)).cast(dtypes.index)
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<32)),)).cast(dtypes.index)
+      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<32)),))
       to_uops_list([ld1])
 
       ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<64)),))
       with self.assertRaises(RuntimeError): to_uops_list([ld1])
 
+  def test_bounds_with_loaded_bool(self):
+    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
+    glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(8), (), 0)
+    gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
+    ld0 = glbl0.index(gidx0).load()
+    ld1 = glbl1.index(gidx0.valid(ld0)).load()
+    with self.assertRaises(RuntimeError): to_uops_list([ld1])
+
   def test_fold_gated_load(self):
     glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)

test/test_uops.py

@@ -544,86 +544,6 @@ class TestUopsObject(unittest.TestCase):
     with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)]
     assert len(ret) == 10000
 
-
-class TestShapeSpec(unittest.TestCase):
-  # ** CONST is CONST(VIEW(DEVICE)) -> RESHPAE -> EXPAND
-
-  def test_expanded_const(self):
-    a = Tensor(1).uop
-    self.assertEqual(a.st, ShapeTracker.from_shape(()))
-    a = Tensor.ones((4, 4)).uop
-    self.assertEqual(a.st, ShapeTracker.from_shape(()).reshape((1,1)).expand((4,4)))
-
-  # NOTE: CONST ShapeTracker comes from its source
-  def test_scalar_const(self):
-    a = Tensor(0).uop
-    self.assertEqual(a.st, ShapeTracker.from_shape(()))
-
-  def test_scalar_var(self):
-    vv = UOp.variable("a", 1, 4).bind(2)
-    t = Tensor(vv).uop
-    self.assertEqual(t.st, ShapeTracker.from_shape(()))
-
-  # ** ASSIGN is ASSIGN(VIEW(BUFFER), new_val)
-
-  def test_assign_flat(self):
-    buffer = Tensor.arange(4).realize()
-    a = buffer.assign(Tensor.zeros((4,), dtype=dtypes.int))
-    assign_pattern = UPat(Ops.ASSIGN, src=(UPat(Ops.BUFFER), UPat()))
-    assert assign_pattern.match(a.uop, {})
-    a.realize()
-    self.assertEqual(buffer.tolist(), [0, 0, 0, 0])
-
-  def test_assign_permuted(self):
-    buffer = Tensor.arange(4).reshape(2, 1, 2).contiguous().realize()
-    a = buffer.permute((1, 2, 0)).assign(Tensor.arange(4).reshape(1, 2, 2).contiguous())
-    a.realize()
-    self.assertEqual(buffer.tolist(), [[[0, 2]], [[1, 3]]])
-
-  def test_assign_reshaped(self):
-    buffer = Tensor.ones((4,)).contiguous().realize()
-    a = buffer.reshape((2, 2)).assign(Tensor.zeros((2, 2)))
-    assign_pattern = UPat(Ops.ASSIGN, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER))), UPat()))
-    assert assign_pattern.match(a.uop, {})
-    a.realize()
-    self.assertEqual(buffer.tolist(), [0, 0, 0, 0])
-
-  # setitem is a partial assign
-  def test_setitem(self):
-    a = Tensor.ones((4,)).contiguous().realize()
-    assign = a.shrink(((1, 2),)).assign(Tensor.zeros((1,)))
-    # the ASSIGN UOp has size=1
-    self.assertEqual(assign.uop.size, 1)
-    # the ASSIGN views the buffer with a shrunk st
-    self.assertEqual(assign.uop.src[0].st, ShapeTracker.from_shape((4,)).shrink(((1, 2),)))
-    # the underlying BUFFER has a size=4
-    self.assertEqual(assign.uop.buf_uop.size, 4)
-    # NOTE: output shape is different from the BUFFER shape
-    self.assertNotEqual(assign.uop.shape, a.uop.shape)
-    assign.realize()
-    self.assertEqual(a.tolist(), [1, 0, 1, 1])
-
-  def test_buffer_st(self):
-    a = UOp.new_buffer(Device.DEFAULT, 10, dtypes.float)
-    self.assertEqual(a.st, ShapeTracker.from_shape((10,)))
-
-  def test_ops_st(self):
-    # view / mop
-    a = Tensor.empty(4, 2, 1).permute((1, 2, 0)).uop
-    self.assertEqual(a.st, ShapeTracker.from_shape((4, 2, 1)).permute((1, 2, 0)))
-    # alu / reduce
-    alu = a*2
-    self.assertEqual(alu.st, ShapeTracker.from_shape((2, 1, 4)))
-    r = Tensor.empty(4, 4).sum(axis=1)
-    self.assertEqual(r.uop.st, ShapeTracker.from_shape((4,)))
-
-  def test_st_wmma_none(self):
-    A = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('a', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 1)))
-    B = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('b', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 2)))
-    C = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('c', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 3)))
-    wmma = UOp(Ops.WMMA, dtypes.float.vec(16), (A, B, C))
-    assert wmma.st is None
-
 class TestUOpChildren(unittest.TestCase):
   def test_children_exist(self):
     a = UOp.variable("weird_name_234", 0, 10)

test/test_uops_stats.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.helpers import getenv, GlobalCounters, EMULATE
+from tinygrad.helpers import getenv, GlobalCounters, EMULATE, RANGEIFY
 from tinygrad.engine.realize import lower_schedule_item, ProgramSpec, get_program
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
@@ -51,7 +51,11 @@ class TestMemoryCount(unittest.TestCase):
     a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
     b = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
     _, mem = get_stats(a+b)
-    self.assertEqual(mem, 1024*1024 + 2*1024)  # 2 lil reads + 1 write
+    if RANGEIFY:
+      # rangeify is smart!
+      self.assertEqual(mem, 1024 + 2*1024)  # 2 lil reads + 1 lil write
+    else:
+      self.assertEqual(mem, 1024*1024 + 2*1024)  # 2 lil reads + 1 write
 
   def test_self_add(self):
     a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)

test/unit/test_attention.py

@@ -1,5 +1,6 @@
 import unittest
 from tinygrad import Tensor, dtypes, TinyJit, UOp
+from tinygrad.helpers import RANGEIFY
 from tinygrad.apps.llm import apply_rope
 
 # TODO: test_scheduler, but just in uint
@@ -12,7 +13,7 @@ class TestAttention(unittest.TestCase):
     attn = q.scaled_dot_product_attention(k, v)
     sched = attn.schedule()
     # attention has 5 kernels now
-    self.assertEqual(len(sched), 5)
+    self.assertEqual(len(sched), 4 if RANGEIFY else 5)
     softmax_inputs = sched[1:4]
     for si in softmax_inputs:
       assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=}"
@@ -42,4 +43,4 @@ class TestAttention(unittest.TestCase):
     self.assertEqual(prune_size, 1)
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

test/unit/test_dtype_spec.py

@@ -593,14 +593,6 @@ class TestAutoCastType(unittest.TestCase):
 
     dtypes.default_float = old_default_float
 
-  @unittest.skipIf(CI, "TODO: broken RuntimeError: Attempting to relocate against an undefined symbol 'fmaxf'")
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  def test_backward_sum_acc_dtype(self):
-    # test acc of sum in the backward is upcasted to float
-    t = Tensor([5, -5], dtype=dtypes.half, requires_grad=True)
-    t.reshape(2, 1).expand(2, 10001).max().backward()
-    np.testing.assert_allclose(t.grad.numpy(), [1, 0])
-
   @unittest.skipIf(Device.DEFAULT == "PYTHON", "very slow")
   @unittest.skipIf(CI and Device.DEFAULT == "AMD", "very slow")
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Binding size is larger than the maximum storage buffer binding size")
@@ -611,6 +603,7 @@ class TestAutoCastType(unittest.TestCase):
     t = Tensor([[x]], dtype=dtypes.half, requires_grad=True).expand(N, N).contiguous()
     np.testing.assert_allclose(t.mean(axis=1).numpy(), np.array([x] * N, dtype=np.float16), rtol=1e-3)
 
+  @unittest.skip("this test only works with SPLIT_REDUCEOP=1")
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
   def test_mean_half_precision_overflow(self):
     N = 256

test/unit/test_kernelize.py

@@ -1,6 +1,7 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.uop import Ops
+from tinygrad.helpers import RANGEIFY
 
 class TestKernelize(unittest.TestCase):
   def test_add_reshaped(self):
@@ -17,7 +18,11 @@ class TestKernelize(unittest.TestCase):
     a1 = a.sum(axis=1)
     a0 = a1.sum(axis=0)
     a0.kernelize()
-    self.assertIs(a1.uop.base.op, Ops.ASSIGN)
+    self.assertEqual(len([s for s in a0.uop.toposort() if s.op is Ops.KERNEL]), 2 if RANGEIFY else 3)
+    self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS if RANGEIFY else Ops.ASSIGN)
+    # input Tensor and user contiguous kernelize
+    self.assertIs(a0.uop.base.op, Ops.ASSIGN)
+    self.assertIs(a.uop.base.op, Ops.ASSIGN)
 
   def test_two_reduce_w_add(self):
     a = Tensor.ones(16,16).contiguous()

test/unit/test_linalg.py

@@ -1,29 +1,26 @@
-import numpy as np
-import unittest
+import unittest, functools
 from tinygrad import Tensor
-from typing import List
-import functools
+import numpy as np
 
-def orthogonality_helper(A:Tensor,tolerance=1.0e-5):
+def orthogonality_helper(A:Tensor, tolerance=1e-5):
   b_shape,m = A.shape[0:-2],A.shape[-2]  #outer dimension should be the dim along orthogonality
-  A_identity = (Tensor.eye(m).reshape((1,) * len(b_shape)+(m,m)).expand(b_shape+(m,m)))
+  A_identity = (Tensor.eye(m).reshape((1,)*len(b_shape)+(m,m)).expand(b_shape+(m,m)))
   np.testing.assert_allclose((A @ A.transpose(-2,-1)).numpy(),A_identity.numpy(),atol=tolerance,rtol=tolerance)
 
-def reconstruction_helper(A:List[Tensor],B:Tensor, tolerance=1.0e-5):
+def reconstruction_helper(A:list[Tensor],B:Tensor, tolerance=1e-5):
   reconstructed_tensor = functools.reduce(Tensor.matmul, A)
   np.testing.assert_allclose(reconstructed_tensor.numpy(),B.numpy(),atol=tolerance,rtol=tolerance)
 
 class TestLinAlg(unittest.TestCase):
-
   def test_svd_general(self):
     sizes = [(2,2),(5,3),(3,5),(3,4,4),(2,2,2,2,3)]
     for size in sizes:
       a = Tensor.randn(size).realize()
-      U,S,V = Tensor.svd(a)
+      U,S,V = a.svd()
       b_shape,m,n = size[0:-2],size[-2],size[-1]
       k = min(m,n)
       s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
-      s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
+      s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([None]*len(b_shape) + [(0,m-k), (0,n-k)]))
       orthogonality_helper(U)
       orthogonality_helper(V)
       reconstruction_helper([U,s_diag,V],a)
@@ -32,7 +29,7 @@ class TestLinAlg(unittest.TestCase):
     sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
     for size in sizes:
       a = Tensor.randn(size).realize()
-      U,S,V = Tensor.svd(a,full_matrices=False)
+      U,S,V = a.svd(full_matrices=False)
       b_shape,m,n = size[0:-2],size[-2],size[-1]
       k = min(m,n)
       s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
@@ -45,20 +42,20 @@ class TestLinAlg(unittest.TestCase):
   def test_svd_large(self):
     size = (1024,1024)
     a = Tensor.randn(size).realize()
-    U,S,V = Tensor.svd(a)
+    U,S,V = a.svd()
     b_shape,m,n = size[0:-2],size[-2],size[-1]
     k = min(m,n)
     s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
-    s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
-    orthogonality_helper(U,tolerance=1.0e-3)
-    orthogonality_helper(V,tolerance=1.0e-3)
-    reconstruction_helper([U,s_diag,V],a,tolerance=1.0e-3)
+    s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([None]*len(b_shape) + [(0,m-k), (0,n-k)]))
+    orthogonality_helper(U,tolerance=1e-3)
+    orthogonality_helper(V,tolerance=1e-3)
+    reconstruction_helper([U,s_diag,V],a,tolerance=1e-3)
 
   def test_qr_general(self):
     sizes = [(3,3),(3,6),(6,3),(2,2,2,2,2)]
     for size in sizes:
       a = Tensor.randn(size).realize()
-      Q,R = Tensor.qr(a)
+      Q,R = a.qr()
       orthogonality_helper(Q)
       reconstruction_helper([Q,R],a)
 
@@ -68,9 +65,9 @@ class TestLinAlg(unittest.TestCase):
     for coefs in coefficients:
       for size in sizes:
         a = Tensor.randn(size)
-        b = Tensor.newton_schulz(a, steps=20, params=coefs, eps=0.0)
+        b = a.newton_schulz(steps=20, params=coefs, eps=0.0)
         # ns(A) = U @ Vt -> (U @ Vt) @ (U @ Vt)t = I
-        orthogonality_helper(b if size[-1] > size[-2] else b.transpose(-2, -1), tolerance=1e-1)
+        orthogonality_helper(b if size[-1] > size[-2] else b.transpose(-2, -1), tolerance=1e-3)
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_shapetracker.py

@@ -814,20 +814,6 @@ class TestShapeTrackerSize(unittest.TestCase):
     st = ShapeTracker.from_shape((10,10)).pad(((2,4), (3,1))).flip((True, True))
     self.assertEqual(st.real_size(), 100)
 
-class TestRender(unittest.TestCase):
-  def test_render(self):
-    st = ShapeTracker.from_shape((2, 3))
-    valid_idx = st.to_valid_uop()
-    idx, valid = valid_idx.get_idx(), valid_idx.get_valid()
-    self.assertEqual(idx.render(), "((ridx0*3)+ridx1)")
-    self.assertEqual(valid.render(), "True")
-
-    st = st.pad(((0, 1), (0, 0)))
-    valid_idx = st.to_valid_uop()
-    idx, valid = valid_idx.get_idx(), valid_idx.get_valid()
-    self.assertEqual(idx.render(), "((ridx0*3)+ridx1)")
-    self.assertEqual(valid.render(), "(ridx0<2)")
-
 class TestVariableShrink(unittest.TestCase):
   def test_shrink(self):
     st = ShapeTracker.from_shape((10,))

test/unit/test_simplify_valid_idx.py

@@ -71,8 +71,8 @@ class TestValidIdxSimplification(unittest.TestCase):
     idx = ridx0+ridx1+ridx2+ridx3
     load = get_gated_load_uop(valid, idx)
     self.check(load,
-      "(((ridx0+ridx1)+ridx2)+ridx3)",
-      "((((ridx0*3)+ridx1)<8)&((((ridx2*3)+ridx3)%4)<2))")
+      "(((r0+r1)+r2)+r3)",
+      "((((r0*3)+r1)<8)&((((r2*3)+r3)%4)<2))")
 
   def test_simplify_within_valid2(self):
     gidx0 = Special("gidx0", 56)
@@ -85,8 +85,8 @@ class TestValidIdxSimplification(unittest.TestCase):
     ridx0 = Range(0, 2)
     v0 = ridx0<1
     v1 = ((ridx0*5+1)%6)<5
-    self.assertEqual(simplify_valid(v0&v1).render(), "(ridx0<1)")
-    self.assertEqual(simplify_valid(v1&v0).render(), "(ridx0<1)")
+    self.assertEqual(simplify_valid(v0&v1).render(), "(r0<1)")
+    self.assertEqual(simplify_valid(v1&v0).render(), "(r0<1)")
 
   def test_valid_order_matters2(self):
     gidx0 = Special("gidx0", 13)
@@ -128,8 +128,8 @@ class TestValidIdxSimplification(unittest.TestCase):
     valid = ((((((ridx2*2)+(ridx3*3))+3)%4)<2)!=True) # noqa: E712
     load = get_gated_load_uop(valid, idx)
     self.check(load,
-      "(((ridx0*2)+(ridx3*-1))+1)",
-      "(ridx2<1)")
+      "(((r0*2)+(r3*-1))+1)",
+      "(r2<1)")
 
   def test_load_in_valid(self):
     # from FUSE_ARANGE=1 python test/test_ops.py TestOps.test_scatter_add
@@ -154,8 +154,8 @@ class TestValidIdxSimplification(unittest.TestCase):
     valid = (ridx2<1)&(ridx1<6)
     load = get_gated_load_uop(valid, idx)
     self.check(load,
-      "(ridx0*1568)",
-      "((ridx2<1)&(ridx1<6))")
+      "(r0*1568)",
+      "((r2<1)&(r1<6))")
 
   def test_valid_becomes_const1_z3(self):
     from z3 import Ints, Solver, And, If, Not, unsat
@@ -195,7 +195,7 @@ class TestValidIdxSimplification(unittest.TestCase):
     load = get_gated_load_uop(valid, idx)
     self.check(load,
       "1",
-      "((((ridx0+ridx1)<1)!=True)&(((ridx2+ridx3)<1)!=True))")
+      "((((r0+r1)<1)!=True)&(((r2+r3)<1)!=True))")
 
   def test_valid_with_non_const_rhs(self):
     ridx0 = Range(0, 2**16)
@@ -205,8 +205,8 @@ class TestValidIdxSimplification(unittest.TestCase):
     idx = ridx0%1024
     load = get_gated_load_uop(valid, idx)
     self.check(load,
-      "ridx0",
-      "(ridx0<((ridx1*4)+ridx2))")
+      "r0",
+      "(r0<((r1*4)+r2))")
 
 class TestImageSimplification(unittest.TestCase):
   def check(self, load, svalid, sidx0, sidx1):
@@ -304,7 +304,7 @@ class TestImageSimplification(unittest.TestCase):
     idx = ((alu4+1530)%1536, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))
 
     load = get_load_image_uop(shape, valid, idx)
-    self.check(load, None, "((((idx1*48)+(ridx2*6))+ridx0)+-6)", "(((idx2*2)+ridx1)+-1)")
+    self.check(load, None, "((((idx1*48)+(r2*6))+r0)+-6)", "(((idx2*2)+r1)+-1)")
 
   def test_openpilot_conv2(self):
     # conv in test/external/external_test_valid_remove.py
@@ -325,7 +325,7 @@ class TestImageSimplification(unittest.TestCase):
     idx = ((alu3+765)%768, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))
     load = get_load_image_uop(shape, valid, idx)
 
-    self.check(load, None, "((((idx1*24)+(ridx2*3))+ridx0)+-3)", "(((idx2*2)+ridx1)+-1)")
+    self.check(load, None, "((((idx1*24)+(r2*3))+r0)+-3)", "(((idx2*2)+r1)+-1)")
 
   def test_openpilot_conv3(self):
     # in openpilot 0.9.7
@@ -346,9 +346,9 @@ class TestImageSimplification(unittest.TestCase):
     load = get_load_image_uop(shape, valid, idx)
 
     self.check(load,
-               "((((idx2*2)+ridx0)<11)&((((idx1*8)+ridx1)<3)!=True))",
-               "(((idx0+((idx1*512)+(ridx1*64)))+832)%1024)",
-               "((((idx2*2)+ridx0)+(((idx1+((ridx1+5)//8))+1)//2))+-4)")
+               "((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))",
+               "(((idx0+((idx1*512)+(r1*64)))+832)%1024)",
+               "((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)")
 
   def test_simplify1(self):
     # idx has the form (A % m, A // m + k) and valid has (c0 < A) and (A < c1)
@@ -424,7 +424,7 @@ class TestImageSimplification(unittest.TestCase):
     alu1 = ((idx2*1536)+(ridx4*768)+ridx3+(idx1*24)+(ridx5*3)+-771)//768
     valid = (((idx2+ridx4)<1)!=1)&(((idx1+ridx5)<1)!=1)
     load = get_load_image_uop((128, 768, 4), valid, (alu0, alu1))
-    self.check(load, None, "((((idx1*24)+ridx3)+(ridx5*3))+-3)", "(((idx2*2)+ridx4)+-1)")
+    self.check(load, None, "((((idx1*24)+r3)+(r5*3))+-3)", "(((idx2*2)+r4)+-1)")
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_tensor_uop_representation.py

@@ -32,8 +32,7 @@ class TestTensorMutates(unittest.TestCase):
     d.realize()
     is_pattern_uop(d.uop.base, realized_pattern)
     is_pattern_uop(c.uop.base, realized_pattern)
-    # NOTE: we keep movement ops on top of the buffer view
-    is_pattern_uop(c.uop, UPat(Ops.BUFFER))
+    is_pattern_uop(c.uop.base, realized_pattern)
     assert d.uop is not d.uop.base
 
   def test_reshape_is_same_child(self):
@@ -56,7 +55,8 @@ class TestTensorUopRepresentation(unittest.TestCase):
     b = Tensor([4.,5,6]).realize()
     c = a+b
     print(c.uop)
-    is_pattern(c, UPat(Ops.ADD, src=(realized_pattern, realized_pattern)))
+    is_pattern(c, UPat(Ops.ADD))
+    for s in c.uop.src: is_pattern_uop(s.base, realized_pattern)
 
   def test_empty_buf(self):
     a = Tensor.empty(3, 3)

test/unit/test_uop_spec.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import unittest
 
 from tinygrad import Tensor
-from tinygrad.helpers import DEBUG
+from tinygrad.helpers import DEBUG, RANGEIFY
 from tinygrad.uop.ops import UOp, Ops, print_uops
 from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -75,6 +75,7 @@ class TestUOpSpec(unittest.TestCase):
     st = UOp.store(buf.view(ShapeTracker.from_shape(())), a.cast(dtypes.float))
     helper_test_verify_ast(st)
 
+  @unittest.skipIf(RANGEIFY, "RANGEIFY does not push views")
   def test_assert_masked_view_in_const(self):
     t = Tensor(6).uop
     a = t.replace(src=(t.src[0].replace(arg=t.st.reshape((1,)).pad(((0, 1),))),))

test/unit/test_uop_symbolic.py

@@ -246,7 +246,7 @@ class TestSymbolic(unittest.TestCase):
   def test_range_mod_its_symbolic_bound(self):
     a = Variable("a", 1, 10, dtypes.index)
     ridx = UOp.range(a+2, 0)
-    self.helper_test_variable(ridx%(a+2), 0, 11, "ridx0")
+    self.helper_test_variable(ridx%(a+2), 0, 11, "r0")
 
   def test_div_min_max(self):
     self.helper_test_variable(Variable("a", 2, 7) // 2, 1, 3, "(a//2)")

test/unit/test_winograd.py

@@ -42,7 +42,7 @@ class TestWinograd(unittest.TestCase):
     out = Tensor.conv2d(x,w, padding=1)
     out.mean().backward()
     backward_schedule = Tensor.schedule(x.grad, w.grad)
-    self.assertEqual(len(backward_schedule), 6 if RANGEIFY else 9)
+    self.assertEqual(len(backward_schedule), 3 if RANGEIFY else 9)
 
   def test_counters(self):
     IC, OC, X, Y = 4,4,9,9

tinygrad/apps/llm.py

@@ -118,7 +118,7 @@ class TransformerBlock:
     return h + self.ffn_down(gated)
 
   def __call__(self, x: Tensor, start_pos: int|UOp):
-    return self._feed_forward(self._attention(x, start_pos))
+    return self._feed_forward(self._attention(x, start_pos)).contiguous()
 
 class Transformer:
   def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, max_context):
@@ -156,6 +156,8 @@ class Transformer:
                         n_heads=kv[f'{arch}.attention.head_count'], n_kv_heads=kv[f'{arch}.attention.head_count_kv'],
                         norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'], vocab_size=len(kv['tokenizer.ggml.tokens']), max_context=max_context)
     nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
+    # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
+    for s in nn.state.get_parameters(model): s.replace(s.contiguous())
     return model, kv
 
   def generate(self, tokens:list[int], start_pos=0):

tinygrad/codegen/__init__.py

@@ -95,7 +95,7 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
   extra_matcher = opts.extra_matcher if opts.extra_matcher is not None else PatternMatcher([])
 
   # lower the index dtype to a concrete int
-  ret.append(RewriteStep(load_store_indexing+pm_lower_index_dtype, lambda _: opts.device, name="lower all index dtypes"))
+  ret.append(RewriteStep(pm_lower_index_dtype+load_store_indexing, lambda _: opts.device, name="lower all index dtypes"))
 
   # optional pre matcher
   if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher"))

tinygrad/codegen/late/devectorizer.py

@@ -53,6 +53,8 @@ def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp,
 load_store_indexing = PatternMatcher([
   # image load valid idx simplification
   (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
+  # simplify away long after index has been lowered
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x", dtypes.long), UPat.var("c", dtypes.bool))), lambda buf,x,c: simplify_valid_load(buf, x, c)),
   # drop true gate
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x"), UPat.const(dtypes.bool, True)),), lambda buf,x: buf.index(x)),
   # delete_redundant_gates (after expand)
@@ -97,7 +99,7 @@ def expand_index(buf:UOp, vec:UOp):
       global_offset += len(grp)
   assert None not in idxs, f"some idxs are missing {idxs}"
   # this base thing is for image, we want the CAT to be a normal pointer
-  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(vec.dtype.count), tuple(ret))
+  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(global_offset), tuple(ret))
   return post_cat.gep(tuple(cast(list[int], idxs)))
 
 def cat_after_store(cat:UOp, data:UOp, sto:UOp):
@@ -126,7 +128,7 @@ load_store_folding = PatternMatcher([
   (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), allow_any_len=True, name="sto"), gep_on_store),
   # put PTRCAT after LOAD
   (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
-   lambda cat,ld: UOp(Ops.CAT, ld.dtype, tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
+   lambda cat,ld: UOp(Ops.CAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
   # put PTRCAT after STORE
   (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), allow_any_len=True, name="sto"), cat_after_store),
 ])

tinygrad/codegen/opt/postrange.py

@@ -71,6 +71,13 @@ class Scheduler:
                         or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
     for ls in local_store_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
 
+    # filter any not in reduces
+    # TODO: enable this
+    """
+    reduce_rngs = [x.ranges for x in self.ast.toposort() if x.op is Ops.REDUCE]
+    for ls in reduce_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
+    """
+
     return [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE and x.arg[1] == AxisType.LOOP] if store_rngs else []
 
   def convert_loop_to_global(self):

tinygrad/codegen/simplify.py

@@ -17,20 +17,24 @@ pm_flatten_range = PatternMatcher([
 
 def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
+  reduce_ranges = [x.ranges for x in u.sparents if x.op is Ops.REDUCE]
   i = range_start[u.op]
   while i < len(u.src)-1:
     r0, r1 = u.src[i], u.src[i+1]
     # check same type
     if r0.arg[-1] == r1.arg[-1]:
-      s0, s1 = r0.src[0], r1.src[0]
-      # do the merge
-      new_range = r0.replace(src=(s0*s1,))
-      nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
-                           name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
-      # check if it simplifies
-      if count_divmod(nidx) <= count_divmod(u):
-        u = nidx
-        continue
+      # check if the ranges to merge are in the same reduces
+      if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges):
+        s0, s1 = r0.src[0], r1.src[0]
+        # do the merge
+        new_range = r0.replace(src=(s0*s1,))
+        nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+                             name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
+
+        # check if it simplifies
+        if count_divmod(nidx) <= count_divmod(u):
+          u = nidx
+          continue
     i += 1
   return u
 

tinygrad/device.py

@@ -327,8 +327,8 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
     if device == "METAL": return not CI
     if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX")
     if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"}
-    return device in {"AMD", "PYTHON"}
-  if dtype in dtypes.fp8s: return device == "PYTHON"
+    return device in {"AMD", "PYTHON", "NULL"}
+  if dtype in dtypes.fp8s: 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/helpers.py

@@ -147,6 +147,7 @@ EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, (os.cpu_count() or 1) // (4 if ARCH_X86 else 2))) # take 1/2 of the cores, accounting HT
 CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1)
 VIZ = PROFILE = ContextVar("VIZ", 0)
+SPEC = ContextVar("SPEC", 0)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/renderer/llvmir.py

@@ -107,7 +107,8 @@ base_rewrite = PatternMatcher([
    f"  {ctx[x]} = phi {ldt(x.dtype)} [ 0, %loop_entry_{x.arg[0]} ], [ {ctx[x]}phi, %loop_latch_{x.arg[0]} ]"),
   (UPat(Ops.ENDRANGE, name="x"), lambda ctx,x:
    f"  br label %loop_latch_{x.src[0].arg[0]}\nloop_latch_{x.src[0].arg[0]}:\n"
-   f"  {ctx[x.src[0]]}phi = add i32 {ctx[x.src[0]]}, 1\n  {ctx[x]} = icmp ult i32 {ctx[x.src[0]]}phi, {ctx[x.src[0].src[0]]}\n"
+   f"  {ctx[x.src[0]]}phi = add {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, 1\n"
+   f"  {ctx[x]} = icmp ult {ldt(x.src[0].dtype)} {ctx[x.src[0]]}phi, {ctx[x.src[0].src[0]]}\n"
    f"  br i1 {ctx[x]}, label %loop_body_{x.src[0].arg[0]}, label %loop_exit_{x.src[0].arg[0]}\nloop_exit_{x.src[0].arg[0]}:"),
 
   # if

tinygrad/schedule/multi.py

@@ -1,7 +1,7 @@
 from typing import cast, TypeVar
 import functools, itertools, operator
 from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, getenv, unwrap
-from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, resolve
+from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, resolve, track_rewrites, graph_rewrite_map
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.device import Device
 
@@ -239,3 +239,6 @@ multi_pm = PatternMatcher([
   (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE),
         src=(UPat(Ops.MULTI, name="multi"), ), name="root"), passthrough_multi),
 ])+replace_allreduce
+
+@track_rewrites()
+def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]: return graph_rewrite_map(big_sink, multi_pm, name="multi_pm")

tinygrad/schedule/rangeify.py

@@ -1,19 +1,21 @@
 from typing import Any, cast
 import functools, operator
 from dataclasses import dataclass, field
-from collections import defaultdict
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify, graph_rewrite_map
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify
 from tinygrad.uop.symbolic import sym, symbolic_simple
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup, unwrap
-from tinygrad.schedule.multi import multi_pm
-
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup
 from tinygrad.schedule.kernelize import Kernel
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType
+from tinygrad.codegen.simplify import pm_flatten_range
 
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
+                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
+                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL}
+
 double_reshape = PatternMatcher([
   # RESHAPE on RESHAPE is the second reshape
   (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE),), name="x"),
@@ -34,24 +36,34 @@ earliest_rewrites = double_reshape+PatternMatcher([
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)),
    lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None),
 
+  # remove contiguous on movement ops before a copy on disk
+  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, allow_any_len=True, name="copy"),
+   lambda x,copy: copy.replace(src=(x,)+copy.src[1:]) if isinstance(x.device, str) and x.device.startswith("DISK") else None),
+  # push copy past movement ops to disk
+  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, allow_any_len=True, name="copy"),
+   lambda x,copy: x.replace(src=(copy.replace(src=(x.src[0],)+copy.src[1:], tag=None),)+x.src[1:], tag=copy.tag) \
+      if isinstance(x.device, str) and x.device.startswith("DISK") else None),
+
   # COPY and source size need to match
   # TODO: expand after copy creates issues with tagging
   (UPat(Ops.COPY, src=(UPat(GroupOp.Movement, name="r"), UPat(name="d")), name="c"),
    lambda c,r,d: c.replace(src=(r.contiguous(), d)) if r.size != r.base.size else None),
 
-  # assign only to buffer
+  # make inputs to mstack contiguous
+  (UPat(Ops.MSTACK, name="ms"), lambda ms: ms.replace(src=tuple(s if s.op in ALWAYS_CONTIGUOUS else s.contiguous() for s in ms.src))),
+
+  # assign only to buffer, otherwise make it a CONTIGUOUS
   (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}, name="target"), UPat(name="x")), name="assign"),
-   lambda x,target,assign: x.f(Ops.NOOP, tag=assign.tag) if target.base.op is not Ops.BUFFER else None),
+   lambda x,target,assign: x.f(Ops.CONTIGUOUS, tag=assign.tag) if ((t:=target.base).op is not Ops.BUFFER and \
+       not (t.op is Ops.MSTACK and all(s.op is Ops.BUFFER for s in t.src))) else None),
+
+  # realize before assign if input permutes the target buffer
+  (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), lambda a,b,assign: assign.replace(src=(a, b.contiguous())) \
+      if any(x.base is a.base and x is not a for x in b.toposort(gate=lambda x:x.op not in ALWAYS_CONTIGUOUS)) else None),
 
   # copy only to different device
   (UPat(Ops.COPY, src=(UPat.var("x"), UPat()), name="copy"), lambda x,copy: x.f(Ops.NOOP, tag=copy.tag) if x.device == copy.device else None),
 
-  # handle disk
-  # TODO: this doesn't need to use st.views
-  (UPat.var("x").f((Ops.BITCAST, Ops.CONTIGUOUS), name="t"),
-   lambda x,t: UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (t.size, x.st.views[0].offset), tag=t.tag).reshape(t.shape) if isinstance(x.device, str) \
-       and x.device.startswith("DISK") else None),
-
   # contiguous/buffer/copy/assign is already contiguous
   #(UPat(Ops.CONTIGUOUS, name="root", src=(UPat((Ops.CONTIGUOUS, Ops.BUFFER, Ops.COPY, Ops.ASSIGN)),)), lambda root: root.src[0]),
 ])
@@ -59,10 +71,6 @@ earliest_rewrites = double_reshape+PatternMatcher([
 # *****************
 # 1. add realize where we have to
 
-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
-                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
-                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD}
-
 def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
 
 def realize_parents(ctx:dict[UOp, None], rb:UOp) -> None:
@@ -83,26 +91,19 @@ do_realize = PatternMatcher([
   (UPat(Ops.ASSIGN, name="a"), realize_assign),
 ])
 
-
 class WrappedContig:
   def __init__(self, x): self.x = x
   def __repr__(self): return f"C({self.x})"
-add_contiguous = PatternMatcher([
-  (UPat(GroupOp.All, name="x"),
-   lambda ctx,x: x.replace(tag=WrappedContig(x.tag)).realize() if x in ctx and not isinstance(x.tag, WrappedContig) else None),
-])
+add_contiguous = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda ctx,x: x.replace(tag=WrappedContig(x.tag)).realize() if x in ctx else None),])
 remove_contig_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=x.tag.x) if isinstance(x.tag, WrappedContig) else None)])
 
 # *****************
 # 2. mark all children
 
 @dataclass
-class ChildrenContext:
-  children: dict[UOp, list[UOp]]|None = None
-  realize_roots: defaultdict[UOp, list[UOp]] = field(default_factory=lambda: defaultdict(list))
+class ChildrenContext: children: dict[UOp, list[UOp]]|None = None
 def extract_children(ctx:ChildrenContext, x:UOp):
   if ctx.children is not None: return
-
   children_map = x.get_children_map()
   ctx.children = {}
   for k,v in children_map.items():
@@ -111,16 +112,6 @@ def extract_children(ctx:ChildrenContext, x:UOp):
     # NOTE: this gate shouldn't be here
     if any(x.op is Ops.REDUCE_AXIS for x in k.toposort()) and any(x.op in {Ops.BUFFER, Ops.CONTIGUOUS} for x in k.toposort()):
       ctx.children[k] = non_sink_children
-  # if a node is in the toposort of multiple realizes, it will be indexed by different indices and we can bufferize early
-  # this prevents index_child: "children not making progress" error on big graphs
-  for r in [u for u in x.toposort() if u.op is Ops.REALIZE and (RANGEIFY<2 or u.arg is None)]:  # ignore partial realizes
-    for u in r.toposort(gate=lambda x: x is not Ops.REALIZE or (RANGEIFY>1 and u.arg is not None)):
-      ctx.realize_roots[u].append(r)
-
-def bufferize_early(ctx:ChildrenContext, x:UOp):
-  # this will also change the sources such that mark_children wont add the children/child uops anymore
-  new_srcs = [s.realize() if s in unwrap(ctx.children) and len(ctx.realize_roots[s])>1 else s for s in x.src]
-  return x.replace(src=tuple(new_srcs))
 
 def mark_children(ctx:ChildrenContext, x:UOp):
   assert ctx.children is not None
@@ -130,8 +121,7 @@ def mark_children(ctx:ChildrenContext, x:UOp):
 
 pm_children = PatternMatcher([
   (UPat(Ops.SINK, name="x"), extract_children),
-  (UPat(GroupOp.All-{Ops.CHILDREN, Ops.SINK, Ops.REALIZE}, name="x"), bufferize_early),
-  (UPat(GroupOp.All-{Ops.CHILDREN, Ops.SINK, Ops.REALIZE}, name="x"), mark_children),
+  (UPat(GroupOp.All-{Ops.CHILD, Ops.CHILDREN, Ops.SINK}, name="x"), mark_children),
 ])
 
 # *****************
@@ -191,7 +181,10 @@ def map_expand(r:UOp, idx:UOp):
     else:
       ending_ranges.extend(axis_to_range)
       new_rngs.append(a.const_like(0))
-  ending_ranges = [x.arg for x in ending_ranges if x not in non_ending_ranges]
+  # if RANGEIFY >= 2, we are aggressive about not ending ranges
+  if RANGEIFY >= 2: ending_ranges = [x.arg for x in ending_ranges if x not in non_ending_ranges]
+  # if RANGEIFY=1, if it's ending at all we end it
+  else: ending_ranges = [x.arg for x in ending_ranges]
   if idx.arg is not None: ending_ranges.append(idx.arg)
   return r.src[0].index(*new_rngs, arg=min(ending_ranges) if ending_ranges else None)
 
@@ -275,11 +268,16 @@ def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
     end_ranges = []
     idx_ranges = []
     # NOTE: locals aren't working, so we only fully bufferize here (unless RANGEIFY > 1)
-    all_all_same = all(all_same(r) for r in all_rngs)
-    for i,valid_rngs in enumerate(all_rngs):
+    rngs_valids = []
+    for valid_rngs in all_rngs:
       rngs, valids = zip(*[(r.get_idx(), r.get_valid()) for r in valid_rngs])
+      # if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
+      same_rngs = [x if x.op is not Ops.RANGE or resolve(x.src[0] != 1) else UOp.const(dtypes.index, 0) for x in rngs]
+      rngs_valids.append((rngs, valids, all_same(same_rngs)))
+    all_all_same = all(same_rngs for _,_,same_rngs in rngs_valids)
+    for i,(rngs,valids,same_rngs) in enumerate(rngs_valids):
       # we compare the ranges without their valids
-      if all_same(rngs) and (all_all_same or RANGEIFY > 1):
+      if same_rngs and (all_all_same or RANGEIFY > 1):
         # the new valid is the OR of all the children valids
         minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False))
         out_rngs.append(minimum_valid.where(rngs[0], UOp.invalid()).simplify())
@@ -315,7 +313,8 @@ def might_end_axis(idx:UOp):
   if all(x.op not in {Ops.REDUCE_AXIS} for x in idx.toposort()): return None
   to_end_axis = []
   for i,a in enumerate(idx.src[1:]):
-    if any(x.arg > idx.arg for x in a.toposort() if x.op is Ops.RANGE):
+    # in RANGEIFY=1, always realize
+    if not (RANGEIFY > 1) or any(x.arg > idx.arg for x in a.toposort() if x.op is Ops.RANGE):
       to_end_axis.append(i)
   if to_end_axis: return idx.replace(src=(idx.src[0].realize(arg=tuple(to_end_axis)),)+idx.src[1:], arg=None)
   return idx.replace(arg=None)
@@ -346,7 +345,8 @@ pm_rangeify = pm_mops+PatternMatcher([
 
   # handle assign
   (UPat(Ops.INDEX, src=(UPat(Ops.ASSIGN, name="assign"),), allow_any_len=True, name="x"),
-    lambda x,assign: assign.replace(src=tuple([s.index(*x.src[1:]) for s in assign.src])+(assign.src[0],))),
+    lambda x,assign: assign.replace(src=tuple([s.index(*x.src[1:]) for s in assign.src])+(assign.src[0],)) \
+        if assign.src[1].op is not Ops.KERNEL else None),
 
   # move MAP through elementwise ALU / reduce. these are the items with cost
   (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union(
@@ -361,12 +361,19 @@ pm_rangeify = pm_mops+PatternMatcher([
 # *****************
 # 3.5 cleanups
 
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN}
+
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
+  # don't optimize ALWAYS_RUN_OPS
+  if b.src[0].op in ALWAYS_RUN_OPS: return None
+
   new_rng = []
   hit = False
   reshape: list[sint] = []
   for s,rng in zip(b.shape, b.src[1:]):
+    # skip for symbolic. TODO: fix this
+    if rng.op is Ops.RANGE and rng.src[0].op is not Ops.CONST: return None
     if rng not in b.src[0].sparents and rng.op is Ops.RANGE:
       reshape.append(1)
       hit = True
@@ -374,7 +381,8 @@ def cleanup_dead_axes(b:UOp):
       reshape.append(s)
       new_rng.append(rng)
   if hit:
-    return b.replace(src=b.src[0:1]+tuple(new_rng)).reshape(tuple(reshape)).expand(b.shape)
+    # move the tag to the expand
+    return b.replace(src=b.src[0:1]+tuple(new_rng), tag=None).reshape(tuple(reshape)).expand(b.shape).replace(tag=b.tag)
 
 # if a buffer is being stored just for permutes or something, remove it
 # we want to reexpress the indexes of idx2 in terms of the implied b1
@@ -384,17 +392,32 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   assert all(x.op is Ops.RANGE for x in buf.src[1:])
 
   # if it's user contiguous, we never remove it
-  if src.op is Ops.CONTIGUOUS: return None
+  if src.op in ALWAYS_RUN_OPS: return None
 
-  # here is where we compute the cost
-  # for now just no REDUCE, COPY, or ASSIGN
-  ran = src.toposort(gate=lambda x: x.op not in {Ops.INDEX})
   # we don't want to bufferize threefry, also causes problems because not all platforms support long
-  if any(x.op in {Ops.REDUCE, Ops.COPY, Ops.BUFFER_VIEW, Ops.ASSIGN} for x in ran) and src.op is not Ops.THREEFRY: return None
+  if src.op is not Ops.THREEFRY:
+    # *** here is where we compute the cost ***
+    # if we return None, the bufferize is kept
 
-  # simple, matching old behavior
-  #if src.op is not Ops.INDEX: return None
+    accessed_buffers = []
+    def red_gate(x):
+      if x.op is Ops.INDEX:
+        accessed_buffers.append(x)
+        return False
+      return True
+    ran = src.toposort(gate=red_gate)
 
+    # if this is generated from multiple buffers, don't remove this buffer
+    if len(dedup([x.src[0] for x in accessed_buffers])) > 2: return None
+
+    # const reduce is okay
+    # TODO: move the reduce folder to before this to prevent the need for this
+    def okay_reduce(x:UOp): return all(y.op not in {Ops.BUFFER, Ops.COPY} for y in x.sparents)
+
+    # always run this list of ops
+    if any(x.op is Ops.REDUCE and not okay_reduce(x) for x in ran): return None
+
+  # if it makes it here, the bufferize is removed
   # this is the ranges replaced
   return src.substitute(dict(zip(buf.src[1:], idx.src[1:])))
 
@@ -402,8 +425,8 @@ def pre_bufferize(b:UOp, x:UOp, copy:UOp):
   nb = b.replace(src=(b.src[0].contiguous(),)+b.src[1:])
   return copy.replace(src=(x.replace(src=(nb,)+x.src[1:]), copy.src[1]))
 
-pm_cleanups = double_reshape+pm_mops+PatternMatcher([
-  #(UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes),
+pm_cleanups = pm_mops+PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes),
   (UPat(GroupOp.All-{Ops.BUFFERIZE, Ops.BUFFER}, name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
   (UPat((Ops.BUFFERIZE), name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType)
     and (resolve(prod(x.dtype.shape)!=prod(x.shape)) or x.shape[-1]%4!=0) else None),
@@ -416,12 +439,33 @@ pm_cleanups = double_reshape+pm_mops+PatternMatcher([
   (UPat.var("src").f(Ops.BUFFERIZE, allow_any_len=True, name="buf").f(Ops.INDEX, allow_any_len=True, name="idx"), remove_bufferize),
   # no buffers for const
   (UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.arg).rtag(b.tag)),
-  # if any CONST with DEVICE make it here (symbolic/copy issue), remove it
-  #(UPat(Ops.DEVICE).f(Ops.CONST, name="c"), lambda c: c.replace(src=())),
   # copy on CONST is CONST
   (UPat(Ops.COPY, src=(UPat.cvar("x"), UPat()), name="copy"), lambda copy,x: copy.const_like(x.arg)),
   (UPat(Ops.COPY, src=(UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.COPY}).f(Ops.BUFFERIZE, allow_any_len=True, name="b")
                        .f(Ops.INDEX, allow_any_len=True, name="x"), UPat()), name="copy"), pre_bufferize),
+  # mstack on CONST is CONST
+  (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True),
+   lambda s: UOp.const(c.dtype, c.arg) if (c:=s.base).op is Ops.CONST else None),
+])
+
+def late_buffer_view(t:UOp, b:UOp):
+  if isinstance(b.device, str) and b.device.startswith("DISK"):
+    rngs = b.src[1:]
+    size = prod(shape := [int(r.vmax+1) for r in rngs])
+
+    # walk up for the INDEX
+    x = t
+    while not any(u.op is Ops.INDEX for u in x.src): x = x.src[0]
+    x = next(u for u in x.src if u.op is Ops.INDEX)
+
+    if len(shape) == 0: offset = x.src[1].arg
+    else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0)
+
+    return b.replace(src=(UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (size, offset), tag=t.tag),) + b.src[1:])
+  return b
+to_bufferview = PatternMatcher([
+  (UPat((Ops.BITCAST, Ops.CONTIGUOUS), name="t").f(Ops.BUFFERIZE, allow_any_len=True, name="b"), late_buffer_view),
+  (UPat((Ops.BITCAST, Ops.CONTIGUOUS)).f(Ops.BUFFER_VIEW, name="b"), lambda b: b.replace(src=b.src[0].src)),
 ])
 
 # *****************
@@ -471,7 +515,7 @@ def bufferize_to_store(x:UOp):
   # TODO: how is this unified?
   return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype)
 
-pm_add_buffers = pm_mops+PatternMatcher([
+pm_add_buffers = pm_mops+to_bufferview+PatternMatcher([
   (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
 
   # move RESHAPEs through MSELECT/MSTACK
@@ -488,6 +532,7 @@ class LocalAddBufferContext:
   map:dict = field(default_factory=dict)
   vars:dict = field(default_factory=dict)
   range:int = 0
+  parent_tags:list = field(default_factory=list)
 
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
   ret = UOp(Ops.DEFINE_GLOBAL, buf.dtype.ptr(buf.arg), arg=ctx.dg)
@@ -522,6 +567,9 @@ to_define_global = PatternMatcher([
   # this is only needed if you are using symbolic
   (UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
 
+  # remove RANGE with 0 size
+  (UPat(Ops.RANGE, name="r"), lambda r: UOp.const(dtypes.index, 0) if r.vmax == 0 else None),
+
   # renumber the ranges starting with 0 so that kernel deduping works
   (UPat(Ops.RANGE, name="r"), renumber_range),
 ])
@@ -548,20 +596,30 @@ rangeify_codegen = PatternMatcher([
    lambda src, barrier, gate: src.load(UOp(Ops.IF, src=(gate, barrier)))),
 ])
 
+def remove_metadata_tags(ctx:LocalAddBufferContext, x:UOp):
+  if x.tag is None or x.tag == (): return None
+  ctx.parent_tags += list(x.tag)
+  return x.replace(tag=None)
+
+pm_remove_tags = PatternMatcher([
+  # remove all the tags
+  (UPat(GroupOp.All, name="x"), remove_metadata_tags),
+])
+
 def split_store(ctx:list[UOp], x:UOp):
   if len(x.ranges): return None
   if x.src[0].ptrdtype.addrspace is AddrSpace.LOCAL: return None
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()
-  ret = graph_rewrite(x, to_define_global+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True)
+  ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen+pm_remove_tags, ctx=lctx, name="kernel split", bottom_up=True)
 
   # gather the metadata
-  metadatas = [ctx[y].metadata for x in ret.sparents if x.tag is not None for y in x.tag]
+  metadatas = [ctx[y].metadata for y in lctx.parent_tags]
 
   # NOTE: the hack for COPY is here
   ret = ret.sink() if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW} else ret.src[1]
-  kernel_arg = Kernel(ret,tuple(dedup(flatten([x for x in metadatas if x is not None]))))
+  kernel_arg = Kernel(ret,tuple(dedup(flatten([x for x in metadatas if x is not None])))[::-1])
   kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg)
   return x.as_buf().assign(kernel)
 
@@ -578,16 +636,28 @@ add_tags = PatternMatcher([
   (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND}.union(GroupOp.Movement), name="x"), tag_uop),
 ])
 
+# support for using a contiguous permuted view instead of the parent view if one exists
+# modified from kernelize.py to not use ShapeTracker
+
+def found_contiguous(ctx:dict[UOp, UOp], contig:UOp, src:UOp):
+  x = src
+  while x is not src.base:
+    if x.op is Ops.PERMUTE: contig = contig.permute(argsort(x.arg))
+    elif x.op is Ops.RESHAPE: contig = contig.reshape(x.src[0].shape)
+    else: return None
+    x = x.src[0]
+  ctx[src.base] = contig
+replace_contiguous = PatternMatcher([
+  (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement, name="src"),), name="contig"), found_contiguous),
+  (UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
+])
+
 @track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len([u for u in UOp.sink(*ret.values()).toposort() if u.op is Ops.KERNEL]))}", True)
 def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   uop_list: list[UOp] = []
   tsink = graph_rewrite(sink, add_tags, ctx=uop_list, bottom_up=True, name="number the uops")
 
-  # HACKS: handle multi with graph_rewrite_map in order to not have to add all the tag logic to multi
-  msink = graph_rewrite_map(tsink, multi_pm, name="multi")
-  tsink = msink[tsink].substitute({v:v.rtag(k.tag) for k,v in msink.items() if v.tag is None and k.tag is not None})
-
-  tsink = graph_rewrite(tsink, earliest_rewrites, name="earliest rewrites")
+  tsink = graph_rewrite(tsink, earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
   realize_map: dict[UOp, UOp] = {}
   graph_rewrite(tsink, do_realize, ctx=realize_map, name="Input Graph")
   # NOTE: we don't use contiguous here, contiguous is a user op

tinygrad/tensor.py

@@ -8,7 +8,7 @@ from tinygrad.dtype import _from_np_dtype, _to_np_dtype
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup
 from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, RANGEIFY, FUSE_ATTENTION
 from tinygrad.gradient import compute_gradient
-from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, MathTrait, identity_element, all_metadata, index_to_concrete_int, sint_to_uop, \
+from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, MathTrait, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, \
   srender
 from tinygrad.uop.spec import tensor_uop_spec, type_verify
 from tinygrad.device import Device, Buffer
@@ -16,6 +16,7 @@ from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.memory import memory_planner
 from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
 from tinygrad.schedule.rangeify import get_rangeify_map
+from tinygrad.schedule.multi import get_multi_map
 from tinygrad.schedule.kernelize import get_kernelize_map
 
 # *** all in scope Tensors are here. this gets relevant UOps ***
@@ -142,7 +143,7 @@ class Tensor(MathTrait):
     if isinstance(data, UOp):
       assert dtype is None or dtype==data.dtype, "dtype doesn't match, and casting isn't supported"
       # if data is dtype.index that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
-      if data.dtype==dtypes.index: data = index_to_concrete_int(data)
+      if data.dtype==dtypes.index: data = _index_to_concrete_int(data)
       if data.op is Ops.BIND:  # type: ignore  # mypy type narrowing is bugged here
         var, val = data.unbind()  # type: ignore
         # give the bound constant a device
@@ -241,6 +242,10 @@ class Tensor(MathTrait):
     # verify Tensors match the spec
     if __debug__: type_verify(list(big_sink.toposort()), tensor_uop_spec)
 
+    if RANGEIFY and any(isinstance(x._device, tuple) for x in big_sink.toposort()):
+      _apply_map_to_tensors(get_multi_map(big_sink), "Apply Multi Map")
+      big_sink = UOp.sink(*flatten([x.uop.src if x.uop.op is Ops.MULTI else [x.uop] for x in (self,)+lst]))
+
     becomes_map = get_rangeify_map(big_sink) if RANGEIFY else get_kernelize_map(big_sink)
     _apply_map_to_tensors(becomes_map, name="Apply Kernelize Map")
     return self
@@ -535,7 +540,7 @@ class Tensor(MathTrait):
         device=device, dtype=dtypes.uint32, requires_grad=False)
       Tensor._device_rng_counters[device] = Tensor([num], device=device, dtype=dtypes.uint32, requires_grad=False)
     # increment rng counter for devices
-    else: Tensor._device_rng_counters[device].assign(Tensor._device_rng_counters[device] + num).contiguous()
+    else: Tensor._device_rng_counters[device].assign(Tensor._device_rng_counters[device] + num)
 
     # threefry random bits
     bits_count = Tensor._device_rng_counters[device] - num
@@ -4085,24 +4090,24 @@ class Tensor(MathTrait):
     """
     assert self.ndim > 1, "NS only works for two or more dims"
     G = self / (self.square().sum(axis=(-2, -1), keepdim=True).sqrt() + eps)
-    G = G.transpose(-2, -1) if self.shape[-2] > self.shape[-1] else G
+    if (swap := self.shape[-2] > self.shape[-1]): G = G.transpose(-2, -1)
     for _ in range(steps): G = sum(p * functools.reduce(lambda x, y: (y @ y.transpose(-2, -1)) @ x, [G]*i, G) for i,p in enumerate(params))
-    return G.transpose(-2, -1) if self.shape[-2] > self.shape[-1] else G
+    return G.transpose(-2, -1) if swap else G
 
   def qr(self) -> tuple[Tensor, Tensor]:
     assert self.ndim > 1, f"expected two or more dimensions, got {self.ndim}"
+    b_shape, m, n = self.shape[:-2], int(self.shape[-2]), int(self.shape[-1])
     R = self.clone()
-    b_shape, m, n = self.shape[0:self.ndim - 2], int(R.shape[-2]), int(R.shape[-1])
-    Q = Tensor.eye(m, dtype = self.dtype).reshape((1,) * (len(self.shape) - 2) + 2 * (m,)).expand(b_shape + 2 * (m,)).contiguous()
-    for i in range(int(min(m, n))):
-      x = R[..., i:m, i]
+    Q = Tensor.eye(m, dtype=self.dtype).reshape((1,) * len(b_shape) + (m, m)).expand(b_shape + (m, m)).contiguous()
+    for i in range(min(m, n)):
+      x = R[..., i:m, i].contiguous()  # TODO: without contigous this can silently be wrong, should at least assert
       s = -x[..., 0].sign()
       u1 = x[..., 0] - s * x.square().sum(-1).sqrt()
-      w = x.unsqueeze(-1) / u1.reshape(b_shape + 2 * (1,))
+      w = x.unsqueeze(-1) / u1.reshape(b_shape + (1, 1))
       w[..., 0, 0] = 1
-      tau = (-s * u1 / x.square().sum(-1).sqrt()).reshape(b_shape + 2 * (1,)).expand(w.shape)
+      tau = (-s * u1 / x.square().sum(-1).sqrt()).reshape(b_shape + (1, 1))
       R[..., i:m, :] = R[..., i:m, :] - (w * tau) @ (w.transpose(-2, -1) @ R[..., i:m, :])
-      Q[..., :, i:m] = Q[..., :, i:m] - (Q[..., :, i:m] @ w) @ (tau.transpose(-2, -1) * w.transpose(-2, -1))
+      Q[..., :, i:m] = Q[..., :, i:m] - (Q[..., :, i:m] @ w) @ (tau * w).transpose(-2, -1)
     return Q,R
 
   def svd(self, full_matrices = True) -> tuple[Tensor, Tensor, Tensor]:
@@ -4110,14 +4115,14 @@ class Tensor(MathTrait):
     assert self.ndim > 1, f"expected two or more dimensions, got {self.ndim}"
     b_shape, m, n = self.shape[:-2], int(self.shape[-2]), int(self.shape[-1])
     #preprocess the matrix
-    Q, R = (Tensor.qr(self) if m >= n else Tensor.qr(self.transpose(-2, -1)))
-    num, q_num = int(min(m, n)), int(max(m, n))
-    U = R.shrink(tuple([(0, self.shape[i]) for i in range(self.ndim - 2)] + [(0, num), (0, num)])).contiguous()
-    V = Tensor.eye(num, dtype = self.dtype).reshape((1,) * (self.ndim - 2) + (num, num)).expand(b_shape + 2 * (num,)).contiguous()
+    Q, R = (self.qr() if m >= n else self.transpose(-2, -1).qr())
+    num, q_num = min(m, n), max(m, n)
+    U = R.shrink(tuple([None] * len(b_shape) + [(0, num), (0, num)])).contiguous()
+    V = Tensor.eye(num, dtype=self.dtype).reshape((1,) * len(b_shape) + (num, num)).expand(b_shape + (num, num)).contiguous()
     #prepare round robin pairing
-    permute, inverse_permute = Tensor.arange(0, num, dtype = dtypes.int), Tensor.zeros(num, dtype = dtypes.int).contiguous()
+    permute, inverse_permute = Tensor.arange(0, num, dtype=dtypes.int), Tensor.zeros(num, dtype=dtypes.int).contiguous()
     permute[num//2:num] = permute[num//2:num].flip(0)
-    inverse_permute[permute] = Tensor.arange(num, dtype = dtypes.int)
+    inverse_permute[permute] = Tensor.arange(num, dtype=dtypes.int)
     def one_round_jacobi(U, V,permute,inverse_permute):
       #pair all the columns
       V_permuted, runoff_V = (V[..., permute].split(num - 1, -1)) if num % 2 == 1 else (V[..., permute], None)
@@ -4141,15 +4146,15 @@ class Tensor(MathTrait):
       else: permute = permute[0].reshape(1).cat(((permute[1:num] - 2) % (num - 1)) + 1)
       inverse_permute = inverse_permute.scatter(0,permute,Tensor.arange(num,dtype=dtypes.int32))
       return U, V, permute, inverse_permute
-    max_iterations, iterations_per_round = 1, int((num) * math.log2(num) * 2 + 2)#sorta heuristic, most use num*log2(num)
+    max_iterations, iterations_per_round = 1, int(num * math.log2(num) * 2 + 2)#sorta heuristic, most use num*log2(num)
     for _ in range(max_iterations * iterations_per_round): U, V, permute, inverse_permute = one_round_jacobi(U, V, permute, inverse_permute)
     #extract singular values and sort. construct U from Q
     S, indices = U.square().sum(-2).sqrt().sort(dim = -1, descending=True)
-    new_indices = Tensor.arange(num).reshape((1,) * (self.ndim - 1) + (num,)).expand(b_shape + 2 * (num,)).contiguous()
-    new_indices[..., :num] = indices.reshape(b_shape + (1,) + (num,)).expand(b_shape + 2 * (num,))
-    U,V = U.gather(-1, new_indices[...,0:num,0:num]) / S.unsqueeze(-2), V.gather(-1, new_indices[..., 0:num, 0:num]).realize()
+    new_indices = Tensor.arange(num).reshape((1,) * (self.ndim - 1) + (num,)).expand(b_shape + (num, num)).contiguous()
+    new_indices[..., :num] = indices.reshape(b_shape + (1, num)).expand(b_shape + (num, num))
+    U, V = U.gather(-1, new_indices[...,0:num,0:num]) / S.unsqueeze(-2), V.gather(-1, new_indices[..., 0:num, 0:num]).realize()
 
-    padded_u = Tensor.eye(q_num, dtype = U.dtype).reshape((1,) * (self.ndim - 2) + 2 * (q_num,)).expand(b_shape + 2 * (q_num,)).contiguous()
+    padded_u = Tensor.eye(q_num, dtype=U.dtype).reshape((1,) * len(b_shape) + (q_num, q_num)).expand(b_shape + (q_num, q_num)).contiguous()
     padded_u[..., 0:num, 0:num] = U
     U = Q @ padded_u
     if not full_matrices: U, V = U[..., 0:num], V[..., 0:num]

tinygrad/uop/ops.py

@@ -7,7 +7,8 @@ from tinygrad.uop import Ops, GroupOp
 from tinygrad.uop.mathtraits import MathTrait
 from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
-from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, RANGEIFY, VIZ
+from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, RANGEIFY, VIZ, SPEC
+from tinygrad.helpers import strip_parens
 if TYPE_CHECKING:
   from tinygrad.shape.shapetracker import ShapeTracker
   from tinygrad.device import Buffer, MultiBuffer
@@ -66,6 +67,10 @@ class UOpMetaClass(type):
     if _buffer is not None:
       assert op is Ops.BUFFER, f"trying to set Buffer {_buffer} for {op}"
       buffers[created] = _buffer
+    if SPEC:
+      from tinygrad.uop.spec import full_spec
+      with Context(IGNORE_OOB=1): ret = full_spec.rewrite(created)
+      if cast(bool|None, ret) is not True: raise RuntimeError(f"SPEC ISSUE {ret}: {created}")
     return created
 
 # some uops map to other stuff
@@ -150,7 +155,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def st(self) -> ShapeTracker|None:
-    if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG,
+    if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.MSTACK,
                                                    Ops.BUFFER, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.STORE}:
       return None
     if self.op is Ops.BARRIER: return None
@@ -212,8 +217,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # determine what ranges this is in
   @functools.cached_property
-  def ranges(self) -> dict[UOp, None]:
-    if self.op is Ops.RANGE: return {self:None}
+  def _ranges(self) -> dict[UOp, None]:
     ret: dict[UOp, None] = {}
     if self.op in range_start.keys():
       for s in self.src[:range_start[self.op]]: ret.update(s.ranges)
@@ -223,6 +227,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       for s in self.src: ret.update(s.ranges)
     return ret
 
+  @property
+  def ranges(self) -> dict[UOp, None]:
+    if self.op is Ops.RANGE: return {self:None}
+    return self._ranges
+
   # *** uop evaluation ***
 
   def simplify(self, tracked=False):
@@ -465,7 +474,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.op is Ops.MSTACK: return UOp(Ops.MSTACK, self.dtype, src=tuple(x.as_buf() for x in self.src))
     # TODO: this should be the only one of these. this is the one RANGEIFY uses
     s = self
-    while len(s.src) and s.op is not Ops.BUFFER: s = s.src[0]
+    while len(s.src) and s.op not in {Ops.BUFFER, Ops.MSTACK}: s = s.src[0]
     return s
 
   @property
@@ -624,7 +633,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return fxn(**{k:v for k,v in var_vals.items() if k in varnames})
 
   def render(self, simplify=True, pm:PatternMatcher|None=None) -> str:
-    with Context(TRACK_MATCH_STATS=0):
+    with Context(TRACK_MATCH_STATS=0, SPEC=0):
       ret = graph_rewrite(self.simplify() if simplify else self, renderer if pm is None else pm)
     return ret.arg if ret.op is Ops.NOOP else str(ret)
 
@@ -998,7 +1007,8 @@ class RewriteContext:
     return ret
 
   def unified_rewrite(self, root:UOp) -> UOp:
-    stack: list[tuple[UOp, int, UOp]] = [(root, 0, root)]
+    stack: collections.deque[tuple[UOp, int, UOp]] = collections.deque([(root, 0, root)])
+    on_stack = {root}  # all UOps either on the stack or in self.replace, i.e. dont have to be placed again
     while stack:
       if len(stack) >= 200000: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
       n, stage, new_n = stack.pop()
@@ -1016,7 +1026,10 @@ class RewriteContext:
                 seen.add(test_n)
                 new_n, test_n = test_n, self.cached_bpm_rewrite(test_n)
             stack.append((n, 1, new_n))
-            for x in reversed(new_n.src): stack.append((x, 0, x))
+            for x in reversed(new_n.src):
+              if x in on_stack: continue
+              stack.append((x, 0, x))
+              on_stack.add(x)
           # if the bpm matching raised a gate, we are done with this node and dont continue down the srcs
           except BottomUpGate: self.replace[n] = new_n
         elif stage == 1:
@@ -1039,7 +1052,7 @@ class RewriteContext:
           except KeyError: raise RewriteNotReady
       except RewriteNotReady:
         # retry this later
-        stack.insert(0, (n, stage, new_n))
+        stack.appendleft((n, stage, new_n))
     return self.replace[root]
 
 @track_matches
@@ -1064,31 +1077,29 @@ def sint_to_uop(x:sint) -> UOp: return UOp.const(dtypes.index, x) if isinstance(
 def select_dtype(u): return (dtypes.long if u.overflows(dtypes.int32) else dtypes.int).vec(u.dtype.count)
 pm_lower_index_dtype = PatternMatcher([
   # There are no Unary ops at this point in symbolic, those are introduced later
-  (UPat(GroupOp.Binary, dtypes.index, name="u", src=(UPat.var("x"), UPat.var("y"))), lambda u,x,y:
-    x.cast(dt:=least_upper_dtype(select_dtype(u), x.dtype, y.dtype)).alu(u.op, y.cast(dt))),
-  # comparison ops might now have different dtypes in their sources
-  (UPat(GroupOp.Comparison, name="u", src=(UPat.var("x",dtypes.ints), UPat.var("y", dtypes.ints))), lambda u,x,y:
-    x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)).alu(u.op, y.cast(dt)) if x.dtype!=y.dtype else None),
-  (UPat(Ops.WHERE, dtypes.index, src=(UPat(), UPat.var("x"), UPat(Ops.CONST, arg=Invalid)), name="u"), lambda u,x: u.replace(dtype=x.dtype)),
-  (UPat(Ops.WHERE, dtypes.index, src=(UPat.var("cond"), UPat.var("x"), UPat.var("y"))), lambda cond,x,y:
-    cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt))),
-  (UPat((Ops.CONST, Ops.VCONST), dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=select_dtype(u))),
-  (UPat((Ops.RANGE,), dtype=dtypes.index, src=(UPat.var("end")), name="r"), lambda ctx,r,end:
-    r.replace(dtype=(dt:=select_dtype(r)), src=(end.cast(dt),))),
-  (UPat(Ops.CAST, dtype=dtypes.index, src=(UPat.var("x", dtypes.ints),), name="u"), lambda u,x: x),
-  (UPat(Ops.VECTORIZE, dtype=dtypes.index, name="u"), lambda u: u.replace(
-    dtype=(dt:=least_upper_dtype(*[x.dtype for x in u.src])).vec(u.dtype.count), src=tuple(x.cast(dt) for x in u.src))),
-  (UPat(Ops.VECTORIZE, dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=(dt:=(dtypes.long if any(v.overflows(dtypes.int) for v in u.src)
-    else dtypes.long)).vec(u.dtype.count),src=tuple(x.cast(dt) for x in u.src))),
-  (UPat((Ops.SPECIAL,Ops.DEFINE_VAR), dtypes.index, name="u"), lambda u: u.replace(dtype=dtypes.int)),
-  (UPat((Ops.BIND), dtypes.index, name="u"), lambda u: u.replace(dtype=u.src[0].dtype)),
+  (UPat(GroupOp.Binary, name="u", src=(UPat.var("x").cast(dtypes.index), UPat.var("y").cast(dtypes.index))), lambda u,x,y:
+    x.cast(dt:=least_upper_dtype(select_dtype(u), x.dtype, y.dtype)).alu(u.op, y.cast(dt)).cast(u.dtype)),
+  (UPat((Ops.CONST, Ops.VCONST), dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype) if u.arg!=Invalid else None),
+  (UPat(Ops.WHERE, dtypes.index, src=(UPat.var("cond"), UPat.var("x").cast(dtypes.index), UPat.var("y").cast(dtypes.index))), lambda cond,x,y:
+    cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt)).cast(dtypes.index)),
+  (UPat(Ops.RANGE, src=(UPat.var("end").cast(dtypes.index)), name="r"), lambda r,end: r.replace(dtype=end.dtype, src=(end,)).cast(dtypes.index)),
+  (UPat(Ops.VECTORIZE, src=UPat().cast(dtypes.index), name="v"),
+    lambda v: v.replace(dtype=(dt:=select_dtype(v)), src=tuple(s.src[0].cast(dt.scalar()) for s in v.src)).cast(dtypes.index)),
+  # special can only be int32
+  (UPat(Ops.SPECIAL, src=(UPat.var("var").cast(dtypes.index),), name="u"), lambda u,var: u.replace(dtype=dtypes.int, src=(var,)).cast(dtypes.index)),
+  (UPat(Ops.DEFINE_VAR, dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=dtypes.int).cast(dtypes.index)),
+  (UPat(Ops.BIND, src=(UPat.var("var").cast(dtypes.index), UPat.cvar("val").cast(dtypes.index))), lambda var,val: var.bind(val).cast(dtypes.index)),
+  (UPat(Ops.CAST, src=(UPat(name="x").cast(dtypes.index),), name="c"), lambda x,c: x.cast(c.dtype)),
   # lower Invalid
   (UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid))), lambda buf,idx,cond: buf.index(idx, cond)),
-  # remove hanging cast
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast()),), lambda buf,idx: buf.index(idx)),
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast(), UPat.var("valid"))), lambda buf,idx,valid: buf.index(idx, valid)),
+  # remove hanging casts
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast()),), lambda buf,idx: buf.index(idx)),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("valid"))), lambda buf,idx,valid: buf.index(idx, valid)),
+  (UPat((Ops.STORE, Ops.LOAD), src=(UPat(), UPat(), UPat().cast(dtypes.index)), allow_any_len=True, name="s"),
+    lambda s: s.replace(src=s.src[:2]+tuple(u.src[0] for u in s.src[2:]))),
+  (UPat((Ops.SINK, Ops.NOOP), src=UPat().cast(dtypes.index), name="n"), lambda n: n.replace(src=tuple(s.src[0] for s in n.src))),
 ])
-def index_to_concrete_int(u:UOp): return graph_rewrite(u, pm_lower_index_dtype)
+def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]
 
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
 
@@ -1098,7 +1109,7 @@ syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<"
 renderer = PatternMatcher([
   (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: UOp(Ops.NOOP, arg=x.arg[0])),
   (UPat((Ops.SPECIAL), name="x"), lambda x: UOp(Ops.NOOP, arg=x.arg)),
-  (UPat(Ops.RANGE, name="x"), lambda x: UOp(Ops.NOOP, arg=f"ridx{x.arg[0]}" if x.arg[0] >= 0 else f"ridxm{-x.arg[0]}")),
+  (UPat(Ops.RANGE, name="x"), lambda x: UOp(Ops.NOOP, arg=f"r{x.arg[0]}" if x.arg[0] >= 0 else f"rm{-x.arg[0]}")),
   (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: UOp(Ops.NOOP, arg=str(x.arg))),
   (UPat(Ops.UNROLL, name="x"), lambda x: UOp(Ops.NOOP, arg=f"UNROLL({x.src[0].arg}, {x.arg})")),
   (UPat(Ops.CAST, name="x"), lambda x: UOp(Ops.NOOP, arg=f"({str(x.dtype)[7:]})({x.src[0].arg})")),
@@ -1111,6 +1122,8 @@ renderer = PatternMatcher([
   (UPat(Ops.WHERE, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[1].arg} if {x.src[0].arg} else {x.src[2].arg})")),
   (UPat(set(syms.keys()), src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[0].arg}{syms[x.op]}{x.src[1].arg})")),
   (UPat(Ops.VIEW, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.view({x.arg})")),
+  (UPat(Ops.INDEX, name="x"), lambda x:
+   UOp(Ops.NOOP, arg=''.join([f"[{strip_parens(y.arg)}]" for y in x.src[1:]])) if all(y.op is Ops.NOOP for y in x.src[1:]) else None),
 ])
 renderer_infer = PatternMatcher([
   (UPat(Ops.MOD, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"cmod({x.src[0].arg}, {x.src[1].arg})")),
@@ -1137,6 +1150,7 @@ pm_pyrender = PatternMatcher([
    lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.f({x.op}, dtype=dtypes.bool)")),
 ])
 
+@Context(SPEC=0)
 def pyrender(ast:UOp) -> list[str]:
   cmap = ast.get_children_map()
   to_render = set()

tinygrad/uop/spec.py

@@ -1,7 +1,7 @@
 from typing import cast, Callable
 from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite, AxisType
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid
-from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context, cpu_profile
+from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context, cpu_profile, RANGEIFY
 from tinygrad.shape.shapetracker import ShapeTracker
 try:
   import z3
@@ -29,9 +29,9 @@ try:
     (UPat(Ops.SPECIAL, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(x.arg, 0, x.src[0].arg[1]-1, ctx[0])))),
     (UPat(Ops.DEFINE_VAR, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(x.arg[0], x.arg[1], x.arg[2], ctx[0])))),
     (UPat(Ops.RANGE, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(f"ridx{x.arg}", 0, x.src[0].arg[1]-1, ctx[0])))),
-    # float loads only become a variable when they get cast to int/bool
-    (UPat(Ops.LOAD, dtypes.ints, name="x"),
-      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(f"load{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0])))),
+    # loaded bools become a z3 int with min max of 0-1
+    (UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), name="x"), lambda x,ctx:
+      UOp(Ops.NOOP, arg=(ctx[0],create_bounded(f"load{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0]))).cast(x.dtype)),
     (UPat(Ops.CONST, dtype=dtypes.ints+(dtypes.bool,dtypes.index), name="x"),
       lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],(z3.BoolVal if dtypes.is_bool(x.dtype) else z3.IntVal)(x.arg, ctx=ctx[0].ctx)))),
     # z3 can cast from bool to int automatically
@@ -49,7 +49,7 @@ try:
   ])
 
   def uops_to_z3(solver, *uops: UOp) -> 'list[z3.ExprRef]':
-    with Context(TRACK_MATCH_STATS=0):  # cant pickle z3 objects
+    with Context(TRACK_MATCH_STATS=0, SPEC=0):  # cant pickle z3 objects, and these UOps don't follow spec
       return [s.arg[1] for s in graph_rewrite(uops[0].sink(*uops[1:]), z3_renderer, ctx=(solver, {})).src]
 
   z3_imported = True
@@ -124,7 +124,8 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
 
 # ***** uop type spec *****
 
-def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)):
+def validate_index(idx:UOp, gate:UOp|None=None):
+  if gate is None: gate = UOp.const(dtypes.bool, True)
   # TODO: check for overflow
   if IGNORE_OOB or isinstance(idx.dtype, ImageDType) or (sz := idx.src[0].ptrdtype.size) == -1: return True
   # We can use UOp min/max to do a faster check, but it can give false positive since its not an exact bound and doesn't consider the mask
@@ -146,7 +147,8 @@ def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)):
       return False
   return True
 
-def validate_store(idx:UOp, val:UOp, gate:UOp=UOp.const(dtypes.bool, True)):
+def validate_store(idx:UOp, val:UOp, gate:UOp|None=None):
+  if gate is None: gate = UOp.const(dtypes.bool, True)
   if gate.op is Ops.IF: gate = gate.src[0]
   # we need to find the implicit gates, inverse of delete_redundant_gates
   for u in val.toposort():
@@ -226,7 +228,7 @@ spec = PatternMatcher([
 
   (UPat(Ops.REDUCE_AXIS, name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) >= 2 and x.arg[0] in {Ops.ADD, Ops.MUL, Ops.MAX}),
   (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
-  (UPat(Ops.VECTORIZE, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.count and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
+  (UPat(Ops.VECTORIZE, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
   (UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: x.arg is None),
   (UPat(Ops.BARRIER, dtypes.void, src=UPat(Ops.STORE, allow_any_len=True)), lambda: True), # NOTE: all pointers must be local
   (UPat(Ops.BARRIER, dtypes.void), lambda: True), # BARRIERs can also happen at the end of loops
@@ -250,6 +252,69 @@ ast_spec = PatternMatcher([
   (UPat(GroupOp.All-{Ops.SINK}, name="root"), lambda root: all_same([x.shape for x in root.src if x.st is not None])),
 ])
 
+# *** this spec should match all UOps ever created ***
+
+full_non_rangeify_spec = PatternMatcher([]) if RANGEIFY else PatternMatcher([
+  # in non rangeify const can still have a View, and sometimes a FUSE while propagating
+  (UPat((Ops.VIEW, Ops.FUSE)).f(Ops.CONST), lambda: True),
+])
+
+full_spec = PatternMatcher([
+  # 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, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.index))), lambda: True),
+
+  # all children is fine
+  (UPat(Ops.CHILDREN), lambda: True),
+  # child must have CHILDREN parent
+  (UPat(Ops.CHILD, src=(UPat(Ops.CHILDREN),)), lambda: True),
+
+  # all rewrite error are okay
+  (UPat(Ops.REWRITE_ERROR), lambda: True),
+
+  # rangeify: buffer view with index or load is okay
+  (UPat(Ops.BUFFER_VIEW, src=(UPat((Ops.INDEX, Ops.LOAD)),)), lambda: True),
+  # bufferize (must be on ranges)
+  (UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.op is Ops.RANGE for y in x.src[1:])),
+  # realize with one src is fine
+  (UPat(Ops.REALIZE, src=(UPat(),)), lambda: True),
+  # intermediate index
+  (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:]) or None),
+  (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
+  # copy on index
+  (UPat(Ops.COPY, src=(UPat(Ops.INDEX), UPat())), lambda: True),
+  # assign on index. the third op is the shape
+  (UPat(Ops.ASSIGN, src=(UPat(Ops.INDEX), UPat(), UPat(GroupOp.Movement))), lambda: True),
+
+  # expander: unroll/contract/gep/ptrcat/cat
+  (UPat((Ops.UNROLL, Ops.CONTRACT), src=(UPat(),)), lambda: True),
+  # GEP multi is supported here
+  (UPat(Ops.GEP, name="gep"), lambda gep: gep.dtype is dtypes.void or gep.dtype.vcount == len(gep.arg)),
+  # PTRCAT is like VECTORIZE, but it functions on ptrs
+  (UPat(Ops.PTRCAT, name="x"), lambda x: x.dtype.vcount == sum([y.dtype.base.count for y in x.src])),
+  # CAT is like VECTORIZE, but the srcs can be vectors
+  (UPat(Ops.CAT, name="x"), lambda x: x.dtype.vcount == sum([y.dtype.vcount for y in x.src])),
+  # vectorized index
+  (UPat(Ops.INDEX, src=(UPat((Ops.VECTORIZE, Ops.CAST)), UPat())), lambda: True),
+
+  # linearizer: outputs + intermediate KERNELs
+  (UPat((Ops.BLOCKSTART, Ops.BLOCK, Ops.BLOCKFINAL, Ops.BLOCKEND, Ops.KERNEL), dtype=dtypes.void), lambda: True),
+
+  # allow index dtype on a restricted set of UOps
+  (UPat((Ops.ADD, Ops.MUL, Ops.MOD, Ops.IDIV, Ops.MAX, Ops.WHERE,
+         Ops.SPECIAL, Ops.CAST, Ops.RANGE, Ops.VCONST, Ops.VECTORIZE), dtype=dtypes.index), lambda: True),
+
+  # all loads/stores
+  (UPat((Ops.LOAD, Ops.STORE)), lambda: True),
+  # all ifs
+  (UPat(Ops.IF), lambda: True),
+  # all DEFINE_VAR to deal with the floats used in reduce collapse
+  (UPat(Ops.DEFINE_VAR), lambda: True),
+  # reshape on STORE
+  (UPat(Ops.RESHAPE, src=(UPat(Ops.STORE),)), lambda: True),
+])+full_non_rangeify_spec+tensor_uop_spec+spec
+
 # ***** uop helpers *****
 
 def type_verify(uops:list[UOp], extra_spec:PatternMatcher|None=None):

tinygrad/uop/upat.py

@@ -153,7 +153,8 @@ def _get_code(self:UPat, has_ctx:bool):
 @functools.cache
 def upat_compile(self:UPat, fxn) -> Callable|None:
   real_fxn = types.FunctionType(*deconstruct_function(fxn))
-  code = _get_code(self, 'ctx' in inspect.signature(real_fxn).parameters)
+  # UOps used here don't follow the spec
+  with Context(SPEC=0): code = _get_code(self, 'ctx' in inspect.signature(real_fxn).parameters)
   if code is None: return None
   code_str, dyn_lookup = code
   globs = dyn_lookup.copy()

tinygrad/viz/serve.py

@@ -83,6 +83,8 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
         label += f"\n{shape_to_str(u.shape)}"
       elif len(rngs:=u.ranges):
         label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[-1]]) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})"
+      if u.op is Ops.INDEX:
+        label += f"\n{u.render()}"
     except Exception:
       label += "\n<ISSUE GETTING LABEL>"
     if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"
@@ -251,7 +253,9 @@ class Handler(BaseHTTPRequestHandler):
       except FileNotFoundError: status_code = 404
     elif (query:=parse_qs(url.query)):
       if url.path == "/disasm": ret, content_type = get_disassembly(**query), "application/json"
-      else: return self.stream_json(get_details(traces[i:=int(query["ctx"][0])][1][int(query["idx"][0])], i))
+      else:
+        try: return self.stream_json(get_details(traces[i:=int(query["ctx"][0])][1][int(query["idx"][0])], i))
+        except KeyError: status_code = 404
     elif url.path == "/ctxs": ret, content_type = json.dumps(ctxs).encode(), "application/json"
     elif url.path == "/get_profile" and profile_ret: ret, content_type = profile_ret, "application/octet-stream"
     else: status_code = 404