set testpath on pytest

* enable rangeify const folding

* renumber ranges for kernel deduping

.github/workflows/test.yml

@@ -526,6 +526,8 @@ jobs:
           -k "not test_symbolic_arange_sym_step and not test_threefry_doesnt_use_long" \
           test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \
           test/test_outerworld_range.py test/test_sample.py test/test_randomness.py
+    - name: Test multitensor
+      run: RANGEIFY=1 PYTHONPATH="." python3 test/test_multitensor.py TestMultiTensor.test_matmul_shard_1_1 TestMultiTensor.test_simple_add_W
     - name: Test GPU=1 RANGEIFY=1
       run: GPU=1 RANGEIFY=1 pytest -n auto test/test_ops.py
     - name: Test CPU=1 RANGEIFY=2

examples/mamba.py

@@ -279,9 +279,15 @@ def generate(model, tokenizer, prompt: str, n_tokens_to_gen: int = 10, temp: boo
   # Loading in the prompt tokens
   logits = model.forward(Tensor([tks]))[:, -1, :]
   for _ in tqdm(range(n_tokens_to_gen), desc="Speed Gen"):
-    # TODO: topk
     if sample:
-      tok_Tens = (logits/temp).softmax().multinomial()
+      scaled_logits = logits / temp
+      if top_k is not None:
+        topk_values, topk_indices = scaled_logits.topk(top_k)
+        filtered_logits = Tensor.full_like(scaled_logits, -float("inf"))
+        filtered_logits = filtered_logits.scatter(dim=-1, index=topk_indices, src=topk_values)
+        tok_Tens = filtered_logits.softmax().multinomial()
+      else:
+        tok_Tens = scaled_logits.softmax().multinomial()
     else:
       tok_Tens = logits.argmax(axis=-1).unsqueeze(0)
     tok = tok_Tens.item()
@@ -298,6 +304,7 @@ if __name__ == "__main__":
   parser.add_argument("--size", type=str, default="370m",
                       help=f"Size of model to use [{', '.join([k for k in MODELS.keys()])}]")
   parser.add_argument("--n_tokens", type=int, default=10, help="Number of tokens to generate")
+  parser.add_argument("--top_k", type=int, help="Limit sampling to the top k most likely tokens")
   parser.add_argument("--sample", dest="sample", action="store_true", help="Sample flag")
   parser.add_argument("--temp", type=float, default=1.0, help="Sampling temp has to be <=1.0")
   args = parser.parse_args()
@@ -308,8 +315,9 @@ if __name__ == "__main__":
   num_toks = args.n_tokens
   sample = args.sample
   temp = args.temp
+  top_k = args.top_k
   s = time.time()
-  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp)
+  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp, top_k=top_k)
   print(tinyoutput)
   print('TIME: ', time.time() - s)
   TORCHOUTPUT = "Why is gravity \nso important?\nBecause it's the only"

examples/mlperf/scripts/stable_diffusion_downloads.sh

@@ -0,0 +1,57 @@
+#!/usr/bin/env bash
+# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md
+
+# setup dirs
+
+DATA=/raid/datasets/stable_diffusion
+
+LAION=$DATA/laion-400m/webdataset-moments-filtered 
+COCO=$DATA/coco2014
+mkdir -p $LAION $COCO
+
+CKPT=/raid/weights/stable_diffusion
+mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception
+
+# download data
+
+# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below
+#RCLONE=""
+RCLONE="rclone"
+
+## VAE-encoded image latents, from 6.1M image subset of laion-400m
+## about 1 TB for whole download
+$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
+$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P
+$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P
+cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \
+  echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; }
+
+## prompts and FID statistics from 30k image subset of coco2014
+## 33 MB
+$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
+$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P
+
+$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
+$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P
+
+# download checkpoints
+
+## clip (needed for text and vision encoders for validation)
+CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin"
+CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4"
+CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json"
+wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL}
+wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL}
+echo "${CLIP_WEIGHTS_SHA256}  ${CKPT}/clip/open_clip_pytorch_model.bin"                    | sha256sum -c
+
+## sd (needed for latent->image decoder for validation, also has clip text encoder for training)
+SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt'
+SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824"
+wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL}
+echo "${SD_WEIGHTS_SHA256}  ${CKPT}/sd/512-base-ema.ckpt"                    | sha256sum -c
+
+## inception (needed for validation)
+FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
+FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff"
+wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL}
+echo "${FID_WEIGHTS_SHA1}  ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth"                    | sha1sum -c

examples/sdxl.py

@@ -437,8 +437,8 @@ if __name__ == "__main__":
     im.show()
 
   # validation!
-  if args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024 \
-    and not args.weights:
+  is_default = args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024
+  if is_default and not args.weights and not args.fakeweights:
     ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "sdxl_seed0.png")))
     distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
     assert distance < 4e-3, colored(f"validation failed with {distance=}", "red")

extra/models/inception.py

@@ -270,8 +270,10 @@ class FidInceptionV3:
     self.Mixed_7b = inception.Mixed_7b
     self.Mixed_7c = inception.Mixed_7c
 
-  def load_from_pretrained(self):
-    state_dict = torch_load(str(fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")))
+  def load_from_pretrained(self, path=None):
+    if path is None:
+      path = fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")
+    state_dict = torch_load(str(path))
     for k,v in state_dict.items():
       if k.endswith(".num_batches_tracked"):
         state_dict[k] = v.reshape(1)

extra/to_movement_ops.py

@@ -35,7 +35,7 @@ def to_movement_ops(st: ShapeTracker) -> List[Tuple[MovementOps, Tuple]]:
   to_apply:List[Tuple[MovementOps, Tuple]] = []
   for i, v in enumerate(st.views):
     real_shape = tuple(y-x for x,y in v.mask) if v.mask else v.shape
-    offset = v.offset + sum(st*(s-1) for s,st in zip(real_shape, v.strides) if st<0)
+    offset = (v.offset or 0) + sum(st*(s-1) for s,st in zip(real_shape, v.strides) if st<0)
     real_offset = offset + (sum(x*st for (x,_),st in zip(v.mask, v.strides)) if v.mask else 0)
     real_real_shape = [s for s,st in zip(real_shape, v.strides) if st]
     strides: List[int] = [abs(st) if isinstance(st,int) else st for st in v.strides if st]

extra/torch_backend/backend.py

@@ -177,22 +177,28 @@ def cached_to_movement_ops(shape, st) -> list:
 
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from extra.to_movement_ops import to_movement_ops, apply_mop, MovementOps
+
+@wrap_view_op
+def _as_strided(tensor:Tensor, size, stride, storage_offset=None):
+  # multiple as_strided do not compound
+  base = canonical_base(tensor)
+  # TODO: this is heavyweight
+  st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),))
+  ret = base
+  if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st)
+  if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size)
+  for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo)
+  return ret
+
 @torch.library.impl("aten::as_strided", "privateuseone")
 def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
   storage_offset = storage_offset or tensor.storage_offset()
-  @wrap_view_op
-  def _as_strided(tensor:Tensor, size, stride, storage_offset=None):
-    # multiple as_strided do not compound
-    base = canonical_base(tensor)
-    # TODO: this is heavyweight
-    st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),))
-    ret = base
-    if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st)
-    if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size)
-    for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo)
-    return ret
   return _as_strided(tensor, size, stride, storage_offset)
 
+@torch.library.impl("aten::_reshape_alias", "privateuseone")
+def _reshape_alias(tensor:torch.Tensor, size, stride):
+  return _as_strided(tensor, size, stride)
+
 @torch.library.impl("aten::empty_strided", "privateuseone")
 def empty_strided(size, stride, dtype, layout=None, device=None, pin_memory=False):
   if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}")

pytest.ini

@@ -3,3 +3,4 @@ norecursedirs = extra
 timeout = 180
 timeout_method = thread
 timeout_func_only = true
+testpaths = test

setup.py

@@ -9,13 +9,12 @@ with open(directory / 'README.md', encoding='utf-8') as f:
 
 testing_minimal = [
   "numpy",
-  "torch==2.7.1",
+  "torch==2.8.0",
   "pytest",
   "pytest-xdist",
   "pytest-timeout",
   "hypothesis",
   "z3-solver",
-  "ml_dtypes"
 ]
 
 setup(name='tinygrad',
@@ -60,7 +59,7 @@ setup(name='tinygrad',
         'triton': ["triton-nightly>=2.1.0.dev20231014192330"],
         'linting': [
             "pylint",
-            "mypy==1.13.0",
+            "mypy==1.18.1",
             "typing-extensions",
             "pre-commit",
             "ruff",

test/external/external_test_speed_llama.py

@@ -20,7 +20,7 @@ class TestLLaMASpeed(unittest.TestCase):
   def test_llama_compile(self):
     backup_program = Device[Device.DEFAULT].runtime
     backup_allocator = Device[Device.DEFAULT].allocator
-    backup_compiler = Device[Device.DEFAULT].compiler
+    backup_compiler = Device[Device.DEFAULT].compiler.compile_cached
     Device[Device.DEFAULT].runtime = FakeProgram
     Device[Device.DEFAULT].allocator = FakeAllocator(Device.default)
 
@@ -44,14 +44,14 @@ class TestLLaMASpeed(unittest.TestCase):
     run_llama("codegen(1)")
 
     # test no compiler use for this
-    Device[Device.DEFAULT].compiler = None
+    Device[Device.DEFAULT].compiler.compile_cached = None
     run_llama("methodcache", False)
     with Profiling(sort='time', frac=0.1, fn="/tmp/llama.prof", ts=5):
       run_llama("profile", False)
 
     Device[Device.DEFAULT].runtime = backup_program
     Device[Device.DEFAULT].allocator = backup_allocator
-    Device[Device.DEFAULT].compiler = backup_compiler
+    Device[Device.DEFAULT].compiler.compile_cached = backup_compiler
 
 if __name__ == '__main__':
   TestLLaMASpeed().test_llama_compile()

test/test_dtype.py

@@ -10,7 +10,6 @@ from tinygrad import Device, Tensor, dtypes
 from hypothesis import assume, given, settings, strategies as strat
 from test.helpers import rand_for_dtype
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
-import ml_dtypes
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
 
@@ -129,11 +128,10 @@ class TestDType(unittest.TestCase):
       np.testing.assert_allclose(tin, tor, atol=1e-6, rtol=1e-3)
 
   def test_finfo(self):
-    if self.DTYPE not in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]: return
-    info = ml_dtypes.finfo(ml_dtypes.bfloat16 if self.DTYPE is dtypes.bfloat16 else _to_np_dtype(self.DTYPE))
-    assert info.bits == self.DTYPE.itemsize*8
-    assert info.nexp == dtypes.finfo(self.DTYPE)[0]
-    assert info.nmant == dtypes.finfo(self.DTYPE)[1]
+    if self.DTYPE not in [dtypes.float16, dtypes.float32, dtypes.float64]: return
+    info = np.finfo(_to_np_dtype(self.DTYPE))
+    self.assertEqual(info.bits, self.DTYPE.itemsize*8)
+    self.assertEqual((info.nexp, info.nmant), dtypes.finfo(self.DTYPE))
 
 def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
   target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
@@ -151,7 +149,8 @@ class TestFp8s(unittest.TestCase):
 
 class TestFp8sConversions(unittest.TestCase):
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
-  def test_float_to_fp8e4m3(self, x): np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), ml_dtypes.float8_e4m3fn(x).tobytes()[0])
+  def test_float_to_fp8e4m3(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
 
   def test_float_to_fp8e4m3_extreme_values(self):
     np.testing.assert_equal(float_to_fp8(FP8E4M3_MAX, dtypes.fp8e4m3), 126)
@@ -164,7 +163,8 @@ class TestFp8sConversions(unittest.TestCase):
     np.testing.assert_equal(float_to_fp8(-math.nan, dtypes.fp8e4m3), 255)
 
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2_MAX, max_value=FP8E5M2_MAX))
-  def test_float_to_fp8e5m2(self, x): np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), ml_dtypes.float8_e5m2(x).tobytes()[0])
+  def test_float_to_fp8e5m2(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
 
   def test_float_to_fp8e5m2_extreme_values(self):
     np.testing.assert_equal(float_to_fp8(FP8E5M2_MAX, dtypes.fp8e5m2), 123)
@@ -177,10 +177,12 @@ class TestFp8sConversions(unittest.TestCase):
     np.testing.assert_equal(float_to_fp8(-math.nan, dtypes.fp8e5m2), 254)
 
   @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e4m3_to_float(self, x): np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3), np.uint8(x).view(ml_dtypes.float8_e4m3fn).item())
+  def test_fp8e4m3_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fn).float().item())
 
   @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e5m2_to_float(self, x): np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), np.uint8(x).view(ml_dtypes.float8_e5m2).item())
+  def test_fp8e5m2_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())
 
 @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16(unittest.TestCase):

test/test_kernel_cache.py

@@ -16,14 +16,14 @@ class TestKernelCache(unittest.TestCase):
 
     a1 = Tensor.rand(4,4).realize()
     b1 = Tensor.rand(4,4).realize()
-    orig_compile_func = Device['CPU'].compiler
-    Device['CPU'].compiler = None # making it not callable
+    orig_compile_func = Device['CPU'].compiler.compile_cached
+    Device['CPU'].compiler.compile_cached = None # making it not callable
 
     try:
       x1 = a1 + b1 + unique_const
       x1.realize() # Same kernel should be from cache.
     finally:
-      Device['CPU'].compiler = orig_compile_func
+      Device['CPU'].compiler.compile_cached = orig_compile_func
 
 if __name__ == "__main__":
   unittest.main()

test/test_linearizer.py

@@ -482,7 +482,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   assert s[-1].ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {s[-1]}"
   # now all input buffers in s[-1] should be realized
   # create fresh buffers for the outputs
-  bufs = [Buffer((x).device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
+  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
   return push_views(s[-1].ast), bufs
 
 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
@@ -504,7 +504,7 @@ def reset_bufs(bufs:list[Buffer]):
 
 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
                                apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
-  outbufs = [real_bufs[x.src[0].base.arg] for x in realized_ast.src]
+  outbufs = real_bufs[:len(realized_ast.src)]
   device = real_bufs[0].device
   wanna_output = [np.array(x).flatten() for x in wanna_output]
 

test/test_method_cache.py

@@ -5,9 +5,9 @@ from tinygrad.nn.state import get_state_dict
 
 class TestMethodCache(unittest.TestCase):
   def setUp(self):
-    self.backup_compiler = Device[Device.DEFAULT].compiler
+    self.backup_compiler = Device[Device.DEFAULT].compiler.compile_cached
   def tearDown(self):
-    Device[Device.DEFAULT].compiler = self.backup_compiler
+    Device[Device.DEFAULT].compiler.compile_cached = self.backup_compiler
 
   def test_simple_methodcache(self):
     a = Tensor([1])
@@ -15,19 +15,19 @@ class TestMethodCache(unittest.TestCase):
     c = Tensor([3])
     d = Tensor([4])
     (a+b).realize()
-    Device[Device.DEFAULT].compiler = None
+    Device[Device.DEFAULT].compiler.compile_cached = None
     (c+d).realize()
 
   def test_nested_methodcache(self):
     a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4])
     ((a+b)+(a+b)).realize()
-    Device[Device.DEFAULT].compiler = None
+    Device[Device.DEFAULT].compiler.compile_cached = None
     ((c+d)+(c+d)).realize()
 
   def test_nested_methodcache_swap(self):
     a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4])
     ((a+b)+(c+d)).realize()
-    Device[Device.DEFAULT].compiler = None
+    Device[Device.DEFAULT].compiler.compile_cached = None
     ((c+d)+(a+b)).realize()
 
   @unittest.skip("incorrect use of transformer")
@@ -38,7 +38,7 @@ class TestMethodCache(unittest.TestCase):
     # NOTE: you have to do this twice due to the k-v cache
     for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
     for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
-    Device[Device.DEFAULT].compiler = None
+    Device[Device.DEFAULT].compiler.compile_cached = None
     for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
 
 if __name__ == '__main__':

test/test_ops.py

@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM, RANGEIFY
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -234,7 +234,8 @@ class TestOps(unittest.TestCase):
   def test_unfold(self):
     helper_test_op([(8,)], lambda x: x.unfold(0, 2, 1))
     helper_test_op([(8,)], lambda x: x.unfold(0, 2, 2))
-    helper_test_op([(8,)], lambda x: x.unfold(0, 7, 3))
+    # TODO: something is wrong with unfold
+    if not getenv("TINY_BACKEND"): helper_test_op([(8,)], lambda x: x.unfold(0, 7, 3))
     helper_test_op([(3,3,3)], lambda x: x.unfold(2, 2, 8))
     helper_test_op([(3,3,3)], lambda x: x.unfold(1, 0, 8))
     helper_test_op([(3,3,3,3,3)], lambda x: x.unfold(-1, 2, 2))
@@ -3028,6 +3029,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,y.clip(0,1),
                                                                                                         pos_weight=torch.tensor(pos_weight)),
                                        lambda x,y: x.binary_crossentropy_logits(y.clip(0,1),pos_weight=Tensor(pos_weight)))
+
+  @unittest.skipIf(RANGEIFY > 1, "broken on RANGEIFY > 1, TODO: fix")
   def test_cross_entropy_class_probabilities(self):
     helper_test_op([(32,), (32,)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
     helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))

test/test_rangeify.py

@@ -3,6 +3,19 @@ from tinygrad import Tensor, nn
 from tinygrad.helpers import RANGEIFY, Context, GlobalCounters
 from tinygrad.uop.ops import UOp
 
+@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
+class TestRangeifyAssign(unittest.TestCase):
+  def test_assign_permuted(self):
+    A = Tensor.empty(4, 4, dtype='int')
+    B = Tensor.arange(16).reshape(4,4)
+    ret = A.permute(1,0).assign(B)
+    lst = ret.tolist()
+    lst2 = A.tolist()
+    lst3 = B.tolist()
+    print(lst)
+    print(lst2)
+    print(lst3)
+
 N = 256
 
 @unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")

test/test_schedule.py

@@ -12,9 +12,9 @@ from tinygrad import nn, dtypes, Device, Tensor
 from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, ImageDType
 from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
+from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
 from tinygrad.uop.symbolic import symbolic_simple
-from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
+from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp, RANGEIFY
 from tinygrad.schedule.kernelize import merge_views, get_kernelize_map, Kernel
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
@@ -1861,14 +1861,24 @@ class TestSchedule(unittest.TestCase):
     run_schedule(check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2))
     np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]])
 
-  def test_assign_non_contiguous(self):
-    x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize()
-    y = Tensor.randint(4, 2).contiguous().realize()
-    a = Tensor.arange(8).reshape(4, 2)+y
-    x.shrink((None, (0, 2))).assign(a).realize()
-    xref = np.zeros((4, 4), dtype=int)
-    xref[:, :2] = np.arange(8).reshape(4, 2)+y.numpy()
+  def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
+  def test_assign_non_contiguous(self, alt=False):
+    x = (Tensor.arange(16)-100).reshape(4,4).contiguous().realize()
+    xref = x.numpy()
+    if alt:
+      y = Tensor.randint(2, 4).contiguous().realize()
+      a = Tensor.arange(8).reshape(2, 4)+y
+      tst = x.shrink(((0, 2), None)).assign(a).realize()
+      xref[:2, :] = np.arange(8).reshape(2, 4)+y.numpy()
+    else:
+      y = Tensor.randint(4, 2).contiguous().realize()
+      a = Tensor.arange(8).reshape(4, 2)+y
+      tst = x.shrink((None, (0, 2))).assign(a).realize()
+      xref[:, :2] = np.arange(8).reshape(4, 2)+y.numpy()
     np.testing.assert_equal(x.numpy(), xref)
+    if RANGEIFY > 0:
+      # NOTE: this is a bug on non rangeify
+      np.testing.assert_equal(tst.numpy(), a.numpy())
 
   def test_sparse_categorical_crossentropy_simple(self):
     X = Tensor([[0, 2, 3], [1, 2, 3]]).realize()
@@ -2138,84 +2148,6 @@ class TestSimplifier(unittest.TestCase):
     assert UPat(Ops.CONST, arg=False).match(sink, {}), f"expected {sink} to collapse to a const False"
     assert sink.shape == a.shape
 
-tensor_const_pm = PatternMatcher([
-  (UPat(Ops.CONST, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),)),)), lambda: True),
-  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),)))), UPat(Ops.CONST))), lambda: True),
-])
-class TestConst(unittest.TestCase):
-  # ** part 1: basic functionality of a tensor directly created from CONST
-
-  def test_tensor_const(self):
-    a = Tensor(1)
-    print(a.uop)
-    self.assertTrue(tensor_const_pm.rewrite(a.uop))
-
-  def test_tensor_variable(self):
-    vv = UOp.variable("a", 0, 10).bind(1)
-    a = Tensor(vv)
-    print(a.uop)
-    self.assertTrue(tensor_const_pm.rewrite(a.uop))
-
-  def test_const_schedule(self):
-    a = Tensor.ones((4, 4))
-    sched = a.schedule()
-    self.assertEqual(len(sched), 0)
-
-  def test_const_contiguous_schedule(self):
-    # this ends up in the big graph
-    a = Tensor.ones((4,)).contiguous()
-    sched = a.schedule()
-    self.assertEqual(len(sched), 1)
-
-  # ** part 2: scheduler behavior when const folding happens later
-
-  def test_const_folding_no_realize(self):
-    a = Tensor([1, 2, 3, 4])*0
-    sched = a.schedule()
-    self.assertEqual(len(sched), 0)
-
-  def test_src_const_folding(self):
-    with Context(TRACK_MATCH_STATS=0):
-      a = Tensor.full((4,), 1).contiguous().realize()
-      b = Tensor.full((4,), 2).contiguous().realize()
-    mul0 = a*0
-    add = b+mul0
-    sched = add.schedule()
-    self.assertEqual(len(sched), 0)
-    # b+0 and b share the same underlying device memory
-    self.assertIs(add.uop.buffer, b.uop.buffer)
-    self.assertListEqual(add.tolist(), [2, 2, 2, 2])
-
-  def test_src_masked_const_folding(self):
-    with Context(TRACK_MATCH_STATS=0):
-      a = Tensor.full((4,), 1).contiguous().realize()
-      b = Tensor.full((6,), 2).contiguous().realize()
-    mul0 = a*0
-    add = b+mul0.pad((1, 1), value=2)
-    sched = add.schedule()
-    self.assertEqual(len(sched), 1)
-    run_schedule(sched)
-    # add gets assigned to a new buffer
-    self.assertIsNot(add.uop.base.realized, b.uop.base.realized)
-    self.assertListEqual(add.tolist(), [4, 2, 2, 2, 2, 4])
-
-  # ** part 3: Tensor variable bindings
-
-  #@unittest.expectedFailure # TODO: should schedule assert if you try to realize a Variable?
-  def test_var_schedule(self):
-    vv = UOp.variable("a", 0, 10).bind(1)
-    a = Tensor(vv)
-    sched = a.schedule()
-    self.assertEqual(len(sched), 0)
-
-  def test_add_tvar(self):
-    vv = UOp.variable("a", 0, 10).bind(1)
-    a = Tensor(vv)+2
-    sched, var_vals = a.schedule_with_vars()
-    self.assertEqual(len(sched), 1)
-    run_schedule(sched, var_vals)
-    self.assertEqual(a.tolist(), 3)
-
 @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu")
 class TestCopyFolding(unittest.TestCase):
   def test_const_copy_is_free(self):

test/test_setitem.py

@@ -165,6 +165,17 @@ class TestSetitem(unittest.TestCase):
     t[idx] = val
     self.assertEqual(t.tolist(), [val]*idx_size+[idx_size])
 
+  def test_setitem_advanced_indexing(self):
+    # Example from https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing
+    t = Tensor.zeros(10,20,30,40,50).contiguous()
+    ind_1 = Tensor([5,3,7,8])
+    ind_2 = Tensor([[[0],[1],[2]],[[3],[4],[5]]])
+    v = Tensor.arange(2*3*4*10*30*50).reshape(2,3,4,10,30,50)
+    t[:, ind_1, :, ind_2, :] = v
+    n = np.zeros((10,20,30,40,50))
+    n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
+    np.testing.assert_allclose(t.numpy(), n)
+
 class TestWithGrad(unittest.TestCase):
   def test_no_requires_grad_works(self):
     z = Tensor.rand(8, 8)

test/test_symbolic_jit.py

@@ -16,6 +16,18 @@ class TestSymbolicJit(unittest.TestCase):
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
+  @unittest.expectedFailure # TODO: fix, this works without jit
+  def test_plus1_pad(self):
+    def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(a[:, :vi]).numpy()
+      expected = f(a[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
   def test_add(self):
     def f(a, b): return (a+b).realize()
     jf = TinyJit(f)

test/test_symbolic_ops.py

@@ -17,6 +17,15 @@ class TestSymbolicOps(unittest.TestCase):
       expected = f(a[:, :i]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
+  def test_plus1_pad(self):
+    def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
+    a = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = f(a[:, :vi]).numpy()
+      expected = f(a[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
   def test_add(self):
     def f(a, b): return (a+b).realize()
     a = Tensor.rand(3, 10)

test/unit/test_device.py

@@ -64,6 +64,15 @@ class TestDevice(unittest.TestCase):
                         shell=True, check=True, env={**os.environ, "DEV": "AMD", "AMD_HIP": "1", "AMD_LLVM": "1"})
     else: self.skipTest("only run on CPU/AMD")
 
+  def test_compiler_envvar(self):
+    d = Device[Device.DEFAULT]
+    dname = Device.DEFAULT.split(':')[0].upper()
+    assert d._get_compiler_envvar(type("Compiler", (), {})) == f"{dname}_COMPILER"
+    assert d._get_compiler_envvar(type("LLVMCompiler", (), {})) == f"{dname}_LLVM"
+    assert d._get_compiler_envvar(type("RandomCompiler", (), {})) == f"{dname}_RANDOM"
+    assert d._get_compiler_envvar(type(f"{dname}Compiler", (), {})) == f"{dname}_{dname}COMPILER" # do not repeat device name alone
+    assert d._get_compiler_envvar(type(f"{dname}LLVMCompiler", (), {})) == f"{dname}_LLVM" # do not repeat device name
+
 class MockCompiler(Compiler):
   def __init__(self, key): super().__init__(key)
   def compile(self, src) -> bytes: return src.encode()
@@ -92,7 +101,7 @@ class TestCompiler(unittest.TestCase):
 class TestRunAsModule(unittest.TestCase):
   def test_module_runs(self):
     p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE,
-      env={**os.environ, "DEBUG": "1"}, timeout=10,)
+      env={**os.environ, "DEBUG": "1"}, timeout=30,)
     out = (p.stdout + p.stderr).decode()
     self.assertEqual(p.returncode, 0, msg=out)
     self.assertIn("CPU", out) # for sanity check

test/unit/test_dtype_spec.py

@@ -6,7 +6,6 @@ from tinygrad.helpers import getenv, CI, DEBUG
 from hypothesis import given, settings, strategies as strat
 import numpy as np
 import torch
-import ml_dtypes
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -190,7 +189,7 @@ class TestHelpers(unittest.TestCase):
     elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), math.copysign(math.nan, x))
     elif x > FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), FP8E4M3_MAX)
     elif x < -FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), -FP8E4M3_MAX)
-    else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), ml_dtypes.float8_e4m3fn(x))
+    else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), torch.tensor(x, dtype=torch.float8_e4m3fn).float().item())
 
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
   def test_truncate_fp8e5m2(self, x):
@@ -198,7 +197,7 @@ class TestHelpers(unittest.TestCase):
     elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), x)
     elif x > FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), FP8E5M2_MAX)
     elif x < -FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), -FP8E5M2_MAX)
-    else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), ml_dtypes.float8_e5m2(x))
+    else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), torch.tensor(x, dtype=torch.float8_e5m2).float().item())
 
 class TestTypeSpec(unittest.TestCase):
   def setUp(self):

test/unit/test_tensor_uop_representation.py

@@ -34,7 +34,7 @@ class TestTensorMutates(unittest.TestCase):
     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(d.uop, UPat(Ops.VIEW, src=(realized_pattern,)))
+    assert d.uop is not d.uop.base
 
   def test_reshape_is_same_child(self):
     a = Tensor([1,2,3])
@@ -58,40 +58,6 @@ class TestTensorUopRepresentation(unittest.TestCase):
     print(c.uop)
     is_pattern(c, UPat(Ops.ADD, src=(realized_pattern, realized_pattern)))
 
-  def test_const_pattern(self):
-    a = Tensor(1)
-    print(a.uop)
-    is_pattern(a, const_pattern) # const in tensor has a DEVICE and VIEW src
-    is_pattern(a, UPat.cvar("x")) # even cvar works!
-
-  def test_consts_do_not_realize(self):
-    a = Tensor(1)
-    print(a.uop)
-    pre_realize = a.uop
-    a.realize()
-    assert a.uop is pre_realize
-
-  def test_viewed_consts_do_not_realize(self):
-    a = Tensor.ones(10, 10)
-    print(a.uop)
-    a.realize()
-    is_pattern(a, const_pattern)
-    self.assertEqual(a.uop.shape, (10, 10))
-
-  # CONST is EXPAND -> RESHAPE -> CONST -> DEVICE
-  def test_consts_dont_have_buffers(self):
-    a = Tensor.ones(10, 10)
-    buffers_in_parents = [x.op for x in a.uop.toposort() if x.op is Ops.BUFFER]
-    self.assertEqual(len(buffers_in_parents), 0)
-    is_pattern(a, UPat(Ops.EXPAND, src=(UPat(Ops.RESHAPE, src=(const_pattern,)),)))
-
-  # COPY has a copyin source and a device.
-  def test_copyin(self):
-    a = Tensor([1.,2,3]).realize()
-    c = a.to("TEST")   # NOTE: this isn't checked
-    print(c.uop)
-    is_pattern(c, UPat(Ops.COPY, src=(realized_pattern, UPat(Ops.DEVICE)), arg=None))
-
   def test_empty_buf(self):
     a = Tensor.empty(3, 3)
     is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),)))

tinygrad/codegen/__init__.py

@@ -19,7 +19,7 @@ from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, blo
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
 from tinygrad.codegen.opt.postrange import pm_postrange_opt
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range
-from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
+from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen
 
 @dataclass
 class RewriteStep:
@@ -76,7 +76,7 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
   ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
 
   # add locals
-  ret.append(RewriteStep(pm_add_buffers_local+rangeify_codegen, name="add local buffers"))
+  ret.append(RewriteStep(pm_add_buffers+rangeify_codegen, name="add local buffers"))
 
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce

tinygrad/codegen/late/devectorizer.py

@@ -258,7 +258,8 @@ pm_render = PatternMatcher([
   (UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x),
   # give any loads that are masked an alt value
   (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"),
-    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE) else None),
+    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:])
+      if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None),
   # gate any stores that aren't gated with ifs
   (UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="idx").or_casted(), UPat()), name="store", allow_any_len=True),
     lambda store,idx: UOp(Ops.STORE, dtype=store.dtype, src=store.src[:2]+(UOp(Ops.IF, src=(idx.src[2],)),)+store.src[2:]) if \

tinygrad/codegen/opt/postrange.py

@@ -2,14 +2,15 @@ from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad, GroupOp
 from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace, dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod
 from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
-from tinygrad.schedule.rangeify import remove_tags
+
+remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
 axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
@@ -315,12 +316,12 @@ def apply_opts(ctx:Renderer, ast:UOp):
   if ast.tag is not None: return None
   k = Scheduler(ast, ctx)
   k.convert_loop_to_global()
-  if BEAM >= 1:
+  if ast.arg is not None and ast.arg.opts_to_apply is not None:
+    for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
+  elif BEAM >= 1:
     from tinygrad.codegen.opt.search import beam_search
     rawbufs = bufs_from_ast(ast, ctx.device)
     k = beam_search(k, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
-  elif ast.arg is not None and ast.arg.opts_to_apply is not None:
-    for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
   elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()):
     from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
     # NOTE: hand_coded_optimizations doesn't support multiblock opts yet

tinygrad/device.py

@@ -279,22 +279,24 @@ class Compiled:
 
   def __init__(self, device:str, allocator:Allocator, compilers:Sequence[CompilerPairT]|None, runtime, graph=None, group_id=None):
     self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id
-    compilers = cast(list[CompilerPairT], compilers or [(Renderer, Compiler)])
+    self.compilers = cast(list[CompilerPairT], compilers or [(Renderer, Compiler)])
 
-    devname = device.split(':')[0].upper()
-    envnames = [f"{devname}_{unwrap_class_type(c).__name__.removesuffix('Compiler').removeprefix(devname).upper()}" for r,c in compilers]
-
-    enable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, compilers) if en is not None and getenv(en, -1) == 1)
-    disable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, compilers) if en is not None and getenv(en, -1) == 0)
+    envnames = [self._get_compiler_envvar(c) for r,c in self.compilers]
+    enable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 1)
+    disable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 0)
 
     if len(enable_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {enable_comps}")
-    for _, comp_pair in disable_comps: compilers.remove(comp_pair)
+    for _, comp_pair in disable_comps: self.compilers.remove(comp_pair)
 
-    try: self.renderer, self.compiler = next(self._get_available_compilers([list(enable_comps)[0][1]] if len(enable_comps) == 1 else compilers))
+    try: self.renderer, self.compiler = next(self._get_available_compilers([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers))
     except StopIteration as exc: raise RuntimeError(f"no usable compilers for {self.device}") from exc
 
     if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}")
 
+  def _get_compiler_envvar(self, c):
+    compiler_name = f"{unwrap_class_type(c).__name__.upper().removesuffix('COMPILER').removeprefix(devname:=self.device.split(':')[0].upper())}"
+    return f"{devname}_{compiler_name if len(compiler_name) > 0 else unwrap_class_type(c).__name__.upper()}"
+
   def _get_available_compilers(self, compilers) -> Iterator[tuple[Renderer, Compiler]]:
     for renderer, compiler in compilers:
       with contextlib.suppress(Exception): yield renderer(), compiler()
@@ -357,16 +359,22 @@ if PROFILE:
       launch_viz(PROFILE, fn)
 
 if __name__ == "__main__":
+  from tinygrad import Tensor, Device
+
   for device in ALL_DEVICES:
+    compilers_results, any_works = [], False
     try:
-      _ = Device[device].device
-      try:
-        from tinygrad import Tensor
-        with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
-        if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
-        result = colored("PASS", "green")
-      except Exception as e:
-        result = f"{colored('FAIL', 'yellow')} {e}"
+      default_compiler = (d:=Device[device]).compiler
+      for i,(r,c) in enumerate(d.compilers):
+        try:
+          d.renderer, d.compiler = r(), c()
+          with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
+          if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
+          default_text = '(default)' if type(default_compiler) is type(d.compiler) else f'({d._get_compiler_envvar(c)}=1 to make default)'
+          compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
+          any_works = True
+        except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {unwrap_class_type(c).__name__}: {e}")
+      result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results])
     except Exception as e:
       result = f"{colored('FAIL', 'red')} {e}"
     print(f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}")

tinygrad/engine/realize.py

@@ -140,13 +140,13 @@ class BufferXfer(BufferCopy):
 
 # **************** method cache ****************
 
-method_cache: dict[tuple[str, bytes, tuple[int, ...], bool], CompiledRunner] = {}
+method_cache: dict[tuple[str, type, bytes, tuple[int, ...], bool], CompiledRunner] = {}
 def get_runner(device:str, ast:UOp) -> CompiledRunner:
   # TODO: this should be all context relevant to rendering
   context = (BEAM.value, NOOPT.value, DEVECTORIZE.value)
-  ckey = (device, ast.key, context, False)
+  ckey = (device, type(Device[device].compiler), ast.key, context, False)
   if cret:=method_cache.get(ckey): return cret
-  bkey = (device.split(":")[0], ast.key, context, True)
+  bkey = (device.split(":")[0], type(Device[device].compiler), ast.key, context, True)
   if bret:=method_cache.get(bkey):
     method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device), bret.lib)
   else:

tinygrad/runtime/support/compiler_cpu.py

@@ -13,7 +13,7 @@ class ClangJITCompiler(Compiler):
     # x18 is a reserved platform register. It is clobbered on context switch in macos and is used to store TEB pointer in windows on arm, don't use it
     target = 'x86_64' if sys.platform == 'win32' else platform.machine()
     # on arm march means "runs on this arch and superset" instead of "optimize for this arch". x86 march == arm mcpu
-    arch = '-march=native' if platform.machine() in ('x86_64', 'AMD64') else '-mcpu=native'
+    arch = {'x86_64': '-march=native', 'AMD64': '-march=native', 'riscv64': '-march=rv64g'}.get(platform.machine(), "-mcpu=native")
     args = [arch, f'--target={target}-none-unknown-elf', '-O2', '-fPIC', '-ffreestanding', '-fno-math-errno', '-nostdlib', '-fno-ident']
     arch_args = ['-ffixed-x18'] if target == 'arm64' else []
     obj = subprocess.check_output([getenv("CC", 'clang'), '-c', '-x', 'c', *args, *arch_args, '-', '-o', '-'], input=src.encode('utf-8'))
@@ -29,7 +29,7 @@ def expect(x, err, ret=None):
 
 class LLVMCompiler(Compiler):
   jit = True
-  target_arch = {'arm64': 'AArch64', 'aarch64': 'AArch64', 'x86_64': 'X86', 'AMD64': 'X86'}[platform.machine()]
+  target_arch = {'arm64': 'AArch64', 'aarch64': 'AArch64', 'x86_64': 'X86', 'AMD64': 'X86', 'riscv64': 'riscv64'}[platform.machine()]
   def __init__(self, processor:str, feats:str):
     for component in ['Target', 'TargetInfo', 'TargetMC', 'AsmParser', 'AsmPrinter']: getattr(llvm, f'LLVMInitialize{self.target_arch}{component}')()
 

tinygrad/runtime/support/hcq.py

@@ -440,14 +440,19 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     except MemoryError: buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
     return buf, realloced
 
+  def _make_no_iface_error(self, errs:str, err_short:str) -> RuntimeError:
+    # Keep it in a separate function to avoid creating a traceback <-> locals ref cycle
+    e = RuntimeError(f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available")
+    if hasattr(e, "add_note"): e.add_note(errs + err_short)
+    return e
+
   def _select_iface(self, *ifaces:Type):
     errs, err_short = "", ""
     if val:=getenv(f'{type(self).__name__[:-6].upper()}_IFACE', ""): ifaces = tuple(x for x in ifaces if x.__name__.startswith(val.upper()))
     for iface_t in ifaces:
       try: return iface_t(self, self.device_id)
-      except Exception as e: errs, err_short = errs + f"\n{iface_t.__name__}: {traceback.format_exc()}", err_short + f"\n{iface_t.__name__}: {e}"
-    raise RuntimeError(f"{errs}\nNo interface for {type(self).__name__[:-6]}:{self.device_id} is available:{err_short}\n" \
-                       f"\nForce an interface with {type(self).__name__[:-6].upper()}_IFACE={('|'.join(x.__name__[:-5] for x in ifaces))}.")
+      except Exception as e: errs, err_short = errs + f"\n{iface_t.__name__}: {traceback.format_exc()}", err_short + f"\n{iface_t.__name__}: {e}."
+    raise self._make_no_iface_error(errs, err_short)
 
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
 

tinygrad/schedule/rangeify.py

@@ -1,15 +1,16 @@
-from typing import Any
+from typing import Any, cast
 import functools, operator
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute
-from tinygrad.uop.symbolic import sym
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify, graph_rewrite_map
+from tinygrad.uop.symbolic import sym, symbolic_simple
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup
 from tinygrad.schedule.multi import multi_pm
 
 from tinygrad.schedule.kernelize import Kernel
-from tinygrad.uop.ops import track_rewrites, graph_rewrite_map, graph_rewrite, identity_element, sint, AxisType
+from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType
 
+# *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
 double_reshape = PatternMatcher([
@@ -19,30 +20,42 @@ double_reshape = PatternMatcher([
 
 earliest_rewrites = double_reshape+PatternMatcher([
   # non shape changing RESHAPE is NOOP
-  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0] if x.src[0].shape == x.arg else None),
+  #(UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0] if x.src[0].shape == x.arg else None),
+  # DETACH and CONTIGUOUS_BACKWARD are NOOPs here, so is FUSE
+  #(UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0].f(Ops.NOOP, tag=x.tag)),
+
+  # just removing it works...
+  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]),
+
+  # preserve tags?
   # UOp with size 0 is zero
   (UPat(GroupOp.All-{Ops.SINK}, name="root"), lambda root: root.const_like(0) if root.base.st is not None and root.size == 0 else None),
   # reduce of size 0 is the identity element
   (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),
-  # DETACH and CONTIGUOUS_BACKWARD are NOOPs here, so is FUSE
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]),
+
+  # copy reorder
+  # TODO: this is causing many copies wih the replace tag None
   # RESHAPE after COPY
-  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).reshape(r.arg)),
+  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d), tag=None).reshape(r.arg)),
   # TODO: this should be BUFFER_VIEW
-  (UPat(Ops.COPY, src=(UPat(Ops.SHRINK, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).shrink(r.arg)),
+  (UPat(Ops.COPY, src=(UPat(Ops.SHRINK, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d), tag=None).shrink(r.arg)),
+
   # const hacks
-  (UPat(Ops.CONST, name="x"), lambda x:
-   x.replace(src=(x.src[0].src[0],)).reshape((1,)*len(x.shape)).expand(x.shape) if \
-    len(x.src) and x.src[0].op is Ops.VIEW and not any(s == 0 for s in x.shape) else None),
+  #(UPat(Ops.CONST, name="x"), lambda x:
+  # x.replace(src=(x.src[0].src[0],)).reshape((1,)*len(x.shape)).expand(x.shape) if \
+  #  len(x.src) and x.src[0].op is Ops.VIEW and not any(s == 0 for s in x.shape) else None),
+
   # assign only to buffer
-  (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}, name="target"), UPat(name="x"))),
-   lambda x,target: x if target.base.op is not Ops.BUFFER else None),
+  (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),
+
   # 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]),
+  #(UPat(Ops.CONTIGUOUS, name="root", src=(UPat((Ops.CONTIGUOUS, Ops.BUFFER, Ops.COPY, Ops.ASSIGN)),)), lambda root: root.src[0]),
 ])
 
-# 1. add contiguous where we have to
+# *****************
+# 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,
@@ -68,11 +81,17 @@ do_realize = PatternMatcher([
   (UPat(Ops.ASSIGN, name="a"), realize_assign),
 ])
 
-add_contiguous = PatternMatcher([
-  (UPat(GroupOp.All, name="x"), lambda ctx,x: x.replace(tag=1).realize() if x in ctx and x.tag is None else None),
-])
-remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
+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),
+])
+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
@@ -99,7 +118,8 @@ pm_children = PatternMatcher([
   (UPat(GroupOp.All-{Ops.CHILD, Ops.CHILDREN}, name="x"), mark_children),
 ])
 
-# 3. rangeify
+# *****************
+# 3a. rangeify (movement)
 
 @dataclass
 class RangeifyContext:
@@ -175,12 +195,18 @@ pm_mops = PatternMatcher([
   (UPat(Ops.PAD, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), map_pad),
 ])
 
+# *****************
 # 3b. rangeify (ops)
 
+# bufferization can happen in three ways
+#   1. there's an explicit REALIZE in the graph
+#   2. the ranges from the children don't match and we have to create a buffer (only on children)
+#   3. might_end_axis triggers because we should be closing a loop to save compute
+
 @dataclass(frozen=True)
 class BufferizeOpts:
   # on AddrSpace.LOCAL, device is the id
-  device: str|tuple[str, ...]|int
+  device: str|tuple[str, ...]|int|None
   addrspace: AddrSpace = AddrSpace.GLOBAL
 
 def map_partial_realize(ctx:RangeifyContext, x:UOp, idx:UOp):
@@ -195,21 +221,17 @@ def map_partial_realize(ctx:RangeifyContext, x:UOp, idx:UOp):
       ranges.append(idx.src[1+i])
       continue
     passthrough_idx.append(idx.src[1+i])
-    ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.index, 0))
+    ranges.append(ctx.new_range(s))
     new_ranges.append(ranges[-1])
-  ret = x.src[0].index(*ranges).bufferize(*[x for x in new_ranges if x.op is not Ops.CONST], arg=BufferizeOpts(device=x.device))
+  # TODO: this should be able to be global or local
+  ret = x.src[0].index(*ranges).bufferize(*[x for x in new_ranges if x.op is not Ops.CONST],
+                                          arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL))
   return ret.index(*passthrough_idx)
 
 def map_realize(ctx:RangeifyContext, x:UOp):
   if x.arg is not None: return None
-  ranges = []
-  for s in x.shape[len(x.src)-1:]:
-    ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.index, 0))
-  ret = x.src[0].index(*ranges).bufferize(*x.src[1:], *[x for x in ranges if x.op is not Ops.CONST], arg=BufferizeOpts(device=x.device))
-  # was there a shrink? move this before the bufferize?
-  # TODO: do we need this?
-  if resolve(prod(x.shape) != prod(ret.shape)): ret = ret.forced_reshape((prod(ret.shape),)).shrink(((0, prod(x.shape)),))
-  return ret.forced_reshape(x.shape)
+  ranges = [ctx.new_range(s) for s in x.shape]
+  return x.src[0].index(*ranges).bufferize(*x.src[1:], *ranges, arg=BufferizeOpts(device=x.device), tag=x.src[0].tag)
 
 def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp):
   rngs = list(idx.src[1:])
@@ -218,7 +240,7 @@ def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp):
     if i in red.arg[1]:
       rngs[i] = ctx.new_range(s, axistype=AxisType.REDUCE)
       new_ranges.append(rngs[i])
-  return UOp(Ops.REDUCE, red.dtype, src=(red.src[0].index(*rngs),)+tuple(new_ranges), arg=red.arg[0])
+  return UOp(Ops.REDUCE, red.dtype, src=(red.src[0].index(*rngs),)+tuple(new_ranges), arg=red.arg[0], tag=red.tag)
 
 def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
   if c not in ctx.seen_children: ctx.seen_children[c] = {}
@@ -256,7 +278,14 @@ def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
   # index based on the shared ranges
   ret = c.index(*out_rngs)
   # if all ranges aren't the same between children, we have to bufferize
-  if len(idx_ranges) > 0: ret = ret.bufferize(*end_ranges, arg=BufferizeOpts(device=x.device)).index(*[idx.src[1+i] for i in idx_ranges])
+  if len(idx_ranges) > 0:
+    if len(idx_ranges) == len(out_rngs):
+      # this is a global bufferize
+      ret = ret.bufferize(*end_ranges, arg=BufferizeOpts(device=x.device))
+    else:
+      assert RANGEIFY > 1, "this isn't supported with RANGEIFY=1"
+      ret = ret.bufferize(*end_ranges, arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL))
+    ret = ret.index(*[idx.src[1+i] for i in idx_ranges])
   return ret
 
 def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):
@@ -266,7 +295,7 @@ def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):
 def might_end_axis(idx:UOp):
   if idx.arg is None: return None
   # TODO: write a proper cost function here
-  if all(x.op not in {Ops.BUFFER, Ops.CONTIGUOUS, Ops.BUFFERIZE} for x in idx.toposort()): return None
+  if all(x.op not in {Ops.BUFFER, Ops.REALIZE, Ops.BUFFERIZE} for x in idx.toposort()): return None
   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:]):
@@ -275,6 +304,8 @@ def might_end_axis(idx:UOp):
   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)
 
+def unprocessed_index(x:UOp): raise RuntimeError(f"unprocessed index on {x.src[0].op}")
+
 pm_rangeify = pm_mops+PatternMatcher([
   # sink contigs to kick it off
   (UPat(Ops.REALIZE, src=(UPat(),), name="x", allow_any_len=True), map_realize),
@@ -291,27 +322,36 @@ pm_rangeify = pm_mops+PatternMatcher([
   # CONST (or DEFINE_VAR) can't have axes. remove srcs when we INDEX it
   (UPat(Ops.INDEX, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),)), lambda c: c.replace(src=())),
 
+  # copy on CONST is CONST
+  (UPat(Ops.COPY, src=(UPat.cvar("c"), UPat())), lambda c: c),
+
   # handle arg on any op with weight. old endrange stuff
   (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.REDUCE_AXIS})),), allow_any_len=True, name="idx"), might_end_axis),
 
+  # 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],))),
+
   # move MAP through elementwise ALU / reduce. these are the items with cost
   (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union(
-    {Ops.STORE, Ops.ASSIGN, Ops.COPY, Ops.DEVICE, Ops.BIND, Ops.CONTIGUOUS})),), allow_any_len=True, name="x"),
+    {Ops.STORE, Ops.COPY, Ops.DEVICE, Ops.BIND, Ops.CONTIGUOUS, Ops.NOOP})),), allow_any_len=True, name="x"),
    lambda x: x.src[0].replace(src=tuple([s.index(*x.src[1:]) for s in x.src[0].src]))),
   (UPat(Ops.INDEX, src=(UPat(Ops.REDUCE_AXIS, name="red"),), allow_any_len=True, name="idx"), map_reduce),
+
+  # assert if there's any index we didn't process
+  (UPat(GroupOp.All-{Ops.REALIZE, Ops.BUFFERIZE}).f(Ops.INDEX, name="x"), unprocessed_index),
 ])
 
+# *****************
 # 3.5 cleanups
 
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
-# TODO: figure out how to reenable this
 def cleanup_dead_axes(b:UOp):
-  parents = b.src[0].toposort()
   new_rng = []
   hit = False
   reshape: list[sint] = []
   for s,rng in zip(b.shape, b.src[1:]):
-    if rng not in parents and rng.op is Ops.RANGE:
+    if rng not in b.src[0].sparents and rng.op is Ops.RANGE:
       reshape.append(1)
       hit = True
     else:
@@ -327,31 +367,35 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   assert len(buf.src) == len(idx.src), "index on wrong bufferize"
   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
+
   # here is where we compute the cost
   # for now just no REDUCE, COPY, or ASSIGN
-  # TODO: exclude fusion of user contiguous
-  #ran = src.toposort(gate=lambda x: x.op not in {Ops.INDEX})
-  #if any(x.op in {Ops.REDUCE, Ops.COPY, Ops.ASSIGN} for x in ran): return None
+  ran = src.toposort(gate=lambda x: x.op not in {Ops.INDEX})
+  if any(x.op in {Ops.REDUCE, Ops.COPY, Ops.ASSIGN} for x in ran): return None
 
   # simple, matching old behavior
-  if src.op is not Ops.INDEX: return None
+  #if src.op is not Ops.INDEX: return None
 
   # this is the ranges replaced
   return src.substitute(dict(zip(buf.src[1:], idx.src[1:])))
 
-
 pm_cleanups = double_reshape+pm_mops+PatternMatcher([
   #(UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes),
   # remove noop buffers. if we look at the next index we can remove even more of these
   # NOTE: this is mostly the same case as below, but if there's no INDEX this gets more
-  #(UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"),
-  # lambda idx,b2: idx.src[0] if idx.src[1:] == b2.src[1:] else None),
+  (UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"),
+   lambda idx,b2: idx.src[0].replace(tag=nt if len(nt:=(idx.src[0].tag or ()) + (b2.tag or ())) else None) if idx.src[1:] == b2.src[1:] else None),
   # remove reindexing with cost function
   (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: c.reshape((1,)*len(b.shape)).expand(b.shape)),
+  (UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: c.reshape((1,)*len(b.shape)).expand(b.shape)),
+  # 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=())),
 ])
 
+# *****************
 # 4. put in buffers for bufferize
 # TODO: should BUFFERIZE look a lot more like STORE
 # BUFFERIZE has device in arg
@@ -359,27 +403,44 @@ pm_cleanups = double_reshape+pm_mops+PatternMatcher([
 # BUFFERIZE returns the BUFFER ready for INDEXing (doing this will make splitting a lot easier)
 # NOTE: this has been fixed up a bit
 
-def bufferize_to_store(x:UOp, locals_allowed=False):
+def bufferize_to_store(x:UOp):
   rngs = x.src[1:]
   shape = tuple([int(r.vmax+1) for r in rngs])
+  sym_shape = tuple([ssimplify(r.src[0]) for r in rngs])
   size = prod(shape)
   assert size > 0, f"no zero sized buffers {shape}"
+
   sdtype = x.dtype.ptr(size=size, addrspace=x.arg.addrspace)
   if x.src[0].op is Ops.ASSIGN:
-    assign_target, assign_src = x.src[0].src
+    assign_target, assign_src, assign_mops = x.src[0].src
     assert assign_target.op is Ops.INDEX
-    return assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype)
+    # in assign, this is the buffer size, not the bufferize size
+    # TODO: assign_mops here
+    ret = assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=x.dtype)
+    mops = []
+    walk = assign_mops
+    while walk is not assign_mops.base:
+      mops.append((walk.op, walk.arg))
+      walk = walk.src[0]
+    for m in mops[::-1]: ret = ret._mop(*m)
+    return ret.forced_reshape(shape).replace(tag=x.tag)
+
   # NOTE: the DEFINE_LOCAL needs to be disambiguated here
   if sdtype.addrspace == AddrSpace.GLOBAL:
     buf = UOp.new_buffer(x.arg.device, size, x.dtype)
-  else:
-    if not locals_allowed: return None
-    buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=x.arg.device)
-  return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype)
+    ret = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=x.dtype)
+    ret = ret.forced_reshape(shape)
+    # TODO: is this right? what if it's offset
+    if shape is not sym_shape: ret = ret.shrink(tuple([(0,x) for x in sym_shape]))
+    return ret.replace(tag=x.tag)
 
-pm_add_buffers_local = pm_mops+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), lambda x: bufferize_to_store(x, True)),
-])
+  # handle locals
+  tag = x.arg.device
+  if tag is None: tag = UOp.unique().arg # TODO: hack
+  buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=tag)
+  # store has the other dtype here
+  # 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([
   (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
@@ -389,6 +450,7 @@ pm_add_buffers = pm_mops+PatternMatcher([
    lambda m: m.replace(src=tuple([x.src[0] for x in m.src])).reshape(m.src[0].arg)),
 ])
 
+# *****************
 # 5. split into kernels
 
 @dataclass
@@ -396,6 +458,7 @@ class LocalAddBufferContext:
   dg:int = 0
   map:dict = field(default_factory=dict)
   vars:dict = field(default_factory=dict)
+  range:int = 0
 
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
   ret = UOp(Ops.DEFINE_GLOBAL, buf.dtype.ptr(buf.arg), arg=ctx.dg)
@@ -415,6 +478,12 @@ def handle_assign(ctx:LocalAddBufferContext, assign:UOp):
   ctx.map[buf] = assign
   return buf
 
+def renumber_range(ctx:LocalAddBufferContext, r:UOp):
+  if r.tag is not None: return None
+  ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=())
+  ctx.range += 1
+  return ret
+
 to_define_global = PatternMatcher([
   (UPat(Ops.BUFFER, name="buf"), debuf),
   (UPat(Ops.BIND, name="b"), unbind_kernel),
@@ -423,12 +492,18 @@ to_define_global = PatternMatcher([
   # HACK in case any CONSTs were replaced
   # this is only needed if you are using symbolic
   #(UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
+
+  # renumber the ranges starting with 0 so that kernel deduping works
+  (UPat(Ops.RANGE, name="r"), renumber_range),
 ])
 
 rangeify_codegen = PatternMatcher([
-  # no CONTIGUOUS in the kernel graph
+  # no NOOP in the kernel graph
   # TODO: this can be moved into codegen?
-  (UPat(Ops.CONTIGUOUS, name="x"), lambda x: x.src[0]),
+  (UPat((Ops.NOOP, Ops.CONTIGUOUS), name="x"), lambda x: x.src[0]),
+
+  # strip the arg from store
+  (UPat(Ops.STORE, name="x"), lambda x: x.replace(arg=None) if x.arg is not None else None),
 
   # add loads to non ptr indexes
   # TODO: this can be moved into codegen?
@@ -444,41 +519,73 @@ rangeify_codegen = PatternMatcher([
    lambda src, barrier, gate: src.load(UOp(Ops.IF, src=(gate, barrier)))),
 ])
 
-def split_store(x:UOp):
+def split_store(ctx:list[UOp], x:UOp):
   if len(x.ranges): return None
-  ctx = LocalAddBufferContext()
-  ret = graph_rewrite(x, to_define_global+rangeify_codegen, ctx=ctx, name="kernel split", bottom_up=True)
+  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)
+
+  # gather the metadata
+  metadatas = [ctx[y].metadata for x in ret.sparents if x.tag is not None for y in x.tag]
 
   # NOTE: the hack for COPY is here
   ret = ret.sink() if ret.src[1].op is not Ops.COPY else ret.src[1]
-  kernel = UOp(Ops.KERNEL, src=tuple(ctx.map.values())+tuple(ctx.vars.keys()), arg=Kernel(ret,()))
+  kernel_arg = Kernel(ret,tuple(dedup(flatten([x for x in metadatas if x is not None]))))
+  kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg)
   return x.as_buf().assign(kernel)
 
 split_kernels = PatternMatcher([
   (UPat(Ops.STORE, name="x"), split_store),
 ])
 
-@track_rewrites(name=lambda sink,ret: f"Schedule {pluralize('Kernel',len([u for u in ret[sink].toposort() if u.op is Ops.KERNEL]))}", replay=True)
-def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
-  tensor_map = graph_rewrite_map(sink, multi_pm+earliest_rewrites, name="earliest")
-  realize_map: dict[UOp, UOp] = {}
-  graph_rewrite(tensor_map[sink], do_realize, ctx=realize_map, name="Input Graph")
-  tensor_map = graph_rewrite_map(tensor_map[sink], add_contiguous, ctx=realize_map, bottom_up=True, input_map=tensor_map, name="add realize")
-  tensor_map = graph_rewrite_map(tensor_map[sink], remove_tags, input_map=tensor_map, name="remove tags")
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_children, ctx=ChildrenContext(), bottom_up=True, input_map=tensor_map, name="children")
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_rangeify, ctx=RangeifyContext(), bottom_up=True, input_map=tensor_map, name="rangeify")
-  # NOTE: running symbolic can break the graph, leaving RANGE/INDEX/BUFFERIZE in the final graph
-  #tensor_map = graph_rewrite_map(tensor_map[sink], symbolic_simple, input_map=tensor_map, name="symbolic")
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_cleanups, bottom_up=True, input_map=tensor_map, name="buffer cost")
-  if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Rangeify Graph")
+def tag_uop(ctx:list[UOp], x:UOp):
+  if x.tag is not None: return None
+  ctx.append(x)
+  return x.replace(tag=(len(ctx)-1,))
+add_tags = PatternMatcher([
+  # don't tag BUFFERs, they are global
+  (UPat(GroupOp.All-{Ops.BUFFER, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND}, name="x"), tag_uop),
+])
 
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_add_buffers, bottom_up=True, input_map=tensor_map, name="add buffers")
-  tensor_map = graph_rewrite_map(tensor_map[sink], split_kernels, input_map=tensor_map, name="split kernels")
+@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")
+  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
+  tsink = graph_rewrite(tsink, add_contiguous, ctx=realize_map, bottom_up=True, name="add realize")
+  tsink = graph_rewrite(tsink, remove_contig_tags, name="remove contiguous tags")
+  tsink = graph_rewrite(tsink, pm_children, ctx=ChildrenContext(), bottom_up=True, name="get children")
+
+  # rangeify
+  tsink = graph_rewrite(tsink, pm_rangeify, ctx=RangeifyContext(), bottom_up=True, name="rangeify")
+  # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
+  tsink = graph_rewrite(tsink, symbolic_simple, name="symbolic")  # this supports const folding
+  tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
+
+  # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
+  # if it's not tagged by here, it's out
+  tsink = UOp.sink(*[x for x in tsink.parents if x.op is Ops.BUFFERIZE and x.tag is not None])
+
+  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")
+
+  # bufferize -> store
+  tsink = graph_rewrite(tsink, pm_add_buffers, bottom_up=True, name="bufferize to store")
+  tsink = graph_rewrite(tsink, split_kernels, ctx=uop_list, name="split kernels")
 
   # if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign
   kernel_assign: dict[UOp, UOp] = {}
   assign_rep: dict[UOp, UOp] = {}
-  for u in tensor_map[sink].toposort():
+  for u in tsink.toposort():
     if u.op is not Ops.ASSIGN: continue
     kernel_assign[u.buf_uop] = u
     for s in u.src[1].src:
@@ -487,8 +594,14 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
       if any(x.op is Ops.ASSIGN and x.buf_uop is s for x in u.toposort()):
         raise RuntimeError(f"cycle detected in graph, kernel for {u.buf_uop} must either depend on ASSIGN or BUFFER")
       assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
-  if assign_rep:
-    tensor_map = graph_rewrite_map(tensor_map[sink], _substitute, ctx=assign_rep, bottom_up=True, input_map=tensor_map, name="fix_assign")
+  if assign_rep: tsink = graph_rewrite(tsink, _substitute, ctx=assign_rep, bottom_up=True, name="fix_assign")
 
-  if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Kernel Graph")
-  return tensor_map
+  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
+
+  becomes_map: dict[UOp, UOp] = {}
+  for s in tsink.src:
+    assert s.tag is not None
+    for a in s.tag:
+      if a is None: continue
+      becomes_map[uop_list[cast(int, a)]] = s.replace(tag=None)
+  return becomes_map

tinygrad/tensor.py

@@ -1220,8 +1220,8 @@ class Tensor(MathTrait):
       x = (mask.where(x.reshape(reshape_arg), 0)).sum(sum_axis:=tuple(d + len(big_shape) for d in dims), dtype=x.dtype)
 
       # special permute case
-      if dims[0] != 0 and len(dims) != 1 and tuple(dims) != tuple(range(dims[0], dims[-1]+1)):
-        x = x.permute(*range(dims[0], dims[0]+len(big_shape)), *range(0, dims[0]), *range(dims[0]+len(big_shape), x.ndim))
+      if (permuted := dims[0] != 0 and len(dims) != 1 and tuple(dims) != tuple(range(dims[0], dims[-1]+1))):
+        mask, x = (y.permute(*range(dims[0], dims[0]+len(big_shape)), *range(0, dims[0]), *range(dims[0]+len(big_shape), y.ndim)) for y in (mask, x))
 
       # for advanced setitem, returns whole tensor with indices replaced
       if v is not None:
@@ -1229,7 +1229,7 @@ class Tensor(MathTrait):
         # add back reduced dims from sum
         for dim in sum_axis: vb = vb.unsqueeze(dim)
         # run _masked_setitem on tuple of axis that is to be reduced to match self.shape
-        x = _masked_setitem(self, vb, mask, tuple(range(dims[0], dims[0] + len(big_shape))))
+        x = _masked_setitem(self, vb, mask, tuple(range((start := dims[0] if not permuted else 0), start + len(big_shape))))
 
     return x
 

tinygrad/uop/ops.py

@@ -163,6 +163,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     # CONST with a DEVICE has a shape of ()
     if self.op is Ops.CONST and len(self.src) and self.src[0].op is Ops.DEVICE: return ShapeTracker.from_shape(())
     if self.op is Ops.STORE and isinstance(self.dtype, PtrDType): return ShapeTracker.from_shape((self.dtype.size,))
+    if self.op is Ops.STORE and self.dtype is not dtypes.void: return self.src[0].src[0].st
     # BufferOps and ASSIGN flow ShapeTracker from a direct edge
     if self.op in {Ops.STORE, Ops.ASSIGN, Ops.LOAD}: return self.src[0].st
     if self.op in GroupOp.Buffer: return views[0] if (views:=[x.st for x in self.src if x.op is Ops.VIEW]) else None

tinygrad/uop/spec.py

@@ -99,8 +99,11 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
 
   # Tensor const has a device and an unmasked ShapeTracker of stride 0
   # NOTE: variables in shape can cause multiple views in this ShapeTracker and other issues, see TestSymbolicJit.test_ones_sum
-  (UPat(Ops.CONST, src=(UPat(Ops.VIEW, name="st", src=(UPat(Ops.DEVICE),)),)),
+  # TODO: remove after rangeify is default
+  (UPat(Ops.CONST, src=(UPat.any(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),), name="st"),
+                                 UPat(Ops.VIEW, src=(UPat(Ops.DEVICE), UPat(Ops.BIND)), name="st")),)),
    lambda st: len(st.st.views) == 1 and all(v.mask is None for v in st.st.views)),
+  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE),)), lambda: True),
 
   # DETACH and CONTIGUOUS change how we interpret the source UOp
   # CONTIGUOUS ensures the source UOp realizes
@@ -165,7 +168,7 @@ spec = PatternMatcher([
    lambda x,src: isinstance(x.arg, ShapeTracker) and src.op is not Ops.STORE and x.dtype.base == src.dtype.base),
 
   (UPat(Ops.VALID, dtypes.bool, (UPat(Ops.VIEW),)), lambda: True),
-  (UPat(Ops.CONST, name="x"), lambda x: type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))),
+  (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))),
 
   # early LOAD has a <bufview, store?>
   (UPat(Ops.LOAD, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.Defines),)),)), lambda: True),

tinygrad/uop/symbolic.py

@@ -97,9 +97,6 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
   # b.cast(a).cast(b) -> b if a preserves all values in b
   (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_safe_cast(b.dtype, a.dtype) else None),
-  # if the intermediate cast doesnt narrow we can do it in one cast, we have to be carefull with bfloat16
-  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_safe_cast(x.dtype, a.dtype) and
-    not (a.dtype==dtypes.float and (b.dtype==dtypes.bfloat16 or x.dtype==dtypes.bfloat16)) else None),
   # ** pow **
   (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow),
   # positive const ** x
@@ -352,6 +349,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   (UPat.var("x") % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
   (UPat.var("x") % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax <  0 else None),
   # cast/long folding
+  # if the intermediate cast doesnt narrow we can do it in one cast
+  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_safe_cast(x.dtype, a.dtype) else None),
   (UPat.var('x', dtypes.ints+(dtypes.index,)).cast(dtypes.ints+(dtypes.index,), name="a").cast(name="b"),
     lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
   # try to do math in int instead of long

tinygrad/viz/index.html

@@ -134,17 +134,6 @@
   .metadata > * + *, .rewrite-container > * + *, .ctx-list > * + * {
     margin-top: 12px;
   }
-  .stats-list > * + * {
-    margin-top: 8px;
-  }
-  .stats-list > p > * + * {
-    margin-top: 12px;
-  }
-  .stats-list {
-    width: 100%;
-    max-height: 240px;
-    overflow: auto;
-  }
   .ctx-list > ul > * + * {
     margin-top: 4px;
   }

tinygrad/viz/js/index.js

@@ -178,7 +178,7 @@ async function renderProfiler() {
   const u64 = () => { const ret = new Number(view.getBigUint64(offset, true)); offset += 8; return ret; }
   const f32 = () => { const ret = view.getFloat32(offset, true); offset += 4; return ret; }
   const optional = (i) => i === 0 ? null : i-1;
-  const dur = u32(), peak = u64(), indexLen = u32(), layoutsLen = u32();
+  const dur = u32(), tracePeak = u64(), indexLen = u32(), layoutsLen = u32();
   const textDecoder = new TextDecoder("utf-8");
   const { strings, dtypeSize, markers }  = JSON.parse(textDecoder.decode(new Uint8Array(buf, offset, indexLen))); offset += indexLen;
   // place devices on the y axis and set vertical positions
@@ -192,20 +192,20 @@ async function renderProfiler() {
   // color by key (name/device)
   const colorMap = new Map();
   data = {tracks:new Map(), axes:{}};
-  const heightScale = d3.scaleLinear().domain([0, peak]).range([4,maxheight=100]);
+  const heightScale = d3.scaleLinear().domain([0, tracePeak]).range([4,maxheight=100]);
   for (let i=0; i<layoutsLen; i++) {
     const nameLen = view.getUint8(offset, true); offset += 1;
     const k = textDecoder.decode(new Uint8Array(buf, offset, nameLen)); offset += nameLen;
     const div = deviceList.append("div").attr("id", k).text(k).style("padding", padding+"px");
     const { y:baseY, height:baseHeight } = rect(div.node());
     const offsetY = baseY-canvasTop+padding/2;
-    const shapes = [];
+    const shapes = [], visible = [];
     const EventTypes = {TIMELINE:0, MEMORY:1};
     const eventType = u8(), eventsLen = u32();
     if (eventType === EventTypes.TIMELINE) {
       const levelHeight = baseHeight-padding;
       const levels = [];
-      data.tracks.set(k, { shapes, visible:[], offsetY });
+      data.tracks.set(k, { shapes, visible, offsetY });
       let colorKey, ref;
       for (let j=0; j<eventsLen; j++) {
         const e = {name:strings[u32()], ref:optional(u32()), st:u32(), dur:f32(), info:strings[u32()] || null};
@@ -268,10 +268,11 @@ async function renderProfiler() {
       const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
       for (const [num, {dtype, sz, nbytes, y, x:steps}] of buf_shapes) {
         const x = steps.map(s => timestamps[s]);
-        const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}\nnum:${num}`};
+        const dur = x.at(-1)-x[0];
+        const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}\nnum:${num}\nalive for ${formatTime(dur)}`};
         shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
       }
-      data.tracks.set(k, { shapes, visible:[], offsetY, height, peak, scaleFactor:maxheight*4/height });
+      data.tracks.set(k, { shapes, visible, offsetY, height, peak, scaleFactor:maxheight*4/height });
       div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
         const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
         let offset = 0;
@@ -320,15 +321,16 @@ async function renderProfiler() {
         // contiguous rect
         if (e.x>et || e.x+e.width<st) continue;
         const x = xscale(e.x);
+        const y = offsetY+e.y;
         const width = xscale(e.x+e.width)-x;
-        ctx.fillStyle = e.fillColor; ctx.fillRect(x, offsetY+e.y, width, e.height);
-        visible.push({ y0:offsetY+e.y, y1:offsetY+e.y+e.height, x0:x, x1:x+width, arg:e.arg });
+        ctx.fillStyle = e.fillColor; ctx.fillRect(x, y, width, e.height);
+        visible.push({ y0:y, y1:y+e.height, x0:x, x1:x+width, arg:e.arg });
         // add label
         if (e.label == null) continue;
         ctx.textAlign = "left";
         ctx.textBaseline = "middle";
-        let [labelX, labelWidth] = [x+2, 0];
-        const labelY = offsetY+e.y+e.height/2;
+        let labelX = x+2, labelWidth = 0;
+        const labelY = y+e.height/2;
         for (const [i,l] of e.label.entries()) {
           if (labelWidth+l.width+(i===e.label.length-1 ? 0 : ellipsisWidth)+2 > width) {
             if (labelWidth !== 0) ctx.fillText("...", labelX, labelY);
@@ -655,28 +657,6 @@ async function main() {
   const metadata = document.querySelector(".metadata");
   const [code, lang] = ctx.fmt != null ? [ctx.fmt, "cpp"] : [ret[currentRewrite].uop, "python"];
   metadata.replaceChildren(codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeBlock(code, lang, { wrap:false }));
-  if (ctx.runtime_stats != null) {
-    const div = metadata.appendChild(document.createElement("div"));
-    div.className = "stats-list";
-    for (const [i, s] of ctx.runtime_stats.entries()) {
-      const p = div.appendChild(document.createElement("p"));
-      if (ctx.runtime_stats.length > 1) p.innerText = `Run ${i+1}/${ctx.runtime_stats.length}`;
-      const table = div.appendChild(document.createElement("table"));
-      const tbody = table.appendChild(document.createElement("tbody"));
-      for (const { name, value, unit, subunits } of s.data) {
-          const mainRow = appendRow(tbody, name, value, unit, "main-row");
-          if (!subunits?.length) continue;
-          const subunitRow = tbody.appendChild(document.createElement("tr"));
-          subunitRow.style.display = "none";
-          mainRow.onclick = () => subunitRow.style.display = subunitRow.style.display === "none" ? "table-row" : "none";
-          mainRow.style.cursor = "pointer";
-          const td = subunitRow.appendChild(document.createElement("td"));
-          td.colSpan = 2;
-          const table = td.appendChild(document.createElement("table"));
-          for (const u of subunits) appendRow(table, u.name, u.value, unit, "sub-row");
-      }
-    }
-  }
   // ** rewrite steps
   if (step.match_count >= 1) {
     const rewriteList = metadata.appendChild(document.createElement("div"));

tinygrad/viz/serve.py

@@ -36,8 +36,6 @@ def get_metadata(trace_bufs:list[tuple]) -> list[dict]:
       steps = [{"name":s.name, "loc":s.loc, "depth":s.depth, "match_count":len(s.matches), "code_line":printable(s.loc),
                 "query":f"/ctxs?ctx={i}&idx={j}"} for j,s in enumerate(v)]
       ret.append(r:={"name":k.display_name, "steps":steps})
-      # use the first key to get runtime profiling data about this context
-      if getenv("PROFILE_VALUE") >= 2 and k.keys: r["runtime_stats"] = get_runtime_stats(k.keys[0])
       # program spec metadata
       if isinstance(k.ret, ProgramSpec):
         steps.append({"name":"View Disassembly", "query":f"/disasm?ctx={i}"})
@@ -201,13 +199,6 @@ def get_profile(profile:list[ProfileEvent]) -> bytes|None:
   index = json.dumps({"strings":list(scache), "dtypeSize":dtype_size, "markers":[{"ts":int(e.ts-start_ts), **e.arg} for e in markers]}).encode()
   return struct.pack("<IQII", unwrap(end_ts)-start_ts, max(peaks,default=0), len(index), len(ret))+index+b"".join(ret)
 
-def get_runtime_stats(key) -> list[dict]:
-  ret:list[dict] = []
-  for e in profile:
-    if isinstance(e, ProfileRangeEvent) and e.en is not None and e.name == key:
-      ret.append({"device":e.device, "data":[{"name":"Duration", "value":float(e.en-e.st), "unit":"us"}]})
-  return ret
-
 # ** Assembly analyzers
 
 def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict: