remove noops

* use recursive_property to prevent RecursionError

* not slower

* fix tests

* faster

* simpler

.github/workflows/benchmark.yml

@@ -109,7 +109,7 @@ jobs:
     - name: Train MNIST
       run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
     - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=320 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=330 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
     - name: Run 10 CIFAR training steps w HALF
       run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=385 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
     #- name: Run 10 CIFAR training steps w BF16

.github/workflows/test.yml

@@ -270,9 +270,9 @@ jobs:
     - name: Run targetted tests on NULL backend
       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
+      run: 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
+      run: 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
@@ -452,10 +452,12 @@ jobs:
         run: CL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
       - name: Test MLPerf stuff
         run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
+      - name: NULL=1 beautiful_mnist_multigpu
+        run: NULL=1 python examples/beautiful_mnist_multigpu.py
       - name: Test Bert training
-        run: MAX_BUFFER_SIZE=0 NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
       - name: Test llama 3 training
-        run: MAX_BUFFER_SIZE=0 NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 

examples/mlperf/dataloader.py

@@ -511,6 +511,33 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
       # happens with BENCHMARK set
       pass
 
+# stable diffusion callbacks to match mlperf ref; declared here because they're pickled
+def filter_dataset(sample:dict): return {k:v for k,v in sample.items() if k in {'npy', 'txt'}}
+def collate(batch:list[dict]):
+  ret = {"npy": [], "txt": [], "__key__": []}
+  for sample in batch:
+    for k,v in sample.items():
+      ret[k].append(v)
+  return ret
+def collate_fn(batch): return batch
+
+# Reference (code): https://github.com/mlcommons/training/blob/2f4a93fb4888180755a8ef55f4b977ef8f60a89e/stable_diffusion/ldm/data/webdatasets.py, Line 55
+# Reference (params): https://github.com/mlcommons/training/blob/ab4ae1ca718d7fe62c369710a316dff18768d04b/stable_diffusion/configs/train_01x08x08.yaml, Line 107
+def batch_load_train_stable_diffusion(urls:str, BS:int):
+  import webdataset
+  dataset = webdataset.WebDataset(urls=urls, resampled=True, cache_size=-1, cache_dir=None)
+  dataset = dataset.shuffle(size=1000)
+  dataset = dataset.decode()
+  dataset = dataset.map(filter_dataset)
+  dataset = dataset.batched(BS, partial=False, collation_fn=collate)
+  dataset = webdataset.WebLoader(dataset, batch_size=None, shuffle=False, num_workers=1, persistent_workers=True, collate_fn=collate_fn)
+
+  for x in dataset:
+    assert isinstance(x, dict) and all(isinstance(k, str) for k in x.keys()) and all(isinstance(v, list) for v in x.values())
+    assert all(isinstance(moment_mean_logvar, np.ndarray) and moment_mean_logvar.shape==(1,8,64,64) for moment_mean_logvar in x["npy"])
+    assert all(isinstance(caption, str) for caption in x["txt"])
+    yield x
+
 # llama3
 
 class BinIdxDataset:

examples/mlperf/model_train.py

@@ -1493,6 +1493,144 @@ def train_llama3():
           safe_save(get_state_dict(model), fn)
         break
 
+def train_stable_diffusion():
+  from extra.models.unet import UNetModel
+  from examples.mlperf.dataloader import batch_load_train_stable_diffusion
+  from examples.mlperf.lr_schedulers import LambdaLR, LambdaLinearScheduler
+  from examples.mlperf.initializers import init_stable_diffusion
+  from examples.mlperf.helpers import get_training_state
+  import numpy as np
+
+  config = {}
+  GPUS               = config["GPUS"]                   = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
+  seed               = config["seed"]                   = getenv("SEED", 12345)
+  # ** hyperparameters **
+  BS                 = config["BS"]                     = getenv("BS", 1 * len(GPUS))
+  BASE_LR            = config["LEARNING_RATE"]          = getenv("LEARNING_RATE", 2.5e-7)
+  # https://github.com/mlcommons/training_policies/blob/cfa99da479b8d5931f7a3c67612d021dfb47510a/training_rules.adoc#benchmark_specific_rules
+  # "Checkpoint must be collected every 512,000 images. CEIL(512000 / global_batch_size) if 512000 is not divisible by GBS."
+  # NOTE: It's inferred that "steps" is the unit for the output of the CEIL formula, based on all other cases of CEIL in the rules
+  CKPT_STEP_INTERVAL = config["CKPT_STEP_INTERVAL"]     = getenv("CKPT_STEP_INTERVAL", math.ceil(512_000 / BS))
+  CKPTDIR            = config["CKPTDIR"]                = Path(getenv("CKPTDIR", "./checkpoints"))
+  DATADIR            = config["DATADIR"]                = Path(getenv("DATADIR", "./datasets"))
+  UNET_CKPTDIR       = config["UNET_CKPTDIR"]           = Path(getenv("UNET_CKPTDIR", "./checkpoints"))
+  TOTAL_CKPTS        = config["TOTAL_CKPTS"]            = getenv("TOTAL_CKPTS", 0)
+
+  print(f"training on {GPUS}")
+  lr = BS * BASE_LR
+  print(f"BS={BS}, BASE_LR={BASE_LR}, lr={lr}")
+  print(f"CKPT_STEP_INTERVAL = {CKPT_STEP_INTERVAL}")
+  for x in GPUS: Device[x]
+  if (WANDB := getenv("WANDB", "")):
+    import wandb
+    wandb.init(config=config, project="MLPerf-Stable-Diffusion")
+
+  Tensor.manual_seed(seed)  # seed for weight initialization
+  model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-train", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)
+
+  optimizer = AdamW(get_parameters(unet))
+  lambda_lr_callback = LambdaLinearScheduler(1000, 1.0, 1.0, 1e-06, 10000000000000).schedule
+  lr_scheduler = LambdaLR(optimizer, Tensor(lr, dtype=dtypes.float, device=optimizer.device), lambda_lr_callback)
+
+  @TinyJit
+  def train_step(mean:Tensor, logvar:Tensor, tokens:Tensor, unet:UNetModel, optimizer:LAMB, lr_scheduler:LambdaLR) -> Tensor:
+    optimizer.zero_grad()
+
+    timestep = Tensor.randint(BS, low=0, high=model.alphas_cumprod.shape[0], dtype=dtypes.int, device=GPUS[0])
+    latent_randn = Tensor.randn(*mean.shape, device=GPUS[0])
+    noise = Tensor.randn(*mean.shape, device=GPUS[0])
+    for t in (mean, logvar, tokens, timestep, latent_randn, noise):
+      t.shard_(GPUS, axis=0)
+
+    std = Tensor.exp(0.5 * logvar.clamp(-30.0, 20.0))
+    latent = (mean + std * latent_randn) * 0.18215
+
+    sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[timestep].reshape(timestep.shape[0], 1, 1, 1)
+    sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[timestep].reshape(timestep.shape[0], 1, 1, 1)
+    latent_with_noise = sqrt_alphas_cumprod_t * latent + sqrt_one_minus_alphas_cumprod_t * noise
+    v_true = sqrt_alphas_cumprod_t * noise - sqrt_one_minus_alphas_cumprod_t * latent
+
+    context = model.cond_stage_model.embed_tokens(tokens)
+
+    out = unet(latent_with_noise, timestep, context)
+    loss = ((out - v_true) ** 2).mean()
+    del mean, logvar, std, latent, noise, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t
+    del out, v_true, context, latent_randn, tokens, timestep
+    loss.backward()
+
+    optimizer.step()
+    lr_scheduler.step()
+    loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
+    Tensor.realize(loss, out_lr)
+    return loss, out_lr
+    
+  # checkpointing takes ~9 minutes without this, and ~1 minute with this
+  @TinyJit
+  def ckpt_to_cpu():
+    ckpt = get_training_state(unet, optimizer, lr_scheduler)
+    # move to CPU first so more GPU bufs aren't created (can trigger OOM)
+    for k,v in ckpt.items(): ckpt[k] = v.detach().to("CPU")
+    Tensor.realize(*[v for v in ckpt.values()])
+    for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype.base).contiguous()
+    Tensor.realize(*[v for v in ckpt.values()])
+    return ckpt
+
+  # training loop
+  dl = batch_load_train_stable_diffusion(f'{DATADIR}/laion-400m/webdataset-moments-filtered/{{00000..00831}}.tar', BS)
+  # for tests
+  saved_checkpoints = []
+
+  train_start_time = time.perf_counter()
+  t0 = t6 = time.perf_counter()
+  for i, batch in enumerate(dl, start=1):
+    loop_time = time.perf_counter() - t0
+    t0 = time.perf_counter()
+    dl_time = t0 - t6
+    GlobalCounters.reset()
+
+    mean, logvar = np.split(np.concatenate(batch["npy"], axis=0), 2, axis=1)
+    mean, logvar = Tensor(mean, dtype=dtypes.float32, device="CPU"), Tensor(logvar, dtype=dtypes.float32, device="CPU")
+    tokens = []
+    for text in batch['txt']: tokens += model.cond_stage_model.tokenizer.encode(text, pad_with_zeros=True)
+    tokens = Tensor(tokens, dtype=dtypes.int32, device="CPU").reshape(-1, 77)
+
+    t1 = time.perf_counter()
+    loss, lr = train_step(mean, logvar, tokens, unet, optimizer, lr_scheduler)
+    loss_item, lr_item = loss.item(), lr.item()
+    t2 = time.perf_counter()
+
+    if i == 3:
+      for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
+      print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
+      
+    total_train_time = time.perf_counter() - train_start_time
+    if WANDB:
+      wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,
+                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (t2-t1), "train/input_prep_time": t1-t0,
+                    "train/train_step_time": t2-t1, "train/total_time": total_train_time})
+
+      if i == 1 and wandb.run is not None:
+        with open(f"{UNET_CKPTDIR}/wandb_run_id_{wandb.run.id}", "w") as f:
+          f.write(f"wandb.run.id = {wandb.run.id}")
+
+    if i % CKPT_STEP_INTERVAL == 0:
+      # https://github.com/mlcommons/training_policies/blob/cfa99da479b8d5931f7a3c67612d021dfb47510a/training_rules.adoc#benchmark_specific_rules
+      # "evaluation is done offline, the time is not counted towards the submission time."
+      fn = f"{UNET_CKPTDIR}/{i}.safetensors"
+      print(f"saving unet checkpoint at {fn}")
+      saved_checkpoints.append(fn)
+      safe_save({k.replace("model.", ""):v for k,v in ckpt_to_cpu().items() if k.startswith("model.")}, fn)
+      if TOTAL_CKPTS and i == TOTAL_CKPTS * CKPT_STEP_INTERVAL:
+        print(f"ending run after {i} steps ({TOTAL_CKPTS} checkpoints collected)")
+        return saved_checkpoints
+
+    t3 = time.perf_counter()
+    print(f"""step {i}: {GlobalCounters.global_ops * 1e-9 / (t2-t1):9.2f} GFLOPS, mem_used: {GlobalCounters.mem_used / 1e9:.2f} GB,
+  loop_time_prev: {loop_time:.2f}, dl_time: {dl_time:.2f}, input_prep_time: {t1-t0:.2f}, train_step_time: {t2-t1:.2f},
+  t3-t2: {t3-t2:.4f}, loss:{loss_item:.5f}, lr:{lr_item:.3e}, total_train_time:{total_train_time:.2f}
+  """)
+    t6 = time.perf_counter()
+
 if __name__ == "__main__":
   multiprocessing.set_start_method('spawn')
 
@@ -1501,7 +1639,7 @@ if __name__ == "__main__":
   else: bench_log_manager = contextlib.nullcontext()
 
   with Tensor.train():
-    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn").split(","):
+    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn,stable_diffusion").split(","):
       nm = f"train_{m}"
       if nm in globals():
         print(f"training {m}")

examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/stable_diffusion/implementations/tinybox_8xMI300X/dev_run.sh

@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+
+DATETIME=${2:-$(date "+%m%d%H%M")}
+LOGFILE="${HOME}/logs/sd_mi300x_${DATETIME}.log"
+# UNET_CKPTDIR must be set: training saves checkpoints to this path, then a separate eval process scans this path to know which checkpoints to eval
+export UNET_CKPTDIR="${HOME}/stable_diffusion/training_checkpoints/${DATETIME}"
+mkdir -p "${HOME}/logs" "$UNET_CKPTDIR"
+
+# run this script in isolation when using the --bg flag
+if [[ "${1:-}" == "--bg" ]]; then
+  echo "logging output to $LOGFILE"
+  echo "saving UNet checkpoints to $UNET_CKPTDIR"
+  script_path="$(readlink -f "${BASH_SOURCE[0]}")"
+  nohup bash "$script_path" run "$DATETIME" >"$LOGFILE" 2>&1 & disown $!
+  exit 0
+fi
+
+# venv management
+if [[ -d .venv-sd-mlperf ]]; then
+  . .venv-sd-mlperf/bin/activate
+else
+  python3 -m venv .venv-sd-mlperf && . .venv-sd-mlperf/bin/activate
+  pip install --index-url https://download.pytorch.org/whl/cpu torch && pip install tqdm numpy ftfy regex pillow scipy wandb webdataset
+fi
+pip list
+apt list --installed | grep amdgpu
+rocm-smi --version
+modinfo amdgpu | grep version
+
+export BEAM=2 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 IGNORE_JIT_FIRST_BEAM=1 HCQDEV_WAIT_TIMEOUT_MS=300000
+export AMD_LLVM=0 # bf16 seems to require this
+export DATADIR="/raid/datasets/stable_diffusion"
+export CKPTDIR="/raid/weights/stable_diffusion"
+export EVAL_CKPT_DIR=$UNET_CKPTDIR
+export MODEL="stable_diffusion" PYTHONPATH="."
+export GPUS=8 BS=304
+export CONTEXT_BS=816 DENOISE_BS=600 DECODE_BS=384 INCEPTION_BS=560 CLIP_BS=240
+export WANDB=1
+export PARALLEL=4
+export PYTHONUNBUFFERED=1
+sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setperfdeterminism 1500 || exit 1
+
+# Retry BEAM search if script fails before BEAM COMPLETE is printed, but don't retry after that
+run_retry(){ local try=0 max=5 code tmp py pgid kids
+  while :; do
+    tmp=$(mktemp)
+    setsid bash -c 'exec env "$@"' _ "$@" > >(tee -a "$LOGFILE" | tee "$tmp") 2>&1 &
+    py=$!; pgid=$(ps -o pgid= -p "$py" | tr -d ' ')
+    wait "$py"; code=$?
+    [[ -n "$pgid" ]] && { kill -TERM -"$pgid" 2>/dev/null; sleep 1; kill -KILL -"$pgid" 2>/dev/null; }
+    kids=$(pgrep -P "$py" || true)
+    while [[ -n "$kids" ]]; do
+      kill -TERM $kids 2>/dev/null; sleep 0.5
+      kids=$(for k in $kids; do pgrep -P "$k" || true; done)
+    done
+    grep -q 'BEAM COMPLETE' "$tmp" && { rm -f "$tmp"; return 1; }
+    rm -f "$tmp"
+    ((code==0)) && return 0
+    ((try>=max)) && return 2
+    ((try++)); sleep 90; echo "try = ${try}"
+  done
+}
+
+# Power limiting to 400W is only needed if GPUs fall out of sync (causing 2.2x increased train time) at higher power, which has been observed at 450W
+sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setpoweroverdrive 750 && \
+run_retry TOTAL_CKPTS=7 python3 examples/mlperf/model_train.py; (( $? == 2 )) && { echo "training failed before BEAM completion"; exit 2; }
+sleep 90
+
+run_retry EVAL_SAMPLES=600 python3 examples/mlperf/model_eval.py; (( $? == 2 )) && { echo "eval failed before BEAM completion"; exit 2; }
+# Checkpoints will be evaluated in reverse chronological order, even if above training crashed early
+# STOP_IF_CONVERGED=1: Stop the eval after the first time convergence is detected; no more checkpoints will be evaluated after that.
+STOP_IF_CONVERGED=1 python3 examples/mlperf/model_eval.py

pytest.ini

@@ -1,6 +1,6 @@
 [pytest]
 norecursedirs = extra
-timeout = 180
+timeout = 240
 timeout_method = thread
 timeout_func_only = true
 testpaths = test

test/external/external_test_opt.py

@@ -221,7 +221,7 @@ class TestOpt(unittest.TestCase):
       for axis in [0, 1]:
         for n in [4, 8, 16]:
           b = torch.ones(n, n).sum(axis).reshape(n, 1).expand(n, n).sum(axis)
-          with CLCache(allowed=2):
+          with CLCache(allowed=3 if RANGEIFY else 2):
             a = Tensor.ones(n, n).contiguous().sum(axis).reshape(n, 1).expand(n, n).sum(axis)
             a.realize()
           np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
@@ -231,7 +231,7 @@ class TestOpt(unittest.TestCase):
       axis1, axis2 = 0, 1
       for n in [4, 8, 16]:
         b = torch.ones(n, n).sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
-        with CLCache(allowed=2):
+        with CLCache(allowed=3 if RANGEIFY else 2):
           a = Tensor.ones(n, n).contiguous().sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
           a.realize()
         np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)

test/external/mlperf_stable_diffusion/external_test_train.py

@@ -0,0 +1,23 @@
+import unittest, os
+from tempfile import TemporaryDirectory
+from tinygrad import Tensor
+from tinygrad.helpers import getenv
+from examples.mlperf.model_train import train_stable_diffusion
+
+class TestTrain(unittest.TestCase):
+  def test_train_to_ckpt(self):
+    # train for num_steps, save checkpoint, and stop training
+    num_steps = 42
+    os.environ.update({"MODEL": "stable_diffusion", "TOTAL_CKPTS": "1", "CKPT_STEP_INTERVAL": str(num_steps), "GPUS": "8", "BS": "304"})
+    # NOTE: update these based on where data/checkpoints are on your system
+    if not getenv("DATADIR", ""): os.environ["DATADIR"] = "/raid/datasets/stable_diffusion"
+    if not getenv("CKPTDIR", ""): os.environ["CKPTDIR"] = "/raid/weights/stable_diffusion"
+    with TemporaryDirectory(prefix="test-train") as tmp:
+      os.environ["UNET_CKPTDIR"] = tmp
+      with Tensor.train():
+        saved_ckpts = train_stable_diffusion()
+      expected_ckpt = f"{tmp}/{num_steps}.safetensors"
+      assert len(saved_ckpts) == 1 and saved_ckpts[0] == expected_ckpt
+
+if __name__=="__main__":
+  unittest.main()

test/helpers.py

@@ -1,4 +1,4 @@
-import time, struct
+import time, struct, unittest
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
@@ -7,7 +7,7 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.engine.realize import Runner
 from tinygrad.dtype import DType
 from tinygrad.nn.state import get_parameters
-from tinygrad.helpers import T, CI
+from tinygrad.helpers import T, CI, RANGEIFY
 from tinygrad.codegen import full_rewrite
 from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
 
@@ -62,3 +62,6 @@ def not_support_multi_device():
 
 # NOTE: This will open REMOTE if it's the default device
 REAL_DEV = (Device.DEFAULT if Device.DEFAULT != "REMOTE" else Device['REMOTE'].properties.real_device)
+
+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)

test/models/test_efficientnet.py

@@ -101,7 +101,8 @@ class TestResNet(unittest.TestCase):
 
   def test_chicken(self):
     labels = _infer(self.model, chicken_img)
-    self.assertEqual(_LABELS[labels[0]], "hen")
+    # NOTE: logits for these two are close
+    self.assertIn(_LABELS[labels[0]], ("hen", "cock"))
 
   def test_car(self):
     labels = _infer(self.model, car_img)

test/models/test_real_world.py

@@ -94,7 +94,7 @@ class TestRealWorld(unittest.TestCase):
     @TinyJit
     def test(t, v):
       with Context(JIT=0): return model(t, v).realize()
-    helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23 if CI else 0.9, 137 if CI else 396, all_jitted=True)
+    helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23 if CI else 0.9, 160 if CI else 396, all_jitted=True)
 
   @unittest.skipIf(CI and Device.DEFAULT == "CPU", "slow")
   def test_train_mnist(self):
@@ -112,7 +112,7 @@ class TestRealWorld(unittest.TestCase):
         loss.backward()
         optimizer.step()
 
-      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 93)
+      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 102)
 
   @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
   def test_forward_cifar(self):
@@ -176,7 +176,7 @@ class TestRealWorld(unittest.TestCase):
       for v in data.values(): v.to_(Device.DEFAULT)
 
       helper_test("train_bert", lambda: (data["input_ids"], data["segment_ids"], data["input_mask"], data["masked_lm_positions"], \
-          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.25, 347)
+          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.28, 357)
 
 if __name__ == '__main__':
   unittest.main()

test/test_assign.py

@@ -119,6 +119,22 @@ class TestAssign(unittest.TestCase):
     new = a + old_a
     np.testing.assert_allclose(new.numpy(), 4)
 
+  def test_assign_changes_alt(self, realize=False):
+    a = Tensor(1).contiguous()
+    if realize: a.realize()
+    b = a.contiguous()    # b returns a new Tensor
+    b.assign(2)
+    b.realize()
+    self.assertNotEqual(a.item(), b.item())
+  # on a realized Tensor contiguous child changes the source
+  @unittest.expectedFailure
+  def test_assign_changes_realized_alt(self): return self.test_assign_changes_alt(realize=True)
+
+  def test_assign_changes_buffer_alt(self):
+    a, b = [Tensor(Tensor(0).contiguous().realize().uop.as_buf()) for _ in range(2)]
+    Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2))
+    self.assertEqual((a + b).item(), 3)
+
   def test_assign_diamond_cycle(self):
     # NOTE: should *not* raise AssertionError from numpy
     with self.assertRaisesRegex(RuntimeError, "cycle"):

test/test_dtype.py

@@ -7,7 +7,7 @@ from tinygrad.helpers import getenv, DEBUG, CI
 from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad import Device, Tensor, dtypes
-from hypothesis import assume, given, settings, strategies as strat
+from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
 import pytest
@@ -52,8 +52,6 @@ def _test_cast(a:Tensor, target_dtype:DType):
   if target_dtype in dtypes.fp8s: expected = list(map(lambda x: truncate[target_dtype](x), expected))
   _test_op(lambda: a.cast(target_dtype), target_dtype, expected)
 def _test_bitcast(a:Tensor, target_dtype:DType, target=None):
-  if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and a.dtype == dtypes.int8 and target_dtype.itemsize != a.dtype.itemsize:
-    raise unittest.SkipTest("shape changing bitcast of int8 broken on PTX")
   expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype)).tolist()
   if target_dtype in dtypes.fp8s: expected = list(map(lambda x: fp8_to_float(x, target_dtype), expected))
   _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected)
@@ -294,7 +292,6 @@ class TestInt8DType(TestDType):
   def test_int8_to_uint16_negative(self):
     _test_op(lambda: Tensor([-1, -2, -3, -4], dtype=dtypes.int8).cast(dtypes.uint16), dtypes.uint16, [2**16-1, 2**16-2, 2**16-3, 2**16-4])
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken in ptx")
   def test_bitcast_alt(self):
     a = Tensor([72, -90, 27, 40, -53, 70, 96, 51], dtype=dtypes.int8).bitcast(dtypes.short)
     self.assertListEqual(a.tolist(), [-22968, 10267, 18123, 13152])
@@ -308,8 +305,6 @@ class TestUint8DType(TestDType):
 class TestBitCast(unittest.TestCase):
   @given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
   def test_shape_change_bitcast(self, dt1, dt2):
-    # NOTE: this has to be assume to prevent hypothesis from skipping all samples
-    assume(not (isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and dt1 == dtypes.int8)) # TODO: bitcasting int8 fails in PTX
     data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
     expected = torch.tensor(data.tolist(), dtype=_to_torch_storage_type(dt1)).view(_to_torch_dtype(dt2))
     if dt2 in dtypes.fp8s:

test/test_image_dtype.py

@@ -139,7 +139,7 @@ class TestImageDType(unittest.TestCase):
       # NOTE: the w1 grad must realize to a seperate kernel
       assert w1.grad.uop.is_realized, f"never realized {w1.grad}"
       self.assertEqual(w1.grad.uop.base.buffer.dtype, dtypes.float32)
-      self.assertEqual(len(sched), 8 if RANGEIFY else 10)
+      self.assertEqual(len(sched), 9 if RANGEIFY else 10)
 
 @unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported")
 class TestImageRealization(unittest.TestCase):

test/test_linearizer.py

@@ -10,7 +10,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
-from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT
+from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, RANGEIFY
 from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
 from tinygrad.codegen import apply_rewrites, rewrites_for_views
 from tinygrad.renderer.ptx import PTXRenderer
@@ -335,6 +335,7 @@ class TestLinearizer(unittest.TestCase):
     a.realize()
     np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
 
+  @unittest.skipIf(RANGEIFY and isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexes differently. might be ok?")
   def test_where_fold(self):
     a = Tensor.ones(4, 4).contiguous().realize()
     b = a.shrink(((1, 2), None)).pad(((1, 2), None))

test/test_multitensor.py

@@ -7,7 +7,7 @@ from tinygrad.nn.state import get_parameters, get_state_dict
 from tinygrad.engine.realize import lower_schedule, BufferCopy, CompiledRunner, run_schedule
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import REAL_DEV, not_support_multi_device
+from test.helpers import REAL_DEV, not_support_multi_device, expect_rangeify_fails
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -201,6 +201,14 @@ class TestMultiTensor(unittest.TestCase):
     fn = f(n)
     np.testing.assert_allclose(fX.numpy(), fn, rtol=1e-6, atol=1e-6)
 
+  @expect_rangeify_fails # TODO: fix
+  def test_allreduce_shard_ring_sum(self):
+    for axis in (0, 1, None):
+      for use_ring in (0, 2):
+        t = Tensor([1, 2, 3, 4]).reshape(2, 2)
+        with Context(RING=use_ring):
+          np.testing.assert_equal(t.shard(devices_2, axis=axis).sum().item(), 10)
+
   def test_allreduce_naive(self):
     with Context(RING=0):
       a,b = _test_allreduce(Tensor.rand(256, 256))

test/test_nn.py

@@ -333,8 +333,8 @@ class TestNN(unittest.TestCase):
 
       np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
       np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
-      # TODO: is this numerical issue or a bug?
-      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=4e-3, rtol=1e-3)
+      # TODO: is this numerical issue or a bug? RANGEIFY big reduce kernel amplifies numerical issue
+      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=1e-2, rtol=1e-3)
       np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
 
   def test_rmsnorm(self):

test/test_ops.py

@@ -1313,7 +1313,7 @@ class TestOps(unittest.TestCase):
   @unittest.skipIf(CI and Device.DEFAULT in ["NV", "CL", "CUDA"] or (Device.DEFAULT == "CPU" and CPU_LLVM) or IMAGE
   or (Device.DEFAULT == "WEBGPU" and platform.system() == "Windows"), "not supported on these in CI/IMAGE")
   def test_gemm_fp16(self):
-    helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3)
+    helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3, grad_atol=5e-3, grad_rtol=5e-3)
   def test_gemm(self):
     helper_test_op([(64,64), (64,64)], lambda x,y: x.matmul(y))
   @slow_test
@@ -3164,8 +3164,8 @@ 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(RANGEIFY and ((getenv("MOCKGPU") and Device.DEFAULT == "AMD") or Device.DEFAULT == "PYTHON"),
-                   "very slow on MOCKGPU because reduce does not fold")
+  @unittest.skipIf(RANGEIFY and (getenv("MOCKGPU") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
+  @unittest.skipIf(RANGEIFY and Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
   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_schedule.py

@@ -18,6 +18,7 @@ from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context,
 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
+from test.helpers import expect_rangeify_fails, expect_nonrangeify_fails
 
 class KernelCountException(Exception): pass
 def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
@@ -42,9 +43,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
     raise KernelCountException(f"{kernel_cnt} != {allowed}")
   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()
 
@@ -1925,6 +1923,15 @@ class TestSchedule(unittest.TestCase):
     run_schedule(check_schedule(loss, 4))
     np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6)
 
+  def test_const_folding_alt(self):
+    t = Tensor.full((2,), 1.)
+    lt = (t < 0.)
+    a = Tensor.empty(2).assign(t*lt.where(-1., 0.))
+    b = Tensor.empty(2, dtype=dtypes.bool).assign(lt)
+    Tensor.realize(a, b)
+    self.assertEqual(a.tolist(), [0., 0.])
+    self.assertEqual(b.tolist(), [False, False])
+
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Validation error on WebGPU")
   def test_mnist_val(self):
     from tinygrad.nn.datasets import mnist
@@ -2218,6 +2225,11 @@ class TestCopyFolding(unittest.TestCase):
     check_schedule(b, 0, filter_sink=False)
     assert b.item() == 1
 
+  def test_const_copy_multi(self):
+    x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"])
+    check_schedule(x, 0, filter_sink=False)
+    self.assertEqual(x.item(), 1)
+
   def test_late_const_copy_folding(self):
     a = Tensor.arange(3).realize()
     zeros = Tensor.zeros(3).realize()

test/test_tiny.py

@@ -15,6 +15,10 @@ class TestTiny(unittest.TestCase):
     out = Tensor([1.,2,3])
     self.assertListEqual(out.tolist(), [1.0, 2.0, 3.0])
 
+  def test_elu(self):
+    out = Tensor([[1.,2],[3,4]]).sum(axis=1).elu()
+    self.assertListEqual(out.tolist(), [3.0, 7.0])
+
   def test_plus(self):
     out = Tensor([1.,2,3]) + Tensor([4.,5,6])
     self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0])

test/test_uop_graph.py

@@ -417,6 +417,11 @@ class TestUOpGraph(unittest.TestCase):
       uops = to_uops_list([v.bitcast(dt)])
       self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST]), 0, f"dtype = {dt}")
 
+  def test_sub_with_cast_folds(self):
+    a = Variable("a", 0, 5)
+    uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
+    assert uops == [UOp.const(dtypes.int, 0)]
+
   def test_where_on_gated_load_fold(self):
     ridx0 = UOp.range(100, 0)
     d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
@@ -461,6 +466,8 @@ class TestUOpGraph(unittest.TestCase):
       if u.op is Ops.STORE: assert u.src[1].arg==5
 
   def test_load_idx_becomes_int(self):
+    # These loads wont overflow int since we know from the gate that the value is bounded
+    r0 = UOp.range(10, 0)
     d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
     d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 1)
     l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),)).cast(dtypes.index)
@@ -471,6 +478,12 @@ class TestUOpGraph(unittest.TestCase):
     for u in uops:
       if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
 
+    valid = (10*r0<5-l0).ne(True)&(l0<3000)
+    l2 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx.valid(valid)),))
+    uops = to_uops_list([l2])
+    for u in uops:
+      if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
+
   def test_in_out_of_bounds_access(self):
     with Context(IGNORE_OOB=0):
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
@@ -599,12 +612,13 @@ class TestUOpGraph(unittest.TestCase):
       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])
+    with Context(IGNORE_OOB=0):
+      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)

test/unit/test_rewrite_map.py

@@ -28,22 +28,6 @@ class TestRewriteMap(unittest.TestCase):
     self.assertIs(sub_map[a+b], e)
     self.assertIs(sub_map[(a+b)*c], f)
 
-  def test_multistage_substitute(self):
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    c = UOp.variable('c', 0, 10)
-    d = UOp.variable('d', 0, 10)
-    sub1 = {a+b:c}
-    start = (a+b)*c
-    # stage 1: (a+b)*c -> c*c
-    sub_map1 = graph_rewrite_map(start, _substitute, sub1, bottom_up=True)
-    self.assertIs(sub_map1[(a+b)*c], c*c)
-    # stage 2: c*c -> d
-    sub2 = {c*c:d}
-    sub_map2 = graph_rewrite_map(sub_map1[start], _substitute, sub2, input_map=sub_map1, bottom_up=True)
-    # (a+b)*c -> c*c -> d
-    self.assertIs(sub_map2[(a+b)*c], d)
-
   def test_add_zero(self):
     # Build a small graph: add(0, add(const=0, const=5))
     zero_node = UOp.const(dtypes.index, 0)
@@ -144,11 +128,11 @@ class TestRewriteMap(unittest.TestCase):
       yz_sum_zero  = yz_sum + zero_node   -> rewrites to yz_sum
       yz_neg       = -yz_sum_zero        -> -(y+z)
       yz_dneg      = -yz_neg             -> y+z    (double neg gone)
-      x_plus_yz    = x_var + yz_dneg     -> x + (y+z)
-      double_neg_x = -(-x_plus_yz)       -> x + (y+z)
-      final_expr   = double_neg_x * one_node -> x + (y+z)
+      x_plus_yz    = x_var + yz_dneg     -> (x+y)+z  (add nodes get sorted)
+      double_neg_x = -(-x_plus_yz)       -> (x+y)+z
+      final_expr   = double_neg_x * one_node -> (x+y)+z
 
-    We expect the final result to be (x + (y+z)).
+    We expect the final result to be ((x+y)+z).
     Each original node should map to the final node that replaces it,
     which might be structurally equivalent but not the same reference.
     """
@@ -163,9 +147,9 @@ class TestRewriteMap(unittest.TestCase):
     yz_sum_zero = yz_sum + zero_node    # (y + z) + 0
     yz_neg = -yz_sum_zero               # -(y+z)
     yz_dneg = -yz_neg                   # -(-(y+z)) -> (y+z)
-    x_plus_yz = x_var + yz_dneg         # x + (y+z)
-    double_neg_x = -(-x_plus_yz)        # neg(neg(x+(y+z))) -> x+(y+z)
-    final_expr = double_neg_x * one_node  # (x+(y+z)) * 1 -> x+(y+z)
+    x_plus_yz = x_var + yz_dneg         # x + (y+z) -> (x+y)+z
+    double_neg_x = -(-x_plus_yz)        # neg(neg(x+(y+z))) -> (x+y)+z
+    final_expr = double_neg_x * one_node  # ((x+y)+z) * 1 -> (x+y)+z
 
     node_map = graph_rewrite_map(final_expr, symbolic)
 
@@ -182,14 +166,15 @@ class TestRewriteMap(unittest.TestCase):
     # -(-(y+z)) => (y+z)
     self.assertEqual(node_map[yz_dneg], yz_sum)
 
-    # x + (y+z) => might get recreated if yz_dneg was changed, so compare to x + yz_sum
-    self.assertEqual(node_map[x_plus_yz], x_var + yz_sum)
+    # x + (y+z) => (x+y)+z
+    expected_xyz = (x_var + y_var) + z_var
+    self.assertEqual(node_map[x_plus_yz], expected_xyz)
 
-    # -(-(x+(y+z))) => x + (y+z)
-    self.assertEqual(node_map[double_neg_x], x_var + yz_sum)
+    # -(-(x+(y+z))) => (x+y)+z
+    self.assertEqual(node_map[double_neg_x], expected_xyz)
 
-    # (x+(y+z)) * 1 => x+(y+z)
-    self.assertEqual(node_map[final_expr], x_var + yz_sum)
+    # ((x+y)+z) * 1 => (x+y)+z
+    self.assertEqual(node_map[final_expr], expected_xyz)
 
     # Unchanged atomic nodes map to themselves
     self.assertEqual(node_map[x_var], x_var)

test/unit/test_shm_tensor.py

@@ -1,10 +1,11 @@
 import unittest
 import multiprocessing.shared_memory as shared_memory
-from tinygrad.helpers import CI
+from tinygrad.helpers import CI, WIN, RANGEIFY
 from tinygrad.tensor import Tensor, Device
 import numpy as np
 
 class TestRawShmBuffer(unittest.TestCase):
+  @unittest.skipIf(WIN and CI and RANGEIFY, "only fails with RANGEIFY on CI windows instance")
   def test_e2e(self):
     t = Tensor.randn(2, 2, 2).realize()
 

test/unit/test_simplify_valid_idx.py

@@ -60,7 +60,7 @@ class TestValidIdxSimplification(unittest.TestCase):
     load = get_gated_load_uop(gate, idx)
     self.check(load,
       "0",
-      "(((lidx0+(gidx0*4))<19)!=True)")
+      "((((gidx0*4)+lidx0)<19)!=True)")
 
   def test_simplify_within_valid1(self):
     ridx0 = Range(0, 4)
@@ -184,7 +184,6 @@ class TestValidIdxSimplification(unittest.TestCase):
     print("The expressions are not equivalent.")
     print(s.model())
 
-  @unittest.expectedFailure  # TODO: improve uop_given_valid
   def test_valid_becomes_const2(self):
     ridx0 = Range(0, 4)
     ridx1 = Range(1, 4)
@@ -304,7 +303,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)+(r2*6))+r0)+-6)", "(((idx2*2)+r1)+-1)")
+    self.check(load, None, "((((idx1*48)+r0)+(r2*6))+-6)", "(((idx2*2)+r1)+-1)")
 
   def test_openpilot_conv2(self):
     # conv in test/external/external_test_valid_remove.py
@@ -325,7 +324,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)+(r2*3))+r0)+-3)", "(((idx2*2)+r1)+-1)")
+    self.check(load, None, "((((idx1*24)+r0)+(r2*3))+-3)", "(((idx2*2)+r1)+-1)")
 
   def test_openpilot_conv3(self):
     # in openpilot 0.9.7
@@ -347,8 +346,8 @@ class TestImageSimplification(unittest.TestCase):
 
     self.check(load,
                "((((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)")
+               "(((idx0+(idx1*512))+(r1*64))+-192)",
+               "((((idx2*2)+(((idx1+((r1+5)//8))+1)//2))+r0)+-4)")
 
   def test_simplify1(self):
     # idx has the form (A % m, A // m + k) and valid has (c0 < A) and (A < c1)
@@ -386,16 +385,16 @@ class TestImageSimplification(unittest.TestCase):
 
     # TODO: can this be simplified further?
     load = get_load_image_uop(shape, alu9, (((alu8+(alu2*8))%64),(alu2//8)))
-    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+8)%64)", "((idx0%8)//2)")
+    self.check(load, "(idx0<256)", "((((idx0//32)+((idx0%8)*32))+8)%64)", "((idx0%8)//2)")
 
     load = get_load_image_uop(shape, alu9, (((alu8+(alu3*8))%64),(alu3//8)))
-    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+16)%64)", "((idx0%8)//2)")
+    self.check(load, "(idx0<256)", "((((idx0//32)+((idx0%8)*32))+16)%64)", "((idx0%8)//2)")
 
     load = get_load_image_uop(shape, alu9, (((alu8+(alu4*8))%64),(alu4//8)))
-    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+24)%64)", "((idx0%8)//2)")
+    self.check(load, "(idx0<256)", "((((idx0//32)+((idx0%8)*32))+24)%64)", "((idx0%8)//2)")
 
     load = get_load_image_uop(shape, alu9, (((alu8+(alu5*8))%64),(alu5//8)))
-    self.check(load, "(idx0<256)", "((((idx0%8)*32)+(idx0//32))%64)", "((idx0%8)//2)")
+    self.check(load, "(idx0<256)", "(((idx0//32)+((idx0%8)*32))%64)", "((idx0%8)//2)")
 
   def test_simplify5(self):
     # openpilot 0.9.7, chunk replacement to simplify

test/unit/test_uop_symbolic.py

@@ -116,6 +116,39 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual((a*b*3+a*b*b).divide_exact(a*b).simplify(), b+3)
     self.assertEqual((((a*-2)+14)*b).divide_exact(((a*-2)+14)).simplify(), b)
 
+  def helper_test_factor(self, expr, *factors):
+    factored = expr.factor(*factors)
+    self.check_equal_z3(expr, factored)
+    for fac in factors: self.assertIn(fac, factored.toposort())
+
+  def test_uop_factor(self):
+    a = Variable("a", 0, 8)
+    b = Variable("b", 0, 8)
+    c = Variable("c", 0, 8)
+    self.helper_test_factor((1400*a+2800*b), (a+2*b))
+    self.helper_test_factor((1400*a+2800*b)%9000, (a+2*b))
+    self.helper_test_factor((a+2*b), (a+2*b))
+    self.helper_test_factor((a+c+2*b), (a+2*b))
+    self.helper_test_factor((1400*a+c+2800*b)%9000, (a+2*b))
+    self.helper_test_factor((1399*a+c+2800*b)%9000+1400*a+2800*b, (a+2*b))
+    self.helper_test_factor((1400*a+c+2800*b)%9000+1400*a+2800*b, (a+2*b))
+    # self.assertIsNone((a+c+3*b).factor(a+2*b))
+    # self.assertIsNone((1399*a+c+2800*b).factor(a+2*b))
+
+  def test_uop_multiple_factors(self):
+    a = Variable("a", 0, 8)
+    b = Variable("b", 0, 8)
+    c = Variable("c", 0, 8)
+    d = Variable("d", 0, 8)
+    self.helper_test_factor((1400*a+2800*b+2*c+d), (a+2*b), (2*c+d))
+    self.helper_test_factor((100*a+200*b+5*c), (a+2*b), (5*c))
+    self.helper_test_factor((3*a+6*b+2*c+4*d), (a+2*b), (c+2*d))
+    self.helper_test_factor((7*a+14*b+3*c+6*d), (a+2*b), (3*c+6*d))
+    self.helper_test_factor((10*a+20*b+10*c+30*d), (a+2*b), (c+3*d))
+    self.helper_test_factor((10*c+(10*a+20*b)//3+30*d), (a+2*b), (c+3*d))
+    self.helper_test_factor((10*c+(10*a+20*b)//3+30*d), (a+2*b), (c+3*d))
+    # self.assertIsNone((7*a+14*b+3*c+6*d).factor((a+8*b), (2*c+6*d)))
+
   def test_divide_exact_not(self):
     a = Variable("a", 1, 8)
     b = Variable("b", 1, 8)
@@ -130,13 +163,13 @@ class TestSymbolic(unittest.TestCase):
     a = Variable("a", 0, 8)
     b = Variable("b", 0, 8)
     self.helper_test_variable(a*2+a*3, 0, 8*5, "(a*5)")
-    self.helper_test_variable(b+a*2+a*3, 0, 8*6, "(b+(a*5))")
+    self.helper_test_variable(b+a*2+a*3, 0, 8*6, "((a*5)+b)")
 
   def test_factorize_no_mul(self):
     a = Variable("a", 0, 8)
     b = Variable("b", 0, 8)
     self.helper_test_variable(a+a*3, 0, 8*4, "(a*4)")
-    self.helper_test_variable((a+b)+a*3, 0, 8*5, "(b+(a*4))")
+    self.helper_test_variable((a+b)+a*3, 0, 8*5, "((a*4)+b)")
     self.helper_test_variable((a*3+b)+b*3, 0, 8*7, "((a*3)+(b*4))")
 
   def test_neg(self):
@@ -159,8 +192,15 @@ class TestSymbolic(unittest.TestCase):
     b = Variable("b", 0, 8)
     self.helper_test_variable(a+a, 0, 16, "(a*2)")
     self.helper_test_variable((a+b)+b, 0, 24, "(a+(b*2))")
-    self.helper_test_variable((a*3+b)+a, 0, 40, "(b+(a*4))")
-    self.helper_test_variable((a+b)+a*3, 0, 40, "(b+(a*4))")
+    self.helper_test_variable((a*3+b)+a, 0, 40, "((a*4)+b)")
+    self.helper_test_variable((a+b)+a*3, 0, 40, "((a*4)+b)")
+
+  def test_add_self_seperated(self):
+    a = Variable("a", 0, 8)
+    b = Variable("b", 0, 8)
+    c = Variable("c", 0, 8)
+    self.helper_test_variable((a+b)+c+a, 0, 32, "(((a*2)+b)+c)")
+    self.helper_test_variable((a*3+b*2)+c*2+a*5, 0, 96, "(((a*8)+(b*2))+(c*2))")
 
   def test_sub_self(self):
     a = Variable("a", 0, 8)
@@ -279,7 +319,7 @@ class TestSymbolic(unittest.TestCase):
   def test_mod_congruence_multiple_vars(self):
     self.helper_test_variable((9+9*Variable("x",0,3)+9*Variable("y",0,3))%10, 3, 9, "(((x*-1)+(y*-1))+9)")
     self.helper_test_variable((7+9*Variable("x",0,2)+9*Variable("y",0,2)+Variable("z",0,2))%10, 3, 9,
-                              ("(((z+(x*-1))+(y*-1))+7)", "(((y*-1)+(z+(x*-1)))+7)"))
+                              ("(((z+(x*-1))+(y*-1))+7)", "(((y*-1)+(z+(x*-1)))+7)", "((((x*-1)+(y*-1))+z)+7)"))
     self.helper_test_variable((10+12*Variable("x",0,2)+Variable("y", 0, 4)%3)%13, 8, 12, "(((x*-1)+(y%3))+10)")
 
   def test_div_congruence(self):
@@ -455,7 +495,7 @@ class TestSymbolic(unittest.TestCase):
     ridx1005 = UOp.variable("ridx1005", 0, 2)
     ridx1006 = UOp.variable("ridx1006", 0, 2)
     self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -2, 20,
-      "(((((lidx1+(((gidx1*18)+(ridx1005*18))+(lidx0*162)))+(gidx0*2))+(ridx1006*2))+-40)//18)")
+      "((((((((gidx0*2)+(gidx1*18))+(lidx0*162))+lidx1)+(ridx1005*18))+(ridx1006*2))+-40)//18)")
 
   def test_add_div(self):
     # careful about the lower bounds and upper bounds
@@ -498,7 +538,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((d1*a*b*d1)//(d1), -1000, 1000, "(a*(b*d1))", test_z3=False)
     self.helper_test_variable((d1*a*d2*b*d1)//(d1*d2),  -1000, 1000, "(a*(b*d1))", test_z3=False)
     self.helper_test_variable((d1*a + b*d1)//(d1), -20, 20, "(a+b)", test_z3=False)
-    self.helper_test_variable((d1*a + b*d1 + c*d1)//(d1), -30, 30, "(c+(a+b))", test_z3=False)
+    self.helper_test_variable((d1*a + b*d1 + c*d1)//(d1), -30, 30, "((a+b)+c)", test_z3=False)
     self.helper_test_variable((3*a*d1 + 9*b*d1)//(3*d1*d2), -40, 40, "(((a+(b*3))//(d2*-1))*-1)", test_z3=False)
     self.helper_test_variable((3*a*d1 + 9*b*d1+3)//(3*d1*d2), -401, 399, "(((((a*d1)+((b*d1)*3))+1)//((d1*d2)*-1))*-1)", test_z3=False)
 
@@ -508,7 +548,7 @@ class TestSymbolic(unittest.TestCase):
     d = Variable("d", 1, 10)
     self.helper_test_variable((d*a+b)//d, 0, 20, "(a+(b//d))")
     self.helper_test_variable((d*a*20+b)//(5*d), 0, 42, "((a*4)+(b//(d*5)))")
-    self.helper_test_variable((d*a*20+b*d*5+10)//(5*d), 0, 52, "((b+(a*4))+(2//d))")
+    self.helper_test_variable((d*a*20+b*d*5+10)//(5*d), 0, 52, "(((a*4)+b)+(2//d))")
 
   def test_mod_gcd_factor_neg(self):
     self.helper_test_variable((Variable("a", 0, 10)*-4+4)%8, -4, 4, "((((a*-1)+1)%2)*4)")
@@ -561,9 +601,7 @@ class TestSymbolic(unittest.TestCase):
     lidx2 = Variable("lidx2", 0, 3)
     alu0 = gidx2*640+gidx1*160+(gidx0//5)*2+lidx0*320+lidx1*10
     self.helper_test_variable((alu0+lidx2*2+1)//20, 0, 8192,
-                              ("((((((gidx0//5)+lidx2)//5)+lidx1)//2)+(((gidx2*32)+(gidx1*8))+(lidx0*16)))",
-                               "(((lidx1+((lidx2+(gidx0//5))//5))//2)+((gidx2*32)+((gidx1*8)+(lidx0*16))))",
-                               "((((gidx1*8)+(gidx2*32))+(lidx0*16))+((lidx1+((lidx2+(gidx0//5))//5))//2))"))
+                              ("((((gidx1*8)+(gidx2*32))+(lidx0*16))+((lidx1+((lidx2+(gidx0//5))//5))//2))",))
 
   def test_sum_div_complex2(self):
     gidx0 = Variable("gidx0", 0, 7)
@@ -641,8 +679,21 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((gidx//4)*4+gidx%4, 0, 124, "gidx")
     self.helper_test_variable(lidx+gidx%4+(gidx//4)*4, 0, 248, "(gidx+lidx)")
     self.helper_test_variable(lidx+(gidx//4)*4+gidx%4, 0, 248, "(gidx+lidx)")
-    self.helper_test_variable(lidx+(gidx//4)*8+2*(gidx%4), 0, 372, "(lidx+(gidx*2))")
-    self.helper_test_variable(lidx+2*(gidx%4)+(gidx//4)*8, 0, 372, "(lidx+(gidx*2))")
+    self.helper_test_variable(lidx+(gidx//4)*8+2*(gidx%4), 0, 372, "((gidx*2)+lidx)")
+    self.helper_test_variable(lidx+2*(gidx%4)+(gidx//4)*8, 0, 372, "((gidx*2)+lidx)")
+
+  def test_div_mod_recombine_seperated(self):
+    gidx = Variable("gidx", 0, 124)
+    lidx = Variable("lidx", 0, 124)
+    a = Variable("a", 0, 3)
+    b = Variable("b", 0, 3)
+    c = Variable("c", 0, 3)
+    self.helper_test_variable(gidx%4+a+b+c+(gidx//4)*4, 0, 133, "(((a+b)+c)+gidx)")
+    self.helper_test_variable((gidx//4)*4+a+b*10+gidx%4, 0, 157, "((a+(b*10))+gidx)")
+    self.helper_test_variable(lidx+gidx%4+a+b+c//2+(gidx//4)*4, 0, 255, "((((a+b)+gidx)+lidx)+(c//2))")
+    self.helper_test_variable(lidx+(gidx//4)*8+b+c+a*8+2*(gidx%4), 0, 402, "(((((a*8)+b)+c)+(gidx*2))+lidx)")
+    # TODO: need better sorting for this one
+    # self.helper_test_variable(lidx+(gidx//4)*4+a*3+b*3+(c*10)%3+gidx%4, , , "")
 
   def test_div_mod_recombine_folded_mod(self):
     a = Variable("a", 0, 2)
@@ -1019,6 +1070,7 @@ class TestSymbolicRealWorld(unittest.TestCase):
       ("((((((((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49))+(gidx0*3211264))+(gidx1*784))+(gidx2*8))+(lidx4*100352))+lidx3)+2207744)",
        '((lidx3+((((((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49))+(gidx0*3211264))+(gidx1*784))+(gidx2*8))+(lidx4*100352)))+2207744)',
        '((lidx3+((lidx4*100352)+((gidx2*8)+((gidx1*784)+((gidx0*3211264)+((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49)))))))+2207744)',
+       '((((((((gidx0*3211264)+(gidx1*784))+(gidx2*8))+lidx3)+(lidx4*100352))+(((lidx5+1)//16)*802816))+(((lidx5+1)%16)*49))+2207744)',
       ))
 
 class TestBounds(unittest.TestCase):

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), 3 if RANGEIFY else 9)
+    self.assertEqual(len(backward_schedule), 4 if RANGEIFY else 9)
 
   def test_counters(self):
     IC, OC, X, Y = 4,4,9,9

tinygrad/apps/llm.py

@@ -135,7 +135,7 @@ class Transformer:
     x = self.token_embd(tokens)                           # (B, T, D)
     for block in self.blk: x = block(x, start_pos)
     # TODO: add temperature
-    return self.output(self.output_norm(x))[:, -1, :].softmax(-1).argmax(-1, keepdim=True)
+    return self.output(self.output_norm(x))[:, -1, :].softmax(-1, dtype="float").argmax(-1, keepdim=True)
 
   def __call__(self, tokens:Tensor, start_pos:int|UOp=0) -> Tensor:
     return (self.forward_jit if getenv("JIT", 1) and tokens.shape[1] == 1 and isinstance(start_pos, UOp) else self.forward)(tokens, start_pos)

tinygrad/codegen/__init__.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable
 import functools
 from dataclasses import dataclass
-from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY
+from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY, getenv
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype
 from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
@@ -10,14 +10,14 @@ from tinygrad.renderer import Renderer
 from tinygrad.codegen.lowerer import pm_lowerer, get_index
 from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
+from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
 from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render
 from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
 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.opt.postrange import pm_postrange_opt, pm_add_local_buffers, pm_pipeline
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range, pm_split_ranges
 from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen
 
@@ -77,16 +77,25 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
   # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+migrate_indexing, name="postopt symbolic"))
 
-  # expand
-  ret.append(RewriteStep(sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander"))
+  # locals
+  if getenv("LOCALS") or getenv("PIPELINE"):
+    ret.append(RewriteStep(pm_add_local_buffers, name="add locals"))
 
   # add locals
   ret.append(RewriteStep(pm_add_buffers+rangeify_codegen, name="add local buffers"))
 
+  # expand
+  ret.append(RewriteStep(sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander"))
+
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
   ret.append(RewriteStep(pm_reduce+gep_pushing, lambda _: ReduceContext(), name="remove_reduce"))
 
+  # pipelining
+  if getenv("PIPELINE"):
+    ret.append(RewriteStep(pm_pipeline, name="pipeline"))
+    ret.append(RewriteStep(sym, name="pipeline sym"))
+
   # add gpu dims (late). this works after devectorize, but it's faster here
   ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
 
@@ -101,6 +110,7 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
 
   # lower the index dtype to a concrete int
   ret.append(RewriteStep(pm_lower_index_dtype+load_store_indexing, lambda _: opts.device, name="lower all index dtypes"))
+  ret.append(RewriteStep(symbolic, name="post index symbolic"))
 
   # optional pre matcher
   if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher"))

tinygrad/codegen/late/devectorizer.py

@@ -50,6 +50,7 @@ def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp,
   # remove the gate from the index
   return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val, *store.src[2:])
 
+def no_load(u:UOp) -> bool: return not any(x.op is Ops.LOAD for x in u.sparents)
 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)),
@@ -60,6 +61,8 @@ load_store_indexing = PatternMatcher([
   # delete_redundant_gates (after expand)
   (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")),
                                   UPat.var("val")), name="store", allow_any_len=True), delete_redundant_gates),
+  # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes a pattern in reduce_collapse
+  (UPat.var("c")<(UPat.var("x", dtypes.index)+UPat.var("y")), lambda x,y,c: (-x < -(c-y)) if no_load(y) and no_load(c) and not no_load(x) else None),
 ])
 
 # ***** load/store grouping *****

tinygrad/codegen/opt/heuristic.py

@@ -35,7 +35,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
       pass
     if good_tc_opt:
       # skip hand-coded TC opts if AMX, upcasting will make kernel slower
-      if rngs is not None and not AMX:
+      if rngs is not None and not AMX and False:
         for tc_dim in [1,0]: # attempt to upcast M and N
           szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
           if szs:
@@ -43,6 +43,11 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
             rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
         if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
           tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
+      #tk.apply_opt(Opt(OptOps.LOCAL, 0, 2))
+      #tk.apply_opt(Opt(OptOps.LOCAL, 1, 2))
+      #tk.apply_opt(Opt(OptOps.UPCAST, 0, 2))
+      #tk.apply_opt(Opt(OptOps.UPCAST, 1, 2))
+      #tk.apply_opt(Opt(OptOps.UNROLL, 0, 8))
       return tk
 
   # make a copy so it does not mutate the input

tinygrad/codegen/opt/postrange.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
-import math, itertools
+import math, itertools, functools, operator
 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, 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, merge_dicts, get_single_element
+from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, dedup
 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 BufferizeOpts
 
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
@@ -257,12 +258,13 @@ class Scheduler:
           except KernelOptError: continue
 
           # we create the warp as a whole thing, in case some of these ranges are moved/removed later
-          warp = UOp.range(tc.threads, -1, AxisType.WARP)
+          warp_num = -10
           ne: list[UOp] = []
           for opt in tc.opts:
             if opt[0] == "l":
-              axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.LOCAL, input_new_rng=warp%2)
-              warp //= 2
+              warp = UOp.range(2, warp_num, AxisType.WARP)
+              axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.WARP, input_new_rng=warp)
+              warp_num += 1
             elif opt[0] == "u":
               axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.UPCAST)
             else: raise RuntimeError(f"unsupported opt {opt[0]} in tensor cores")
@@ -347,3 +349,135 @@ def apply_opts(ctx:Renderer, ast:UOp):
 pm_postrange_opt = PatternMatcher([
   (UPat(Ops.SINK, name="ast"), apply_opts),
 ])
+
+def add_local_buffer(x:UOp):
+  if x.tag is not None: return None
+  # should UPCAST/UNROLL be here?
+  branges = tuple([r for r in x.ranges if r.arg[-1] in {AxisType.WARP, AxisType.LOCAL, AxisType.UPCAST, AxisType.UNROLL}])[::-1]
+  buf = UOp(Ops.BUFFERIZE, x.dtype, src=(x.replace(tag=1),)+branges, arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL))
+  return UOp(Ops.INDEX, x.dtype, src=(buf,)+branges)
+
+pm_add_local_buffers = PatternMatcher([
+  (UPat(Ops.LOAD, name="x"), add_local_buffer),
+])
+
+def add_pipeline(x:UOp):
+  if x.tag == 1: return None
+  if x.arg[-1] == AxisType.REDUCE:
+    # 3 splits
+    #srcs = (x.const_like(0), x.replace(src=(x.src[0]-2,), tag=1)+1, x.src[0]-1)
+
+    # 4 split
+    rng = x.replace(src=((x.src[0]-2)//2,), tag=1)
+    srcs = (x.const_like(0), rng*2+1, rng*2+2, x.src[0]-1)
+
+    return UOp(Ops.SPLIT, x.dtype, src=srcs, arg=1).simplify()
+    #vec = UOp(Ops.VECTORIZE, x.dtype.vec(3), src=(x.const_like(0), x.replace(src=(x.src[0]-2,), tag=1)+1, x.src[0]-1)).simplify()
+    #return UOp(Ops.UNROLL, x.dtype, src=(vec,), arg=())
+
+def do_split(x:UOp):
+  splits = [x for x in x.src if x.op is Ops.SPLIT]
+  if len(splits) == 0: return None
+  if x.op is Ops.SINK: return x.replace(src=x.src[0].src)
+  #if x.op is Ops.REDUCE:
+    #assert x.src[0].op is Ops.SPLIT
+    #rr = [y for y in x.src[1].toposort() if y.op is Ops.RANGE][0]
+    #return x.replace(src=(functools.reduce(operator.add, x.src[0].src), rr))
+  uu = []
+  for i in range(len(splits[0].src)):
+    new_srcs = []
+    for s in x.src:
+      if s.op is Ops.SPLIT:
+        new_srcs.append(s.src[i])
+      else:
+        new_srcs.append(s)
+    uu.append(UOp(x.op, x.dtype, tuple(new_srcs), x.arg, x.tag))
+  if x.op is Ops.STORE and len(splits) == 2:
+    dls = dedup([x for x in uu[0].toposort() if x.op is Ops.DEFINE_LOCAL])
+    uu2 = []
+    # NOTE: here we have to order the STORES and fix the ranges
+    for i,u in enumerate(uu):
+      subs = {}
+      for dl in dls: subs[dl] = dl.replace(arg=(dl.arg, i%2))
+      uu2.append(u.substitute(subs))
+    # TODO: reorder (is the reorder just a toposort question?)
+    # there's 4 barriers and 4 output stores
+    # load 0
+    # barrier (between load 0 and compute 0)
+    #   load 1 (depends on range)
+    #   compute 0
+    #   barrier (between load 0+2 and load 2)
+    #   load 2 (depends on range)
+    #   compute 1
+    #   barrier (between load 1 and load 1)
+    # load 3
+    # compute 2
+    # barrier (between load 3 and compute 3)
+    # compute 3
+
+    # uncouple based on barriers
+    def do_uncouple(ctx, l:UOp, b:UOp):
+      ctx[1].append(b.src[0])
+      return l.replace(src=l.src[0:1]+(UOp(Ops.NOOP, tag=ctx[0]),))
+    uncouple_barrier = PatternMatcher([
+      (UPat(Ops.LOAD, src=(UPat(), UPat(Ops.BARRIER, name='b')), name='l'), do_uncouple),
+      (UPat(Ops.LOAD, src=(UPat(), UPat(), UPat()), name='x'), lambda x: x.replace(src=x.src[0:2])),
+    ])
+
+    loads = []
+    computes = []
+    for i,u in enumerate(uu2):
+      cc = graph_rewrite(u, uncouple_barrier, ctx=(i,tloads:=[]))
+      computes.append(cc)
+      loads.append(UOp(Ops.NOOP, src=tuple(tloads)))
+
+    # remove the range from here
+    computes[2] = computes[2].replace(src=computes[2].src[0:2])
+
+    # pipelined!
+    ret = computes[3]
+
+    # put computes[2] before compute[3]
+    const_store = [x for x in ret.toposort() if x.op is Ops.STORE and x.src[1].op is Ops.CONST][0]
+    ret = ret.substitute({const_store:const_store.replace(tag=1)})
+    ret = ret.substitute({const_store.replace(tag=1):computes[2].barrier()})
+
+    # put loads[3] before compute[2] (both ports)
+    const_store = [x for x in ret.toposort() if x.op is Ops.STORE and x.src[1].op is Ops.CONST][0]
+    ret = ret.substitute({const_store:const_store.replace(tag=1)})
+    ret = ret.substitute({const_store.replace(tag=1):loads[3], UOp(Ops.NOOP, tag=2):loads[3]})
+
+    # put compute[1] before loads[3]
+    load = [x for x in loads[3].toposort() if x.op is Ops.LOAD][0]
+    ret = ret.substitute({load: load.replace(src=load.src+(computes[1].barrier(),))})
+
+    # put loads[2] before compute[1] (both ports)
+    const_store = [x for x in ret.toposort() if x.op is Ops.STORE and x.src[1].op is Ops.CONST][0]
+    ret = ret.substitute({const_store:const_store.replace(tag=1)})
+    ret = ret.substitute({const_store.replace(tag=1):loads[2], UOp(Ops.NOOP, tag=1):loads[2]})
+
+    # put compute[0] before loads[2]
+    load = [x for x in loads[2].toposort() if x.op is Ops.LOAD][0]
+    ret = ret.substitute({load: load.replace(src=load.src+(computes[0].barrier(),))})
+
+    # put loads[1] before compute[0] (one port)
+    ret = ret.substitute({UOp(Ops.NOOP, tag=0):loads[1]})
+
+    # put loads[0] before loads[1]
+    load = [x for x in loads[1].toposort() if x.op is Ops.LOAD][0]
+    ret = ret.substitute({load: load.replace(src=load.src+(loads[0],))})
+
+    pm_remove_noops = PatternMatcher([
+      (UPat(Ops.NOOP, src=(UPat.var('x'),)), lambda x: x),
+    ])
+    return graph_rewrite(ret, pm_remove_noops, name="remove noops")
+  return UOp(Ops.SPLIT, x.dtype, src=tuple(uu))
+
+pm_pipeline = PatternMatcher([
+  (UPat(Ops.RANGE, name="x"), add_pipeline),
+  # do expansion
+  (UPat(GroupOp.All, name="x", custom_early_reject=set([Ops.SPLIT])), do_split),
+  #(UPat(Ops.STORE, src=(UPat(), UPat(), UPat(Ops.CONST)), name="x"), lambda x: x.replace(src=x.src[:2])),
+  (UPat(Ops.STORE, src=(UPat.var('x'), UPat.var('z'), UPat.var('rr')), name='r'),
+   lambda x,rr,r,z: r.replace(src=(x,z)+tuple([y for y in rr.toposort() if y.op is Ops.RANGE]))),
+])

tinygrad/codegen/opt/search.py

@@ -68,7 +68,7 @@ def _try_compile_linearized_w_idx(x:tuple[int,Scheduler], compiler:Compiler) ->
   try:
     p = get_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].opts)
     assert p.uops is not None, "uop list wasn't generated?"
-    if len(p.uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
+    if len(p.uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 6000)) > 0:
       if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(p.uops)=}, {uops_max=}")
       raise RuntimeError("too many uops")
     st = time.perf_counter()

tinygrad/device.py

@@ -125,7 +125,7 @@ class Buffer:
   def allocate(self, opaque=None, external_ptr=None) -> Buffer:
     assert not self.is_initialized(), "can't allocate already allocated buffer"
     if DEBUG >= 7: print(f"buffer: allocate {self.nbytes} bytes on {self.device}")
-    if MAX_BUFFER_SIZE > 0 and self.size > MAX_BUFFER_SIZE: raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
+    if not self.device.startswith("NULL") and self.size > MAX_BUFFER_SIZE > 0: raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
     self.allocator:Allocator = Device[self.device].allocator
     if external_ptr is not None:
       self.options = replace(self.options, external_ptr=external_ptr) if self.options else BufferSpec(external_ptr=external_ptr)

tinygrad/helpers.py

@@ -148,6 +148,8 @@ CPU_COUNT = ContextVar("CPU_COUNT", max(1, (os.cpu_count() or 1) // (4 if ARCH_X
 CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1)
 VIZ = PROFILE = ContextVar("VIZ", 0)
 SPEC = ContextVar("SPEC", 0)
+# TODO: disable by default due to speed
+IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/renderer/cstyle.py

@@ -295,7 +295,7 @@ class MetalRenderer(CStyleLanguage):
   # language options
   kernel_typedef = "kernel void"
   buffer_prefix = "device "
-  smem_prefix = "threadgroup __attribute__((aligned(16))) "
+  smem_prefix = "threadgroup " #__attribute__((aligned(16))) "
   arg_int_prefix = "constant int&"
   barrier = "threadgroup_barrier(mem_flags::mem_threadgroup);"
   float4 = "float4"

tinygrad/renderer/ptx.py

@@ -102,7 +102,7 @@ string_rewrite = PatternMatcher([
   (UPat(Ops.CAST, name="x", dtype=dtypes.bool, src=(UPat.var("a"),)),
    lambda ctx, x, a: f"setp.ne.b{ctx.types[a.dtype][1:]} {ctx.r[x]}, {ctx.r[a]}, {render_val(0, a.dtype)};"),
   (UPat(Ops.CAST, name="x", src=(UPat.var("a"),)),
-   lambda ctx, x, a: f"cvt{modifier(x.dtype, a.dtype)}.{ctx.types[x.dtype]}.{ctx.types[a.dtype]} {ctx.r[x]}, {ctx.r[a]};"),
+   lambda ctx, x, a: f"cvt{modifier(x.dtype, a.dtype)}.{ctx.cast_types[x.dtype]}.{ctx.cast_types[a.dtype]} {ctx.r[x]}, {ctx.r[a]};"),
   (UPat(Ops.LOAD, name="x", src=(UPat.var('loc'), UPat(name='alt'), UPat(name="gate", op=GroupOp.ALU)), allow_any_len=True),
     lambda ctx, x, loc, alt, gate: flatten([
     [f"mov.{ctx.mem_types[x.dtype.scalar()]} {v}, {render_val(0, x.dtype.scalar())};" for v in ctx.r[x]],
@@ -146,12 +146,12 @@ class PTXRenderer(Renderer):
 .address_size 64
 .visible .entry"""
   barrier = "bar.sync\t0;"
-  # HACK: Use s16 and u16 for int8 and uint8 buffers. This can be wrong in cast.
   types: dict[DType, str] = { dtypes.int8: "s16", dtypes.int16: "s16", dtypes.int32: "s32", dtypes.int64: "s64",
                               dtypes.uint8: "u16", dtypes.uint16: "u16", dtypes.uint32: "u32", dtypes.uint64: "u64",
                               dtypes.float16: "f16", dtypes.float32: "f32", dtypes.float64: "f64", dtypes.bool: "pred" }
 
   mem_types: dict[DType, str] = {**types, dtypes.int8: "s8", dtypes.uint8: "u8", dtypes.bool: "u8", dtypes.float16: "b16"}
+  cast_types: dict[DType, str] = {**types, dtypes.int8: "s8", dtypes.uint8: "u8"}
 
   def render_kernel(self, kernel, function_name, bufs, regs, uops) -> str:
     def fmt(line): return line if line[0]=="$" else "\t" + line.replace(" ", "\t" if len(line.split(" ")[0]) > 7 else "\t\t", 1)

tinygrad/runtime/ops_remote.py

@@ -176,7 +176,7 @@ class RemoteHandler:
     self.sessions: defaultdict[SessionKey, RemoteSession] = defaultdict(RemoteSession)
 
     try: self.ib_ctx: IBCtx|None = IBCtx(getenv("IB_DEV", 0))
-    except (IndexError, AttributeError): self.ib_ctx = None
+    except (RuntimeError, IndexError, AttributeError): self.ib_ctx = None
     self.ib_lock = asyncio.Lock()
     self.ib_conns: dict[str, IBConn|None] = {}
     self.iova_cache: dict[tuple[SessionKey, int], tuple[int, int, int]] = {}

tinygrad/runtime/support/ib.py

@@ -10,7 +10,7 @@ DEFAULT_PORT, DEFAULT_GID = getenv("DEFAULT_PORT", 1), getenv("DEFAULT_GID", 3)
 IOVA_ALIGN = resource.getpagesize()
 
 def checkz(x, ret=None):
-  assert x == 0, f'{x} != 0 (errno {ctypes.get_errno()})'
+  if x != 0: raise RuntimeError(f'{x} != 0 (errno {ctypes.get_errno()})')
   return ret
 
 @dataclass(frozen=True)

tinygrad/schedule/rangeify.py

@@ -46,6 +46,9 @@ earliest_rewrites = PatternMatcher([
   # just removing it works...
   (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]),
 
+  # remove CONTIGUOUS if the BUFFER is already contiguous
+  (UPat(Ops.BUFFER).f(Ops.RESHAPE, name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
+
   # split_reduceop
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
 
@@ -62,13 +65,20 @@ earliest_rewrites = PatternMatcher([
    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 rules **
+
+  # early fixup const copy
+  (UPat(Ops.COPY, src=(UPat.var("s"), UPat()), name="c"), lambda c,s: c.const_like(ss.arg) if (ss:=s.base).op is Ops.CONST 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),
 
-  # 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))),
+  # 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),
+
+  # ** assign rules **
 
   # 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"),
@@ -78,11 +88,8 @@ earliest_rewrites = PatternMatcher([
    # realize before assign if input permutes the target buffer
    (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), find_permutes),
 
-  # 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),
-
-  # 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]),
+  # contiguous buffer is buffer, this is for *correctness* of assign, not just speed
+  (UPat(Ops.CONTIGUOUS, name="root", src=(UPat(Ops.BUFFER),)), lambda root: root.src[0].forced_reshape(root.shape).rtag(root.tag)),
 ])
 
 # *****************
@@ -270,12 +277,11 @@ def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp):
 
 def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
   if c not in ctx.seen_children: ctx.seen_children[c] = {}
+  ctx.seen_children[c][x.arg[0]] = idx
   # wait here until we have seen all the children
   if len(ctx.seen_children[c]) != x.arg[1]:
     ctx.progress += 1
     if ctx.progress > 10000: raise RuntimeError("children not making progress")
-    # NOTE: we mark this here
-    ctx.seen_children[c][x.arg[0]] = idx
     raise RewriteNotReady
   ctx.progress = 0
 
@@ -358,7 +364,7 @@ pm_rangeify = pm_mops+PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.REDUCE_AXIS})),), allow_any_len=True, name="idx"), might_end_axis),
 
   # handle size 0
-  (UPat(Ops.INDEX, name="x"), lambda x: x.replace(src=(x.const_like(0),)+x.src[1:]) if x.st is not None and x.size == 0 else None),
+  #(UPat(Ops.INDEX, name="x"), lambda x: x.replace(src=(x.const_like(0),)+x.src[1:]) if x.st is not None and x.size == 0 else None),
 
   # handle assign
   (UPat(Ops.INDEX, src=(UPat(Ops.ASSIGN, name="assign"),), allow_any_len=True, name="x"),
@@ -372,13 +378,14 @@ pm_rangeify = pm_mops+PatternMatcher([
   (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, Ops.MSELECT}).f(Ops.INDEX, name="x"), unprocessed_index),
+  (UPat(GroupOp.All-{Ops.REALIZE, Ops.BUFFERIZE, Ops.MSELECT, Ops.MSTACK}).f(Ops.INDEX, name="x"), unprocessed_index),
 ])
 
 # *****************
 # 3.5 cleanups
 
-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN}
+# Ops.NOOP happens when we have a COPY to the device the Tensor is already on. We treat it like COPY here for MSTACK.
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP}
 
 # 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):
@@ -430,7 +437,7 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
 
     # 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)
+    def okay_reduce(x:UOp): return all(y.op not in {Ops.BUFFER, Ops.BUFFERIZE, 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
@@ -438,7 +445,8 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   # if it makes it here, the bufferize is removed
   # this is the ranges replaced
   # NOTE: if buf src is a const, we don't replace it
-  return src.substitute({k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST})
+  replaces = flatten([(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST])
+  return UOp(Ops.SUBSTITUTE, src=(src, UOp(Ops.NOOP, src=tuple(replaces[0::2])), UOp(Ops.NOOP, src=tuple(replaces[1::2]))))
 
 def pre_bufferize(b:UOp, x:UOp, copy:UOp):
   nb = b.replace(src=(b.src[0].contiguous(),)+b.src[1:])
@@ -713,6 +721,24 @@ replace_contiguous = PatternMatcher([
   (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),
 ])
 
+def do_sub_recurse(s:UOp):
+  x,keys,values = s.src[0], s.src[1].src, s.src[2].src
+  # SUBSTITUTE applied to SUBSTITUTE runs the child SUB on the parents. though this is probably wrong in the generic case
+  if x.op is Ops.SUBSTITUTE:
+    sub_k = UOp(Ops.SUBSTITUTE, src=(x.src[1],)+s.src[1:])
+    sub_v = UOp(Ops.SUBSTITUTE, src=(x.src[2],)+s.src[1:])
+    return UOp(Ops.SUBSTITUTE, src=(x.src[0], sub_k, sub_v))
+  # here we actually do the SUBSTITUTE
+  if x in keys: return values[keys.index(x)]
+  # we filter any keys that aren't in parents. this keeps the algorithm O(output graph size)
+  new_kv = {k:v for k,v in zip(keys,values) if k in x.sparents}
+  # if there's no SUBSTITUTEs left, we can just return x
+  if len(new_kv) == 0: return x
+  # then we add SUBSTITUTE to all parents
+  uop_keys, uop_values = UOp(Ops.NOOP, src=tuple(new_kv.keys())), UOp(Ops.NOOP, src=tuple(new_kv.values()))
+  return x.replace(src=tuple([UOp(Ops.SUBSTITUTE, src=(y,uop_keys,uop_values)) for y in x.src]))
+pm_substitute_recurse = PatternMatcher([(UPat(Ops.SUBSTITUTE, src=(UPat(), UPat(Ops.NOOP), UPat(Ops.NOOP)), name="s"), do_sub_recurse)])
+
 @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] = []
@@ -731,12 +757,14 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
   tsink = graph_rewrite(tsink, symbolic_simple+pm_reduce_unparented, name="symbolic")  # this supports const folding
   tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
+  # TODO: can you substitute and remove costly buffers at the same time?
+  tsink = graph_rewrite(tsink, pm_substitute_recurse, bottom_up=True, name="run substitutes")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rangeify_ctx, name="limit buffers")
 
   # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
   # MSTACK stacks multiple BUFFERIZEs in one tagged tensor
   # if it's not tagged by here, it's out
-  tsink = UOp.sink(*[x for x in tsink.parents if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST} and x.tag is not None])
+  tsink = UOp.sink(*[x for x in tsink.parents if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST, Ops.BUFFER} and x.tag is not None])
 
   if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")
 

tinygrad/uop/__init__.py

@@ -10,6 +10,7 @@ class FastEnum(IntEnum):
 class Ops(FastEnum):
   # uops that aren't rendered
   NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto(); REWRITE_ERROR = auto()  # noqa: E702
+  SENTINEL = auto()
 
   # track children
   CHILD = auto(); CHILDREN = auto() # noqa: E702
@@ -19,6 +20,8 @@ class Ops(FastEnum):
 
   # create buffer
   BUFFERIZE = auto()
+  SUBSTITUTE = auto()
+  SPLIT = auto()
 
   # ops that adjust the behavior of the scheduler
   CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702

tinygrad/uop/ops.py

@@ -8,7 +8,7 @@ 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, SPEC
-from tinygrad.helpers import strip_parens
+from tinygrad.helpers import strip_parens, make_tuple
 if TYPE_CHECKING:
   from tinygrad.shape.shapetracker import ShapeTracker
   from tinygrad.device import Buffer, MultiBuffer
@@ -79,6 +79,20 @@ class UOpMetaClass(type):
 buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER uops to their device Buffers
 all_metadata:weakref.WeakKeyDictionary[UOp, tuple[Metadata, ...]] = weakref.WeakKeyDictionary() # TODO: should this be here?
 
+# recursive_property replaces functools.cached_property in recursive UOp functions to prevent RecursionError
+_NOT_FOUND = object()
+class recursive_property(property):
+  def __init__(self, fxn):
+    self.fxn = fxn
+    self.nm = "_RECURSIVE_PROPERTY_"+fxn.__name__
+    self.__doc__ = fxn.__doc__
+  def __get__(self, x:UOp|None, owner=None):
+    if x is None: return self
+    if (val:=x.__dict__.get(self.nm, _NOT_FOUND)) is _NOT_FOUND:
+      for s in x.toposort(lambda z: not hasattr(z, self.nm)):
+        s.__dict__[self.nm] = val = self.fxn(s)
+    return val
+
 # NOTE: this should be frozen, but frozen is slower
 @dataclass(eq=False, slots=True)
 class UOp(MathTrait, metaclass=UOpMetaClass):
@@ -115,7 +129,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   def f(self, op, **kwargs): return UOp(op, dtype=kwargs.pop("dtype", self.dtype), src=(self,), **kwargs)
 
-  @functools.cached_property
+  @recursive_property
   def parents(self:UOp) -> dict[UOp, None]:
     ret = {s:None for s in self.src}
     for s in self.src: ret.update(s.parents)
@@ -150,6 +164,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def tuplize(self:UOp) -> tuple:
     return (self.op.value, self.arg, self.dtype,)+tuple([x.tuplize for x in self.src])
 
+  @functools.cached_property
+  def order_add(self:UOp) -> tuple:
+    if self.op is Ops.MUL and self.src[1].op in (Ops.CONST, Ops.VCONST): return (self.src[0].tuplize, make_tuple(self.src[1].arg, 1))
+    return (self.tuplize, (0,))
+
   @property
   def ptrdtype(self) -> PtrDType:
     if not isinstance(self.dtype, PtrDType): raise RuntimeError("ptrdtype called on UOp without PtrDType")
@@ -157,7 +176,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # *** uop shape stuff ***
 
-  @functools.cached_property
+  @recursive_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, Ops.MSTACK,
                                                    Ops.MSELECT, Ops.BUFFER, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.STORE}:
@@ -182,7 +201,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
     # BUFFER/BUFFER_VIEW and KERNEL only have a size
     if self.op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return ShapeTracker.from_shape((self.size,))
-    if self.op is Ops.KERNEL: return ShapeTracker.from_shape((self.arg.ast.size,))
+    if self.op is Ops.KERNEL:
+      ast = self.arg.ast
+      return ShapeTracker.from_shape((ast.size,)) if ast.st is not None else None
     if self.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
       sz = self.ptrdtype.size
       return ShapeTracker.from_shape((sz,)) if sz > 0 else None
@@ -241,9 +262,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   def simplify(self, tracked=False):
     # late import!
-    from tinygrad.uop.symbolic import symbolic
+    from tinygrad.uop.symbolic import symbolic_flat
     with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value):
-      return graph_rewrite(self, symbolic, name="simplify")
+      return graph_rewrite(self, symbolic_flat, name="simplify")
   def ssimplify(self) -> UOp|ConstType: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret
   def _eval(self, dtype, expected_type:Type[T]) -> T:
     assert self.dtype in dtype, f"eval with wrong dtype {self}"
@@ -568,6 +589,32 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       if (d0:=self.src[0].divides(v)) is not None: return d0 * self.src[1]
       if (d1:=self.src[1].divides(v)) is not None: return self.src[0] * d1
     return None # generic None if we aren't sure
+  def factor(self, *factors: UOp) -> UOp:
+    # factor out expr from self if possible, might return self
+    # (1400*a + 2800*b + c).factor(a+2*b) -> 1400*(a+2*b) + c
+    if self.dtype in dtypes.floats: return self
+    if self.op is Ops.ADD:
+      factored = []
+      # dict of {term: const_factor}, i.e. {a: 1, b: 2}
+      remainders = dict([(u.divides(f:=u.const_factor()).simplify(),f) for u in self.split_uop(Ops.ADD)])
+      for fac in factors:
+        if fac.dtype not in (dtypes.index,)+dtypes.ints: continue
+        fac_terms = dict((u.divides(f:=u.const_factor()).simplify(),f) for u in fac.split_uop(Ops.ADD))
+        factored_terms  = {k:v for k,v in remainders.items() if k in fac_terms}
+        new_remainders  = {k:v for k,v in remainders.items() if k not in fac_terms}
+
+        if any(u not in factored_terms for u in fac_terms) or any(factored_terms[u]%fac_terms[u]!=0 for u in fac_terms) or not \
+          all_same(mul:=[factored_terms[u]//fac_terms[u] for u in fac_terms]):
+          continue
+
+        remainders = new_remainders
+        factored.append(fac*mul[0])
+      if not factored: return self
+      start = functools.reduce(operator.add, factored)
+      return sum([k.factor(*factors)*v for k,v in remainders.items()], start=start)
+
+    if self.op not in GroupOp.ALU|{Ops.VECTORIZE}: return self
+    return self.replace(src=tuple(s.factor(*factors) for s in self.src))
   def pop_const(self, op=Ops.ADD) -> tuple[UOp, ConstType]:
     return (self.src[0], self.src[1].arg) if self.op is op and self.src[1].op is Ops.CONST else (self, identity_element(op, self.dtype))
   @staticmethod
@@ -992,6 +1039,7 @@ if TRACK_MATCH_STATS or PROFILE:
 
 # *** simple graph rewrite engine ***
 
+SENTINEL = UOp(Ops.SENTINEL)
 class RewriteNotReady(Exception): pass
 class BottomUpGate(Exception): pass
 class RewriteContext:
@@ -1004,45 +1052,58 @@ class RewriteContext:
     self.replace: dict[UOp, UOp] = {}
 
   def cached_pm_rewrite(self, x:UOp):
-    if (ret:=self.pm_cache.get(x,False)) is not False: return ret
+    if (ret:=self.pm_cache.get(x,SENTINEL)) is not SENTINEL: return ret
     ret = self.pm_cache[x] = cast(PatternMatcher, self.pm).rewrite(x, self.ctx)
     return ret
 
   def cached_bpm_rewrite(self, x:UOp):
-    if (ret:=self.bpm_cache.get(x,False)) is not False: return ret
+    if (ret:=self.bpm_cache.get(x,SENTINEL)) is not SENTINEL: return ret
     ret = self.bpm_cache[x] = cast(PatternMatcher, self.bpm).rewrite(x, self.ctx)
     return ret
 
   def unified_rewrite(self, root:UOp) -> UOp:
     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
+    REWRITE_STACK_LIMIT = getenv("REWRITE_STACK_LIMIT", 250000)
     while stack:
-      if len(stack) > getenv("REWRITE_STACK_LIMIT", 250000): raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
+      if len(stack) > REWRITE_STACK_LIMIT: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
       n, stage, new_n = stack.pop()
       if n in self.replace: continue  # skip any nodes we have seen
-      try:
-        if stage == 0:
+      if stage == 0:
+        # if bottom up, we rewrite this node early. in both cases, we add its parents to the stack
+        if self.bpm is not None:
+          # apply rewrite rules until a fixed point is reached. may return `uop` itself if PatternMatcher doesn't match
+          test_n: UOp|None = n
+          seen = set()
           try:
-            # if bottom up, we rewrite this node early. in both cases, we add its parents to the stack
-            if self.bpm is not None:
-              # apply rewrite rules until a fixed point is reached. may return `uop` itself if PatternMatcher doesn't match
-              test_n: UOp|None = n
-              seen = set()
-              while test_n is not None:
-                if test_n in seen: raise RuntimeError("infinite loop in fixed_point_rewrite")
-                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):
-              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:
-          try: new_src = tuple([self.replace[x] for x in new_n.src])
-          except KeyError: raise RewriteNotReady
-          if new_src == new_n.src:
+            while test_n is not None:
+              if test_n in seen: raise RuntimeError("infinite loop in fixed_point_rewrite")
+              seen.add(test_n)
+              new_n, test_n = test_n, self.cached_bpm_rewrite(test_n)
+          except RewriteNotReady:
+            # try the full thing again later
+            stack.appendleft((n, 0, n))
+            continue
+          except BottomUpGate:
+            # if the bpm matching raised a gate, we are done with this node and dont continue down the srcs
+            self.replace[n] = new_n
+            continue
+        stack.append((n, 1, new_n))
+        for x in reversed(new_n.src):
+          if x in on_stack: continue
+          stack.append((x, 0, x))
+          on_stack.add(x)
+      elif stage == 1:
+        tmp = []
+        for x in new_n.src:
+          if (rx:=self.replace.get(x, SENTINEL)) is SENTINEL:
+            # if some new sources aren't ready, we try this again later
+            stack.appendleft((n, 1, new_n))
+            break
+          tmp.append(rx)
+        else:
+          # in stage 1, once all srcs are rewritten, rebuild (if changed) or run top-down rewrite
+          if (new_src:=tuple(tmp)) == new_n.src:
             # if top down, do the rewrite. if no rewrite or bottom up, we are done rewriting this node so we add it to the dict
             if self.pm is None or (new_src_n:=self.cached_pm_rewrite(new_n)) is None:
               self.replace[n] = new_n
@@ -1053,13 +1114,14 @@ class RewriteContext:
           # trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n
           stack.append((n, 2, new_src_n))
           stack.append((new_src_n, 0, new_src_n))
+      else:
+        # in stage 2, we link the result of new_n to the result of n
+        if (replaced_new_n:=self.replace.get(new_n, SENTINEL)) is SENTINEL:
+          # not ready, try the link later
+          stack.appendleft((n, 2, new_n))
         else:
-          # in stage 2, we link the result of new_n to the result of n
-          try: self.replace[n] = self.replace[new_n]
-          except KeyError: raise RewriteNotReady
-      except RewriteNotReady:
-        # retry this later
-        stack.appendleft((n, stage, new_n))
+          # otherwise we are done
+          self.replace[n] = replaced_new_n
     return self.replace[root]
 
 @track_matches

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, RANGEIFY
+from tinygrad.helpers import all_same, prod, DEBUG, IGNORE_OOB, Context, cpu_profile, RANGEIFY
 from tinygrad.shape.shapetracker import ShapeTracker
 try:
   import z3
@@ -55,9 +55,6 @@ try:
   z3_imported = True
 except (ImportError, AttributeError): z3_imported = False
 
-# if you have z3 installed, by default we check the bounds
-IGNORE_OOB = ContextVar("IGNORE_OOB", int(not z3_imported))
-
 buffer_spec = PatternMatcher([
   (UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True),
   (UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d:

tinygrad/uop/symbolic.py

@@ -274,10 +274,16 @@ gep_pushing = PatternMatcher([
   (UPat(Ops.WMMA, name="wmma").f(Ops.GEP, name="gep"), gep_through_wmma),
 ])
 
+def chain_insert(chain, b, op):
+  if chain.op is not op or b.order_add > chain.src[1].order_add: return chain.alu(op, b)
+  return chain_insert(chain.src[0], b, op).alu(op, chain.src[1])
+
 commutative = PatternMatcher([
   # ** COMMUTATIVE flipping (only for index) **
   # NOTE: this can break merging vector math by only flipping some of them
-  (UPat(GroupOp.Commutative, dtype=dtypes.index, name='x'), lambda x: x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize else None),
+  (UPat(GroupOp.Commutative-{Ops.ADD}, dtype=dtypes.index, name='x'), lambda x:
+    x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize else None),
+  (UPat(Ops.ADD, dtype=dtypes.index, name="x"), lambda x: functools.reduce(operator.add, sorted(x.split_uop(Ops.ADD), key=lambda u: u.order_add)))
 ])
 
 symbolic = symbolic_simple+commutative+PatternMatcher([
@@ -373,7 +379,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
 ])+gep_pushing
 
 symbolic_flat = symbolic+PatternMatcher([
-  # ** combine terms (opinionated) **
+  # ** combine terms (opinionated), can make it harder to substitute valids **
   (-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)),  # -(x+y) -> -x + -y
   # (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
   ((UPat.var("x", dtypes.index) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
@@ -405,10 +411,13 @@ def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None:
   # don't simplify any other gates, can lead to OOB, we substitute them back later
   uop = uop.substitute((load_subs:={u: UOp(Ops.NOOP, arg=u) for u in uop.toposort() if u.op is Ops.INDEX}))
 
+  all_candidates = []
   # simplify uop given that valid is True
-  for expr,v in bounds.items():
+  for i, (expr,v) in enumerate(bounds.items()):
     v0, v1 = (expr.vmin if v[0] is None else v[0], expr.vmax if v[1] is None else v[1])
     expr = expr.substitute(load_subs)  # make sure expr appears in same form in the uop
+    # if the expr is an add we try and factorize so its more likely to substitute
+    if expr.op is Ops.ADD: uop = uop.factor(expr)
     # some expr has lower bound > upper bound -> valid is an empty set and we return None
     if v0 > v1: return None
     # whole node became a const
@@ -421,7 +430,9 @@ def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None:
       # if the constraint is a simplex: X0 + X1 + ... > 0, we can check if all Xi > 0 simplify into the same output
       candidates.append([(Xi, UOp.variable("fake", 1, Xi.vmax, Xi.dtype)) for Xi in expr.split_uop(Ops.ADD)])
     # try checking the whole clause
-    if expr in uop.toposort(): candidates.append([(expr, UOp.variable("fake", v0, v1, expr.dtype))])
+    if expr in uop.toposort():
+      candidates.append([tup:=(expr, UOp.variable(f"fake{i}", v0, v1, expr.dtype))])
+      all_candidates.append(tup)
 
     for candidate in candidates:
       # if every branch in candidate gives the same simplified uop, we can rewrite the uop
@@ -431,6 +442,9 @@ def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None:
         if all_same([uops.src[1] for uops in newuops]): uop = uop.replace(src=(uop.src[0], newuops[0].src[1]))
       elif all_same(newuops): uop = newuops[0]
 
+  uop = uop.factor(*(e[0] for e in all_candidates))
+  uop = uop.substitute(sub_dict:=dict(all_candidates)).simplify().substitute({newX:X for X,newX in sub_dict.items()}).simplify()
+
   # put the loads back in
   uop = uop.substitute({v:k for k,v in load_subs.items()})
   return uop

tinygrad/viz/serve.py

@@ -20,7 +20,8 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF",
                Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D", Ops.REALIZE: "#C1C14D",
-               Ops.CHILDREN: "#80ffc0", Ops.CHILD: "#80fff0", Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e"}
+               Ops.CHILDREN: "#80ffc0", Ops.CHILD: "#80fff0", Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e",
+               Ops.SUBSTITUTE: "#ffff00"}
 
 # VIZ API