remove noops

* use recursive_property to prevent RecursionError

* not slower

* fix tests

* faster

* simpler

.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/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

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/test_assign.py

@@ -130,6 +130,11 @@ class TestAssign(unittest.TestCase):
   @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_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_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])

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
@@ -17,7 +17,7 @@ from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_in
   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"))
 

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/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/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

@@ -65,11 +65,21 @@ 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),
 
+  # 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"),
    lambda x,target,assign: x.f(Ops.CONTIGUOUS, tag=assign.tag) if ((t:=target.base).op is not Ops.BUFFER and \
@@ -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,7 +378,7 @@ 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),
 ])
 
 # *****************
@@ -750,7 +756,9 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   tsink = graph_rewrite(tsink, pm_rangeify, ctx=(rangeify_ctx:=RangeifyContext()), bottom_up=True, name="rangeify")
   # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
   tsink = graph_rewrite(tsink, symbolic_simple+pm_reduce_unparented, name="symbolic")  # this supports const folding
-  tsink = graph_rewrite(tsink, pm_cleanups+pm_substitute_recurse, bottom_up=True, name="remove costly buffers")
+  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

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
@@ -20,6 +21,7 @@ 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

@@ -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)
@@ -162,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}:
@@ -187,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
@@ -1023,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:
@@ -1035,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
@@ -1084,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