also skip casted consts in expand_broadcast

# Conflicts:
#	tinygrad/schedule/indexing.py

extra/bench_log.py

@@ -35,7 +35,7 @@ class WallTimeEvent:
     return self
   def __exit__(self, *_):
     self.time = time.monotonic() - self.start
-    _events[self.event]["wall"].append(self.time)
+    _events[self.event]["wall"].append((self.time, BENCHMARK_LOG.value))
     return False
 
 class KernelTimeEvent:
@@ -47,19 +47,19 @@ class KernelTimeEvent:
     self.start = GlobalCounters.time_sum_s
     return self
   def __exit__(self, *_):
-    _events[self.event]["kernel"].append(GlobalCounters.time_sum_s - self.start)
+    _events[self.event]["kernel"].append((GlobalCounters.time_sum_s - self.start, BENCHMARK_LOG.value))
     return False
 
 def log_event_instant(event:InstantBenchEvent, value:float):
-  _events[event].append(value)
+  _events[event].append((value, BENCHMARK_LOG.value))
 
 if BENCHMARK_LOG:
   INFLUXDB_HOST = getenv("INFLUXDB_HOST", "")
   INFLUXDB_ORG = getenv("INFLUXDB_ORG", "tiny")
   INFLUXDB_TOKEN = getenv("INFLUXDB_TOKEN", "")
 
-  def _create_point(run_id, i, attempt, ref, commit, name, value, run):
-    point = Point(BENCHMARK_LOG.value).tag("id", run_id).tag("index", i)
+  def _create_point(run_id, i, attempt, ref, commit, name, value, log_name, run):
+    point = Point(log_name.replace(':', '_').replace('.', '_')).tag("id", run_id).tag("index", i)
     point = point.tag("device", Device.DEFAULT)
     point = point.tag("attempt", attempt).tag("ref", ref).tag("commit", commit)
     point = point.field(name, value).field("x", run)
@@ -91,12 +91,12 @@ if BENCHMARK_LOG:
       run_id = str(uuid.uuid4())
       if isinstance(event, BenchEvent):
         for event_type, values in _events[event].items():
-          for i, value in enumerate(values):
-            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
+          for i, (value, log_name) in enumerate(values):
+            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, log_name, run)
             points.append(point)
       else:
-        for i, value in enumerate(_events[event]):
-          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, run)
+        for i, (value, log_name) in enumerate(_events[event]):
+          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, log_name, run)
           points.append(point)
 
     write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)

test/backend/test_ops.py

@@ -6,6 +6,7 @@ from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.renderer.isa.x86 import X86Renderer
 
 TINY_BACKEND = getenv("TINY_BACKEND")
 if TINY_BACKEND:
@@ -808,6 +809,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
     helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)
 
+  # TODO: x86 PARAM dtype fails SPEC=2
+  @Context(SPEC=1 if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else 2)
   def test_and(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)

test/backend/test_rangeify.py

@@ -1,6 +1,6 @@
 import unittest
-from tinygrad import Tensor, nn, Device, dtypes, Variable
-from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, DEBUG
+from tinygrad import Tensor, Device, dtypes, Variable
+from tinygrad.helpers import Context, GlobalCounters, getenv, DEBUG
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
 from tinygrad.codegen.opt import OptOps, Opt
 from tinygrad.renderer.ptx import PTXRenderer
@@ -14,7 +14,7 @@ class TestDoubleMatmul(unittest.TestCase):
       self.ref = (self.a @ self.b @ self.c).realize()
 
   def _test(self, opts):
-    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
+    with Context(DEBUG=max(2, DEBUG.value)):
       out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()
 
     with Context(DEBUG=0):
@@ -88,16 +88,15 @@ class TestRangeifyEdgeCase(unittest.TestCase):
     res = Tensor.cat(a, c, dim=0)
     self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
 
-  def test_pcontig_multi_gather(self):
+  def test_multi_gather(self):
     # regression test: local bufferize must have device set for const_like to work
-    with Context(PCONTIG=2):
-      # NOTE: with uint type, this will become a long and fail on WEBGPU
-      forest = Tensor(list(range(8)), dtype='int')
-      idx = Tensor([0, 0], dtype='int')
-      node_val = forest.gather(0, idx)
-      idx2 = idx * 2 + 1
-      node_val2 = forest.gather(0, idx2)
-      result = (node_val + node_val2).numpy()
+    # NOTE: with uint type, this will become a long and fail on WEBGPU
+    forest = Tensor(list(range(8)), dtype='int')
+    idx = Tensor([0, 0], dtype='int')
+    node_val = forest.gather(0, idx)
+    idx2 = idx * 2 + 1
+    node_val2 = forest.gather(0, idx2)
+    result = (node_val + node_val2).numpy()
     self.assertEqual(result.tolist(), [1, 1])
 
 if getenv("BIG") > 2:
@@ -118,65 +117,6 @@ def fa():
   GlobalCounters.reset()
   return q.scaled_dot_product_attention(k, v)
 
-def fa_bw():
-  Tensor.manual_seed(1337)
-  with Context(DEBUG=0):
-    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-    attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
-    attn_output.weight.realize()
-    target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
-
-  GlobalCounters.reset()
-  attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
-  attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
-  out = attn_output(attn)
-  loss = (out - target).square().mean()
-  loss.backward()
-  #ret = [out, Tensor.stack(q.grad, k.grad, v.grad, dim=-1)]
-  #ret = [out, Tensor.stack(q.grad, k.grad, dim=-1), v.grad]
-  ret = [out, q.grad, k.grad, v.grad]
-  Tensor.realize(*ret)
-  return ret
-
-@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
-class TestPcontig(unittest.TestCase):
-  def test_flash_attention_bw(self):
-    with Context(PCONTIG=max(2, PCONTIG.value), DEBUG=2):
-      grads = fa_bw()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-
-    with Context(PCONTIG=0, DEBUG=2):
-      cmp_grads = fa_bw()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-
-    with Context(DEBUG=0):
-      mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
-    mse = sum(mses)
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
-  def test_flash_attention(self, opts=None):
-    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
-      ret = fa().realize() if opts is None else fa().contiguous(arg=opts).realize()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-    with Context(DEBUG=2):
-      cmp = fa().realize()
-      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
-    with Context(DEBUG=0):
-      mse = ((cmp-ret)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
-  def test_flash_attention_opt(self):
-    opts = ()
-    # columns in top matrix
-    opts += (Opt(OptOps.UPCAST, 0, 4),)
-    # columns in bottom matrix
-    opts += (Opt(OptOps.UPCAST, 3, 4),)
-    # rows in all the matrix
-    opts += (Opt(OptOps.UPCAST, 4, 4),)
-    self.test_flash_attention(opts)
-
 # contiguous + reduce can support ranges?
 
 @unittest.skip("pm_rangeify no longer exists. test this in a different way")

test/testextra/test_bench_log.py

@@ -21,7 +21,7 @@ class TestBenchLog(unittest.TestCase):
     # check event list
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["wall"]), 1)
-      self.assertGreater(_events[event]["wall"][0], 0)
+      self.assertGreater(_events[event]["wall"][0][0], 0)
 
   def test_log_double_wall_time(self):
     for event in BenchEvent:
@@ -35,8 +35,8 @@ class TestBenchLog(unittest.TestCase):
     # check event list
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["wall"]), 2)
-      self.assertGreater(_events[event]["wall"][0], 0)
-      self.assertGreater(_events[event]["wall"][1], 0)
+      self.assertGreater(_events[event]["wall"][0][0], 0)
+      self.assertGreater(_events[event]["wall"][1][0], 0)
 
   @skipIf(_SKIP_KERNEL_TIMING, "ci timing is not accurate")
   def test_log_single_kernel_time(self):
@@ -52,8 +52,8 @@ class TestBenchLog(unittest.TestCase):
     # check event list
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["kernel"]), 1)
-      self.assertLess(_events[event]["kernel"][0], wall_times[0])
-      self.assertGreater(_events[event]["kernel"][0], 0)
+      self.assertLess(_events[event]["kernel"][0][0], wall_times[0])
+      self.assertGreater(_events[event]["kernel"][0][0], 0)
 
   @skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate")
   def test_interleaved_wall_kernel_time(self):
@@ -74,8 +74,8 @@ class TestBenchLog(unittest.TestCase):
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["wall"]), 1)
       self.assertEqual(len(_events[event]["kernel"]), 1)
-      self.assertLess(_events[event]["kernel"][0], wall_times[0])
-      self.assertGreater(_events[event]["kernel"][0], 0)
+      self.assertLess(_events[event]["kernel"][0][0], wall_times[0])
+      self.assertGreater(_events[event]["kernel"][0][0], 0)
 
   @skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate")
   def test_stacked_wall_kernel_time(self):
@@ -93,10 +93,10 @@ class TestBenchLog(unittest.TestCase):
     for event in BenchEvent:
       self.assertEqual(len(_events[event]["wall"]), 2)
       self.assertEqual(len(_events[event]["kernel"]), 2)
-      self.assertLess(_events[event]["kernel"][0], _events[event]["wall"][0])
-      self.assertGreater(_events[event]["kernel"][0], 0)
-      self.assertLess(_events[event]["kernel"][1], _events[event]["wall"][1])
-      self.assertGreater(_events[event]["kernel"][1], 0)
+      self.assertLess(_events[event]["kernel"][0][0], _events[event]["wall"][0][0])
+      self.assertGreater(_events[event]["kernel"][0][0], 0)
+      self.assertLess(_events[event]["kernel"][1][0], _events[event]["wall"][1][0])
+      self.assertGreater(_events[event]["kernel"][1][0], 0)
 
   def test_log_instant_event(self):
     for event in InstantBenchEvent:
@@ -105,7 +105,7 @@ class TestBenchLog(unittest.TestCase):
     # check event list
     for event in InstantBenchEvent:
       self.assertEqual(len(_events[event]), 1)
-      self.assertEqual(_events[event][0], 1000)
+      self.assertEqual(_events[event][0][0], 1000)
 
 if __name__ == '__main__':
   unittest.main()

tinygrad/engine/jit.py

@@ -7,7 +7,7 @@ from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, rewrite_group, graph_rewrite
 from tinygrad.renderer import Estimates
 from tinygrad.engine.realize import capturing, compile_linear, link_linear, run_linear, graph_cache, estimate_uop, get_runtime
-from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins
+from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_written_bufs
 from tinygrad.schedule.memory import memory_plan_rewrite, _collect_bufs
 from tinygrad.nn.state import get_parameters
 from tinygrad.uop.movement import mop_cleanup
@@ -173,13 +173,7 @@ class CapturedJit(Generic[ReturnType]):
 
   @functools.cached_property
   def _written_uops(self) -> set[UOp]:
-    out: set[UOp] = set()
-    for call in self.linear.toposort():
-      if call.op is not Ops.CALL: continue
-      arg_uops = get_call_arg_uops(call)
-      outs, ins = get_call_outs_ins(call)
-      out |= {b for k in set(outs) - set(ins) if (b:=u if (cv:=(u:=arg_uops[k]).contiguous_view()) is None else cv[0]).op is Ops.BUFFER}
-    return out
+    return {b for call in self.linear.toposort() if call.op is Ops.CALL for b in get_call_written_bufs(call)}
 
   def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:
     concrete = tuple(_copy_input(u) if u in self._written_uops else u for u in input_uops)

tinygrad/engine/realize.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
 import random, itertools, math, weakref, array, decimal
 from dataclasses import dataclass, replace, field
-from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple, tqdm
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple, tqdm, dedup
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
@@ -26,6 +26,10 @@ def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
   return (), ()
 
+def get_call_written_bufs(call:UOp) -> list[UOp]:
+  arg_uops, (outs, ins) = get_call_arg_uops(call), get_call_outs_ins(call)
+  return dedup([b for k in outs if k not in ins and (b:=u if (cv:=(u:=arg_uops[k]).contiguous_view()) is None else cv[0]).op is Ops.BUFFER])
+
 def get_call_kernels(call:UOp) -> list[tuple[str, UOp, tuple[str, Estimates, bytes]|None]]:
   if (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq":
     return [(d, call, (name, estimates, profile_key)) for devices,name,estimates,_,profile_key in call.arg.aux.kernels for d in devices]
@@ -213,14 +217,12 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   dev = cast(Any, Device[(info:= call.arg.aux).device[0]])
-  addrs = [(b.bufs[j] if isinstance(b:=_resolve(ctx.input_uops[k], ctx.input_uops).buffer, MultiBuffer) else b).get_buf(dev_name).va_addr
-           for devs, idxs in info.input_idxs for j, dev_name in enumerate(devs) for k in idxs]
+  addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(d).va_addr for d, u in info.input_addrs]
   dev.rt_buffer()._buf.cpu_view().view(offset=(base:=dev.rt_allocator.alloc(len(addrs) * 8)), fmt='Q')[:len(addrs)] = array.array('Q', addrs)
 
   if info.inputs is not None:
-    tables = [UOp.from_buffer(dev.rt_buffer().view(len(idxs), dtypes.uint64, base + j*len(idxs)*8), HCQ_RUNTIME_DEV.value)
-              for devs, idxs in info.input_idxs for j in range(len(devs))]
-    call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
+    table = UOp.from_buffer(dev.rt_buffer().view(len(info.input_addrs), dtypes.uint64, base), HCQ_RUNTIME_DEV.value)
+    call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*[table]*len(info.device))})
   exec_kernel(replace(ctx, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer()._buf.va_addr + base}), call, ast)
 
   def _prof_tm(device:str, name:str, prof:tuple[int, ...], profile_key:bytes) -> float|None:

tinygrad/function.py

@@ -46,7 +46,7 @@ class _function(Generic[ReturnType]):
     params = get_state_dict((args, kwargs), tensor_type=(Tensor, UOp)).values()
 
     # deduplicate input_uops, keeping the first occurrence index for each unique uop
-    call_uops: list[UOp] = dedup([u for t in params if (u:=t._uop).device is not None or u.is_bound_var])
+    call_uops: list[UOp] = dedup([u for t in params if (u:=t._uop).device is not None])
 
     # disable realize/schedule while this is running
     # run it and do surgery later
@@ -64,10 +64,7 @@ class _function(Generic[ReturnType]):
       raise RuntimeError(f"function return type {type(ret)} not supported")
 
     # replace the known inputs with params (using deduplicated slots)
-    def make_param(x:UOp, i:int) -> UOp:
-      p = x.param_like(i)
-      return p.replace(arg=replace(p.arg, name=f"p{i}")) if x.is_bound_var else p
-    subs = {x:make_param(x, i) for i,x in enumerate(call_uops)}
+    subs = {x:x.param_like(i) for i,x in enumerate(call_uops)}
     uret = uret.substitute(subs)
 
     # the BUFFERs that are left are the implicit inputs

tinygrad/helpers.py

@@ -271,7 +271,6 @@ PROFILE = ContextVar("PROFILE", abs(VIZ.value))
 SPEC = ContextVar("SPEC", 1)
 # TODO: disable by default due to speed
 CHECK_OOB = ContextVar("CHECK_OOB", 0)
-PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
 DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)
 # set to 1, this uses tuplize in the linearizer sort order
 TUPLE_ORDER = ContextVar("TUPLE_ORDER", 1)

tinygrad/mixin/op.py

@@ -549,7 +549,6 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     squares = (self - self.mean(axis=axis, keepdim=True)).square()
     n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
     numerator = squares.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim)
-    if resolve(n == 1, False) and correction >= 1: return self.sum(axis=axis, keepdim=keepdim).cast(output_dtype) * math.nan
     return numerator.div(smax(n - correction, 0)).cast(output_dtype)
 
   def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]:

tinygrad/renderer/__init__.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 from typing import Callable, cast
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from tinygrad.helpers import prod, Target, EMULATED_DTYPES
 from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, smin, GroupOp, PatternMatcher
 from tinygrad.dtype import AddrSpace, DType, dtypes
 from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.device import Compiler
 
+# an access takes its dtype from the buffer it indexes, so accessing at another dtype restates the storage on the buffer that owns it
+def with_storage(x:UOp, dt:DType) -> UOp:
+  if x.op in {Ops.PARAM, Ops.BUFFER}: return x.replace(dtype=None, arg=replace(x.arg, dtype=dt))
+  return x.replace(dtype=None, src=(with_storage(x.src[0], dt),)+x.src[1:])
+
 @dataclass(frozen=True)
 class Estimates:
   # number of FLOPS used in the Kernel

tinygrad/renderer/nir.py

@@ -1,7 +1,7 @@
 from typing import Callable, Any
 from tinygrad.dtype import AddrSpace, DType, dtypes, truncate
 from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target, is_image_shape, round_up
-from tinygrad.renderer import Renderer
+from tinygrad.renderer import Renderer, with_storage
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.runtime.autogen import mesa, libc
@@ -123,11 +123,12 @@ class NIRRenderer(Renderer):
   extra_matcher = PatternMatcher([
     # from ptx
     (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
-    # load/store bool -> uint8
+    # a bool is one bit in NIR but a byte in memory, so every access to a bool buffer goes through a uint8 view of it
     (UPat(Ops.LOAD, dtypes.bool, name="x"),
-     lambda x: x.replace(dtype=dtypes.uint8, src=x.src[0:1]+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())+x.src[2:]).cast(dtypes.bool)),
-    (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
-     lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:])),
+     lambda x: x.replace(dtype=None, src=(with_storage(x.src[0], dtypes.uint8),)+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())
+      +x.src[2:]).cast(dtypes.bool)),
+    (UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
+     lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), x.src[1].cast(dtypes.uint8))+x.src[2:])),
     # NIR requires shift amount to be 32 bit: https://docs.mesa3d.org/nir/alu.html#nir-alu-op-ishl
     (UPat((Ops.SHL, Ops.SHR), name="x"), lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint))) if x.src[1].dtype.bitsize != 32 else None),
     # OpConvertFToU is undefined if Result Type is not wide enough, cast through int32

tinygrad/renderer/ptx.py

@@ -4,7 +4,7 @@ from collections import defaultdict
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
 from tinygrad.dtype import dtypes, DType, AddrSpace
-from tinygrad.renderer import Renderer
+from tinygrad.renderer import Renderer, with_storage
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.helpers import flatten, prod, unwrap, Target
 
@@ -45,12 +45,12 @@ ptx_matcher = PatternMatcher([
   # upcast to float32 all the ops that don't support half
   (UPat(doesnt_support_half, dtype=dtypes.half, name="x"),
     lambda x: (UOp(x.op, src=tuple(vv.cast(dtypes.float32) for vv in x.src), arg=x.arg).cast(dtypes.half))),
-  # load/store bool -> uint8 (only for memory, not registers)
+  # a bool is a predicate register in PTX but a byte in memory, so a bool buffer is accessed through a uint8 view of it
   (UPat(Ops.LOAD, dtypes.bool, src=(UPat(name="idx"),), name="x", allow_any_len=True),
-   lambda x,idx: UOp(x.op, dtypes.uint8, x.src[0:1] + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ()) + x.src[2:]).cast(dtypes.bool) \
-     if idx.addrspace != AddrSpace.REG else None),
+   lambda x,idx: x.replace(dtype=None, src=(with_storage(idx, dtypes.uint8),) + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ())
+     + x.src[2:]).cast(dtypes.bool) if idx.addrspace != AddrSpace.REG else None),
   (UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
-   lambda x,idx: UOp(x.op, src=(x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:]) if idx.addrspace != AddrSpace.REG else None),
+   lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), x.src[1].cast(dtypes.uint8))+x.src[2:]) if idx.addrspace != AddrSpace.REG else None),
   # ptx shr and shl instructions require y to be uint
   (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
   (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),

tinygrad/runtime/support/hcq2.py

@@ -30,9 +30,8 @@ class HCQInfo:
   device:tuple[str, ...]
   estimates:Estimates = Estimates()
 
-  input_idxs:tuple[tuple[tuple[str, ...], tuple[int, ...]], ...] = () # per inputs table: (devices, indexes into input_uops)
-  inputs:int|None = None # index of the inputs table in call.src
-  # per kernel: (devices, name, estimates, timestamps, profile key)
+  inputs:int|None = None
+  input_addrs:tuple[tuple[str, UOp], ...] = () # (device, lane arg uop)
   kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...], bytes], ...] = ()
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
@@ -44,6 +43,8 @@ def unwrap_mstack(u:UOp) -> tuple[UOp, ...]:
 def unwrap_view(v:UOp) -> tuple[UOp, int]:
   return unwrap_view(v.src[0]) if v.op is Ops.BITCAST else (v.src[0], v.src[1].val) if v.op is Ops.SHRINK else (v, 0)
 
+def _lane(u:UOp, lane:int) -> UOp: return u.src[lane] if u.op is Ops.MSTACK else u.mselect(lane) if len(to_tuple(u.device)) > 1 else u
+
 # patches
 
 def is_value_known_at_link(val:UOp) -> bool:
@@ -154,11 +155,12 @@ pm_insert_copy_staging = PatternMatcher([
 class HCQDepsTracker(DepsTracker):
   @staticmethod
   def _key(buf:Any) -> tuple[Any, int, int]:
+    if isinstance(buf, UOp) and buf.op is Ops.MSELECT: buf = buf.src[0]
     return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
 
 def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]:
-  refs = get_call_arg_uops(call)
-  return [[b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs] for lane in range(len(devices))]
+  return [[b if (b:=_lane(a, lane)).op is Ops.PARAM or (b.op is Ops.MSELECT and b.src[0].op is Ops.PARAM) else b.buffer
+           for a in get_call_arg_uops(call)] for lane in range(len(devices))]
 
 def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tuple[str, ...], str, int]) -> list[tuple[tuple, int, int]]:
   dep_lanes:list[tuple[tuple, int, int]] = []
@@ -223,8 +225,8 @@ def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
   if len(calls) == 1: return calls[0]
   devs, queue = get_submit(calls[0]).src[0].arg
   body = make_submit(*[cmd for c in calls for cmd in get_submit(c).src[0].src], devs=devs, queue=queue).sink()
-  return make_call(f"submit {queue} ({len(calls)})", body,
-    replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()).simplify()))
+  return make_call(f"submit {queue} ({len(calls)})", body, replace(calls[0].arg.aux,
+    estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()).simplify()))
 
 def _merge_queues(submits:list[UOp]) -> list[UOp]:
   new_src:list[UOp] = []
@@ -325,18 +327,20 @@ def trim_link_patches(ctx:tuple[list[UOp], list[UOp]], a:UOp) -> UOp|None:
   return a.src[0].after(*kept, *[d for p in afters for d in p.src[1:]]) if links else None
 pm_trim_link_patches = PatternMatcher([(UPat(Ops.AFTER, src=(UPat((Ops.PARAM, Ops.MSTACK)),), allow_any_len=True, name="a"), trim_link_patches)])
 
-def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp, UOp], tuple[UOp, ...], dict[UOp, int]]:
+def _dnum(stride:int) -> UOp: return UOp.variable("_device_num", 0, stride - 1, dtypes.int, param=True) if stride > 1 else UOp.const(0, dtypes.int)
+
+def make_addr_table(call:UOp, gaddrs:list[UOp], name:str, stride:int=1) -> tuple[UOp, dict[UOp, UOp], tuple[UOp, ...], dict[UOp, int]]:
   bare = {g: g.replace(src=(g.src[0].without_after,)) for g in gaddrs}
 
-  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(b.tag)))
-  slots = {g:i for i,g in enumerate(order)}
-  table = UOp.placeholder((len(order),), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name)
+  # slot-major layout: slot i of lane j lives at i*stride+j, every lane reads through the same table base
+  slots = {g:i*stride for i,g in enumerate(sorted(dedup(bare.values()), key=lambda g: g.key))}
+  table = UOp.placeholder((len(slots)*stride,), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name)
 
-  reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(slots[bare[g]], dtypes.int)).load() for g in gaddrs}
-  fills = (table.after(*make_patches(table, [(i*table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else ()
+  reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(_dnum(stride) + slots[bare[g]]).load() for g in gaddrs}
+  fills = (table.after(*make_patches(table, [(i*table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots and stride == 1 else ()
   return table, reads, fills, {g:slots[bare[g]] for g in gaddrs}
 
-def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patches:list[UOp]) -> dict[UOp, UOp]:
+def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patches:list[UOp], stride:int) -> dict[UOp, UOp]:
   (dst,), words = dedup(p.buf_uop for p in patches), [(unwrap_view(p.src[0].src[0])[1] + off.val*(val.dtype.itemsize//p.buf_uop.dtype.itemsize),
                                                        slots[val]) for p in patches for off,val in zip(p.src[0].src[1].src, p.src[1].src)]
 
@@ -344,13 +348,13 @@ def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patc
   pairs = UOp.placeholder((2*len(words),), dtypes.uint32, next(UOp.unique_num), device=dst.device).rtag("systems")
   lt_patches.append(make_binary_patch(pairs, struct.pack(f'<{2*len(words)}I', *itertools.chain(*words))))
   r = UOp.range(len(words), next(UOp.unique_num), dtype=dtypes.int, src=(pairs, dst))
-  off, slot = ((pairs.index(2*r+i).load() % bound).cast(dtypes.int) for i, bound in ((0, dst.max_numel()-1), (1, table.max_numel())))
+  off, slot = ((pairs.index(2*r+i).load() % bound).cast(dtypes.int) for i, bound in ((0, dst.max_numel()-1), (1, table.max_numel()-(stride-1))))
   # SHRINK(offset, length): a const length keeps the end bound from becoming an expression the program spec rejects
   patch = UOp(Ops.SHRINK, src=(dst, off, off.const_like(table.dtype.itemsize//dst.dtype.itemsize))).bitcast(table.dtype).index(0) \
-    .store(table.index(slot).load()).end(r)
+    .store(table.index(slot + _dnum(stride)).load()).end(r)
   return {p: UOp(Ops.NOOP) for p in patches} | {patches[0]: patch}
 
-def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))
+def is_input_addr(g:UOp) -> bool: return any(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))
 
 def split_patches(call:UOp) -> UOp|None:
   rt_patches:list[UOp] = []
@@ -358,20 +362,22 @@ def split_patches(call:UOp) -> UOp|None:
   body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(rt_patches, lt_patches), name=f"trim link-time patches ({call.arg.name})")
 
   # split patches. addresses read in the body go through the tables too
+  lanes = len(to_tuple(call.arg.aux.device))
   inputs, internals = partition(dedup([g for p in rt_patches for g in get_getaddrs(p)] + get_getaddrs(body)), is_input_addr)
   runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))
-  tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
+  tables = [make_addr_table(call, gs, n, lanes if n == "inputs" else 1) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
   reads, fills = {k:v for _,r,_,_ in tables for k,v in r.items()}, [f for t in tables[1:] for f in t[2]] # inputs table is filled by exec
 
   ipatches = [p for p in rt_patches if p.tag == "inputs" and all(v in tables[0][3] for v in p.src[1].src)] # only getaddrs go to the table
-  gathers = make_gather_loop(ipatches, tables[0][0], tables[0][3], lt_patches) if ipatches else {}
+  gathers = make_gather_loop(ipatches, tables[0][0], tables[0][3], lt_patches, lanes) if ipatches else {}
   body = body.substitute({p:p.substitute(gathers | reads) for p in rt_patches}).substitute(reads)
 
   lt_srcs = collections.defaultdict(list)
   for p in lt_patches: lt_srcs[p.buf_uop].append(p)
-  return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills),
-    arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=((to_tuple(inputs[0].arg),
-      tuple(sorted(dedup(b.arg.slot for g in inputs for b in unwrap_mstack(g.buf_uop))))),) if inputs else call.arg.aux.input_idxs)))
+
+  bufs = [u for _, u in sorted(dedup([(i, g.src[0].without_after) for g, i in tables[0][3].items()]))]
+  aux = replace(call.arg.aux, input_addrs=tuple((d, _lane(u, j)) for u in bufs for j,d in enumerate(call.arg.aux.device))) if inputs else call.arg.aux
+  return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills), arg=replace(call.arg, aux=aux))
 pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
 
 # *****************
@@ -440,13 +446,12 @@ def _lane_arg(a:UOp, lane:int, table:UOp) -> UOp: return table if a.tag == "inpu
 
 def merge_batch(batch:list[UOp]) -> UOp:
   tables = UOp.variable("hcq_inputs_ptr", 0, 2**64-1, dtypes.uint64, param=True)
-  lanes = [(c, j, sum(len(idxs) * 8 for _, idxs in c.arg.aux.input_idxs)) for c in batch for j in range(len(c.arg.aux.device))] # (call, lane, bytes)
-  offs = itertools.accumulate((table_bytes for _, _, table_bytes in lanes), initial=0) # every lane owns the next table of the region
+  offs = itertools.accumulate((8 * len(c.arg.aux.input_addrs) for c in batch), initial=0) # every call owns the next table of the region
   cmds = [c.src[0].src[0].call(*[_lane_arg(a.without_after, j, tables + off) for a in c.src[1:]], UOp.variable("_device_num", 0, 1 << 30).bind(j))
-          for (c, j, _), off in zip(lanes, offs)]
+          for c, off in zip(batch, offs) for j in range(len(c.arg.aux.device))]
 
   info = HCQInfo((HCQ_RUNTIME_DEV.value,), sum((c.arg.aux.estimates for c in batch), start=Estimates()).simplify(),
-                 input_idxs=tuple(x for c in batch for x in c.arg.aux.input_idxs), kernels=tuple(k for c in batch for k in c.arg.aux.kernels))
+                 input_addrs=tuple(x for c in batch for x in c.arg.aux.input_addrs), kernels=tuple(k for c in batch for k in c.arg.aux.kernels))
   body = UOp.custom_function("hcq", make_submit(*cmds, devs=HCQ_RUNTIME_DEV.value, queue="SUBMIT:0").sink())
   return body.call(*[s for c in batch for s in c.src[1:] if s.without_after.tag != "inputs"], name=f"hcq_submitter ({len(batch)})", aux=info)
 

tinygrad/schedule/__init__.py

@@ -80,9 +80,8 @@ def create_schedule(sched_sink:UOp) -> UOp:
 
 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
-#from tinygrad.schedule.rangeify import get_kernel_graph
-from tinygrad.schedule.rangeify2 import get_kernel_graph
 from tinygrad.schedule.prepare import prepare_rangeify
+from tinygrad.schedule.rangeify import get_kernel_graph
 from tinygrad.helpers import CAPTURING
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
 from tinygrad.dtype import AddrSpace

tinygrad/schedule/allreduce.py

@@ -58,11 +58,8 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
   return UOp.usum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)
 
 def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None:
-  if output is None:
-    call_output = UOp.invalids(red.max_shape, dtype=red.dtype, device=red.device)
-    output = call_output.shrink_to(red.shape)
-  else: call_output = output
+  if output is None: output = UOp.invalids(red.shape, dtype=red.dtype, device=red.device)
   to = red.param_like(0)
   src = buf.param_like(1)
   red = src.allreduce(*red.arg)
-  return output.after(to.after(to.store(handle_allreduce(src, red))).sink().call(call_output, buf.contiguous(), name="allreduce", precompile=True))
+  return output.after(to.after(to.store(handle_allreduce(src, red))).sink().call(output, buf.contiguous(), name="allreduce", precompile=True))

tinygrad/schedule/indexing.py

@@ -2,10 +2,10 @@ from typing import Iterator
 import functools, itertools
 from dataclasses import dataclass, field, replace
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, rewrite_group, broadcast_axes
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, rewrite_group
 from tinygrad.uop.ops import gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
-from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
+from tinygrad.helpers import argsort, all_same, cpu_profile, colored, Context, SPEC, prod
 
 @dataclass
 class IndexingContext:
@@ -60,11 +60,6 @@ class BufferizeOpts:
   addrspace: AddrSpace = AddrSpace.GLOBAL
   removable: bool = True
 
-def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
-  if x.op not in GroupOp.Broadcastable: return rngs
-  baxes, nleft = broadcast_axes(src.shape, x.shape), len(x.shape)-len(src.shape)
-  return tuple(r.const_like(0) if j in baxes else r for j,r in enumerate(rngs) if j >= nleft)
-
 # TODO: srcs contain (real data srcs, something else, ranges) and the boundary is confusing. see range_start
 def data_srcs(op:Ops, src:tuple[UOp, ...]) -> tuple[UOp, ...]:
   if op in {Ops.PARAM, Ops.BUFFER, Ops.RANGE, Ops.SPECIAL}: return ()
@@ -73,13 +68,17 @@ def data_srcs(op:Ops, src:tuple[UOp, ...]) -> tuple[UOp, ...]:
   if op in GroupOp.Movement|{Ops.INDEX, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
   return src
 
+def truncate_src_rngs(rngs:tuple[UOp, ...], s:UOp) -> tuple[UOp, ...]:
+  # smaller rank srcs (like bare scalar CONSTs) don't iterate the leading ranges
+  return rngs[len(rngs)-len(s_shape):] if (s_shape:=s._shape) is not None else rngs
+
 def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
   new_srcs = []
   # shape/bound/index args that are not data src should not be indexed
   data_src_count = len(data_srcs(x.op, x.src))
   for i, s in enumerate(x.src):
     new_src = s
-    src_rngs = broadcast_rngs(x, s, ctx.range_map[x][0]) if x in ctx.range_map else ()
+    src_rngs = truncate_src_rngs(ctx.range_map[x][0], s) if x in ctx.range_map else ()
     if s.op in {Ops.PARAM, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
       if x in ctx.range_map and i < data_src_count: new_src = new_src.index(*src_rngs)
     elif s in ctx.realize_map:
@@ -202,6 +201,9 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> UOp:
 
   # explicit rangeify
   ending_ranges: dict[UOp, list[UOp]] = {}
+  # ranges ended by an EXPAND don't fire at the first elementwise op below it: that eltwise op is a single-consumer
+  # wrapper, realizing there materializes the wrapper instead of the shared value below it. movement ops forward the deferral.
+  deferred_ending: dict[UOp, list[UOp]] = {}
   for x in reversed(tsink_toposort):
     # no ranges on kernels, they are internal
     if x.op in {Ops.CALL, Ops.FUNCTION, Ops.LINEAR}: continue
@@ -213,19 +215,13 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> UOp:
     if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
 
     ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
-    # ranges the consumers iterate that this node broadcasts over
-    ended = [rctx.range_map[c][0][i] for c in consumer_map[x] if c in rctx.range_map and c.op in GroupOp.Broadcastable
-             for i in broadcast_axes(x.shape, c.shape)]
-    broadcast_ending_ranges = list(UOp.sink(*ended).ranges)
-    # fusion decision: REDUCE before the broadcast
-    if x.op is Ops.REDUCE: ending_ranges[x] += broadcast_ending_ranges
 
     # *** the ranges on the output are
     #  1. new if this op is realized
     #  2. from the single consumer if this op only has one consumer
     #  3. potentially new if this op has 2+ consumers
 
-    consumer_rngs = [broadcast_rngs(c, x, rctx.range_map[c][0]) for c in consumer_map[x] if c in rctx.range_map]
+    consumer_rngs = [truncate_src_rngs(rctx.range_map[c][0], x) for c in consumer_map[x] if c in rctx.range_map]
     if x in rctx.realize_map:
       # if this is in the realize_map, we create new ranges (at the output)
       out_rngs = tuple(rctx.new_range(s) for s in x.shape)
@@ -248,13 +244,12 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> UOp:
         local_rngs, valids = zip(*[(r.get_idx(), r.get_valid()) for r in valid_rngs])
         rngs_valids.append((local_rngs, valids))
 
-      # TODO: in RANGEIFY > 1 all_all_same isn't required
       all_all_same = all(all_same(local_rngs) for local_rngs,_ in rngs_valids)
       _out_rngs = []
       _realize_axis = []
       for i,(local_rngs,valids) in enumerate(rngs_valids):
         # we compare the ranges without their valids
-        if all_all_same or (PCONTIG and all_same(local_rngs)):
+        if all_all_same:
           # the new valid is the OR of all the children valids
           minimum_valid = UOp.const(False).usum(valids)
           _out_rngs.append(graph_rewrite(local_rngs[0].valid(minimum_valid), symbolic, name="minimum_valid"))
@@ -266,19 +261,23 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> UOp:
       # we have to (partially) realize here if there's new ranges
       if len(_realize_axis): rctx.realize_map[x] = _realize_axis
 
+    defer = set()
+    if x in deferred_ending:
+      if x.op in GroupOp.Movement: deferred_ending.setdefault(x.src[0], []).extend(deferred_ending[x])
+      elif x.op in GroupOp.Elementwise and len(consumer_map[x]) == 1 and resolve(prod(x.shape) == 1):
+        # scalar single-consumer wrappers below the EXPAND chain (like broadcasting (x * -1)) can't materialize
+        # anything useful: defer the ended ranges to the first node below that can (the shared value anchor)
+        defer = set(deferred_ending[x])
+
     # if this element is a reduce and there's ended ranges, we might have to end some other ranges
     if len(ending_ranges[x]) and x.op in GroupOp.Elementwise.union({Ops.REDUCE}):
-      _realize_axis = rctx.realize_map.get(x) or []
-      for i,r in enumerate(out_rngs):
-        if i in _realize_axis: continue
-        if not (PCONTIG > 1) or any(any(rr.arg > e.arg for e in ending_ranges[x]) for rr in r.ranges):
-          _realize_axis.append(i)
-      ending_ranges[x] = []
-      if len(_realize_axis):
-        rctx.realize_map[x] = _realize_axis
-        out_rngs = tuple([(rctx.new_range(x.shape[i]) if i in _realize_axis else r) for i,r in enumerate(out_rngs)])
-    ending_ranges[x] += broadcast_ending_ranges
-
+      firing = set(ending_ranges[x]) - defer
+      if len(firing):
+        _realize_axis = list(range(len(out_rngs)))
+        ending_ranges[x] = [r for r in ending_ranges[x] if r in defer]
+        if len(_realize_axis):
+          rctx.realize_map[x] = _realize_axis
+          out_rngs = tuple(rctx.new_range(x.shape[i]) for i in range(len(out_rngs)))
     # TODO: some ops don't have shape, enable this after the `.st` property is removed
     #assert len(out_rngs) == len(x.shape), \
     #  f"shape len mismatch {len(out_rngs)} != {len(x.shape)} on {x.op} with {len(consumer_map[x])} consumers and realize {x in realize_map}"
@@ -297,7 +296,9 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> UOp:
     # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
     # NOTE: this doesn't actually always end a range, but this is why convs are realized, so for now we need it
     if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape):
-      ending_ranges[x] += list(UOp.sink(*out_rngs[:len(x.marg)]).ranges.keys())
+      ended_here = list(UOp.sink(*out_rngs[:len(x.marg)]).ranges.keys())
+      ending_ranges[x] += ended_here
+      deferred_ending.setdefault(x.src[0], []).extend(ended_here)
 
     # REDUCE creates ranges for the axes it is reducing
     if x.op is Ops.REDUCE and x.arg[1]:

tinygrad/schedule/multi.py

@@ -33,8 +33,6 @@ replace_allreduce = PatternMatcher([
     x.mselect(0).copy_to_device(c.device) if isinstance(c.device, str) and isinstance(x.device, tuple) else None),
   # MSELECT on MSTACK is replaced with nothing
   (UPat(Ops.MSELECT, src=(UPat(Ops.MSTACK, name="mstack"),), name="ms"), lambda mstack, ms: mstack.src[ms.arg]),
-  # Identical device-less values do not need a multi wrapper.
-  (UPat(Ops.MSTACK, src=(UPat.var("x"),), allow_any_len=True, name="m"), lambda m,x: x if x.device is None and all_same(m.src) else None),
   # move shrink before MSTACK
   (UPat(Ops.SHRINK, src=(UPat(Ops.MSTACK, name="ms"),), allow_any_len=True, name="shrink"), mstack_early_shrink),
   # move MSELECT before movement/ALU ops

tinygrad/schedule/prepare.py

@@ -1,9 +1,9 @@
 import itertools
 from tinygrad.dtype import dtypes, to_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp
-from tinygrad.uop.ops import graph_rewrite, rewrite_group, shape_to_shape_arg, ParamArg, identity_element
+from tinygrad.uop.ops import graph_rewrite, rewrite_group, shape_to_shape_arg, ParamArg, identity_element, _broadcast_shape
 from tinygrad.uop.movement import mop_cleanup
-from tinygrad.helpers import prod, getenv, all_int, DEBUG, SPLIT_REDUCEOP, OPENPILOT_HACKS, FLOAT16, argsort
+from tinygrad.helpers import prod, getenv, all_int, DEBUG, SPLIT_REDUCEOP, OPENPILOT_HACKS, FLOAT16, argsort, all_same
 from tinygrad.schedule.indexing import apply_movement_op
 from tinygrad.schedule.allreduce import create_allreduce_function
 from tinygrad.schedule.multi import multi_pm
@@ -12,12 +12,6 @@ def walk_mop(u:UOp):
   if u.op in GroupOp.Movement or u.op in {Ops.INDEX, Ops.UNSHARD}: return walk_mop(u.src[0])
   return u
 
-def has_buffer_view(u:UOp) -> bool:
-  # CALL argument lowering currently passes the base allocation, so only an
-  # offset-zero contiguous view backed by a real buffer can avoid a copy.
-  if u.has_buffer_identity(after_ok=True): return True
-  return (cv:=u.contiguous_view()) is not None and cv[1] == 0 and cv[0].has_buffer_identity(after_ok=True)
-
 def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
   if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, after = x.src[0], after.cast(dtypes.float)
   while True:
@@ -124,6 +118,15 @@ def expand_bitcast(bc:UOp) -> UOp|None:
   parts = [tmp>>8*i*ns for i in range(os//ns)]
   return parts[0].stack(*parts[1:], dim=-1).flatten(-2).cast(new_uint).bitcast(bc.dtype)
 
+def expand_broadcast(x:UOp):
+  shapes = [u._shape for u in x.src]
+  if any(s is None for s in shapes) or all_same(shapes): return None
+  shape = _broadcast_shape(*shapes)
+  # don't expand CONSTs (bare or casted): scalar consts pass through rangeify as-is,
+  # and EXPAND of an Invalid const must stay a bare scalar
+  def expanded(u:UOp): return u if u.op is Ops.CONST or (u.op is Ops.CAST and u.src[0].op is Ops.CONST) else u.expand(shape)
+  return x.replace(src=tuple([expanded(u) for u in x.src]))
+
 earliest_rewrites = mop_cleanup+PatternMatcher([
   # resolve FUNCTION calls (inline the body)
   (UPat(Ops.FUNCTION, name="c"), resolve_function),
@@ -156,9 +159,9 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)),
 
   # reshaping on STORE can be a NOOP
-  #(UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
-  #                      UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
-  # lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
+  (UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
+                        UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
+   lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
 
   # ** store rules **
 
@@ -175,7 +178,13 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src"))),
    lambda target, src: target.store(src.bitcast(target.dtype))),
 
+  # expand bitcasts and broadcasts
   (UPat(Ops.BITCAST, name="bc"), expand_bitcast),
+  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), expand_broadcast),
+
+  # move RESHAPEs through MSELECT/MSTACK
+  (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
+   lambda m: m.replace(src=tuple([x.src[0].base for x in m.src])).reshape(m.shape)),
 
   # ** size 0 **
 
@@ -185,19 +194,6 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # handle size 0
   (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and 0 in x.shape else None),
 
-  # ** new prepare **
-
-  # CALL inputs need buffer identity (and to be flat)
-  (UPat(Ops.CALL, name="c"),
-   lambda c: c.replace(src=c.src[0:1]+tuple(x if has_buffer_view(x) else x.contiguous() for x in c.src[1:]))),
-
-  # MSTACK inputs need buffer identity
-  (UPat(Ops.MSTACK, name="c"),
-   lambda c: c.replace(src=tuple(x.contiguous() if not x.has_buffer_identity(after_ok=True) else x for x in c.src))),
-
-  # STORE to () is reshaped to (1,)
-  (UPat(Ops.STORE, name="s"), lambda s: s.src[0].reshape((1,)).store(s.src[1].reshape((1,))) if s.shape == () else None),
-
   # remove movement ops from SINK/AFTER. TODO: should be generic
   (UPat(Ops.SINK, name="s"), lambda s: s.replace(src=tuple(walk_mop(u) for u in s.src if u.op is not Ops.NOOP))),
   (UPat(Ops.AFTER, name="s"), lambda s: s.replace(src=(s.src[0],)+tuple(walk_mop(u) for u in s.src[1:] if u.op is not Ops.NOOP))),
@@ -215,13 +211,12 @@ def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
     return existing_buf.flatten().store(input_src)
   # create the output buffer
   buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
-  buf = buf.shrink_to(input_src.shape)
   # reshape back to input
   return buf.after(buf.store(input_src)).reshape(copy.shape)
 
 pm_copy_to_store = PatternMatcher([
   (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
-  (UPat((Ops.COPY, Ops.CONTIGUOUS), name="copy"), convert_copy_to_store),
+  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
 ])
 
 @rewrite_group(new_ctx=False)

tinygrad/schedule/rangeify.py

@@ -6,7 +6,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, K
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import prod, dedup, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
-from tinygrad.helpers import PCONTIG, partition, get_single_element
+from tinygrad.helpers import get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, apply_movement_op
@@ -83,7 +83,7 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   accessed_buffers = dedup(accessed_buffers)
 
   # if this is generated from multiple buffers, don't remove this buffer
-  if len(accessed_buffers) > 3 and not (PCONTIG > 2): return None
+  if len(accessed_buffers) > 3: return None
 
   # if any reduces access a buffer, don't remove this buffer
   buffer_in_reduce = False
@@ -94,22 +94,7 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   UOp.sink(*[x.src[0] for x in reduces]).toposort(gate=buf_gate)
   del buf_gate
   if buffer_in_reduce:
-    if PCONTIG > 2:
-      out_in_ratio = (prod(buf.shape)+1) / (sum([x.numel() for x in accessed_buffers])+1)
-      if out_in_ratio < 10: return None
-      # here we have to check the indexes, we might do a partial contig here
-      local_indexes = [x for x in indexes if x.src[0].op is Ops.STAGE and x.src[0].arg.addrspace == AddrSpace.LOCAL]
-      exclude_ranges = UOp.group(*[UOp.group(*x.src[1:]) for x in local_indexes]).ranges
-      subs = [(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST]
-      # if it's bufferized or a reduce, it's pcontig
-      is_pcontig, is_subs = partition(subs, lambda x: x[0] in exclude_ranges or any([r.arg[-1] == AxisType.REDUCE for r in x[1].ranges]))
-      if not len(is_subs):
-        return None
-      if len(is_pcontig):
-        ret = src.substitute(dict(is_subs), extra_pm=pm_gate_substitute)
-        return ret.bufferize(*[x[0] for x in is_pcontig], arg=BufferizeOpts(None, AddrSpace.LOCAL)).index(*[x[1] for x in is_pcontig])
-    else:
-      return None
+    return None
 
   # if it makes it here, the bufferize is removed
   # this is the ranges replaced

tinygrad/schedule/rangeify2.py

@@ -1,426 +0,0 @@
-from dataclasses import dataclass, field
-import itertools
-from tinygrad.dtype import AddrSpace, Invalid, strong_dtype
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, GroupOp, KernelInfo
-from tinygrad.uop.ops import graph_rewrite, AxisType, rewrite_group, remove_all_tags, resolve, shape_to_shape_arg
-from tinygrad.helpers import all_int, prod, VIZ, SPEC, Context, panic
-from tinygrad.schedule.indexing import BufferizeOpts, apply_movement_op
-from tinygrad.schedule.prepare import has_buffer_view
-from tinygrad.uop.symbolic import symbolic
-from tinygrad.codegen.simplify import pm_reduce_simplify
-
-# *** preparation ***
-
-fix_mselect_mstack = PatternMatcher([
-  # move RESHAPEs through MSELECT/MSTACK
-  (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
-   lambda m: m.replace(src=tuple([x.src[0].base for x in m.src])).reshape(m.shape)),
-])
-
-from tinygrad.helpers import all_same
-from tinygrad.uop.ops import _broadcast_shape
-
-def expand_broadcast(x:UOp):
-  shapes = [u._shape for u in x.src]
-  if any(s is None for s in shapes) or all_same(shapes): return None
-  shape = _broadcast_shape(*shapes)
-  return x.replace(src=tuple([u.expand(shape) for u in x.src]))
-
-pm_lil_prepare_graph = PatternMatcher([
-  # expand broadcasts first
-  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), expand_broadcast),
-])+fix_mselect_mstack
-
-# *** RANGE creation ***
-
-def rangeify_on_reduce(ctx, inp:UOp, red:UOp, idx:UOp|None=None):
-  if red.arg[1] == 0: return None
-  if idx is None and len(red.shape) > 0: return None
-  # TODO: is AxisType.REDUCE a real thing?
-  rngs = [UOp.range(s, next(ctx), AxisType.REDUCE) for s in inp.shape[:red.arg[1]]]
-  return inp.index(*rngs, *(idx.src[1:] if idx is not None else ())).reduce(*rngs, arg=(red.arg[0], 0))
-
-def rangeify_on_store(ctx, x:UOp):
-  if x.shape == (): return None
-  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
-  return x.src[0].index(*rngs).store(x.src[1].index(*rngs)).end(*rngs)
-
-def rangeify_on_stage(ctx, x:UOp):
-  if x.src[0].shape == (): return None
-  # size 1 dims don't get ranges, they are reshaped out and back in
-  if all_int(x.shape) and 0 < len(sq := tuple(s for s in x.shape if s != 1)) < len(x.shape):
-    return rangeify_on_stage(ctx, x.src[0].reshape(sq).bufferize(arg=x.arg)).reshape(x.shape)
-  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
-  return x.replace(src=(x.src[0].index(*rngs), *rngs))
-
-pm_range_creation = PatternMatcher([
-  # reduce/store are what creates ranges
-  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red").index(name="idx", allow_any_len=True), rangeify_on_reduce),
-  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red"), rangeify_on_reduce),
-  (UPat(Ops.STORE, name="x"), rangeify_on_store),
-  (UPat(Ops.STAGE, name="x"), rangeify_on_stage),
-])
-
-# *** RANGE migration ***
-
-# movement op on INDEX as a PatternMatcher
-def _mop_index(r:UOp, idx:UOp):
-  idxs = idx.src[1:]
-  if len(idxs) == len(r.shape):
-    ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
-    if r.op is Ops.PAD:
-      # NOTE: neither 0 or ret.const_like(0) is correct here.
-      # const_like breaks because it adds casts, and 0 is wrong if ret is a bool
-      invalid_value = UOp.const(ret.dtype.const(0))
-      # insert invalid_value for PAD with where
-      a = UOp.const(True)
-      for s in UOp.sink(*ret.src[1:]).simplify().src:
-        if s.is_invalid: return invalid_value
-        if s.op is Ops.WHERE and s.src[2].op is Ops.CONST and s.src[2].arg == Invalid: a = a & s.src[0]
-      ret = a.where(ret, invalid_value)
-    return ret
-  if r.op is Ops.RESHAPE:
-    src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
-    if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
-      if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
-      ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), dtype=idx.dtype, arg=idx.arg)
-      return ret if ret.shape == idx.shape else None
-
-# TODO: this should be in _mop_index
-def index_on_stack(stack:UOp, idx:UOp):
-  srcs = [s.index(*idx.src[2:]) for s in stack.src]
-  r0 = idx.src[1]
-  ret = srcs[-1]
-  for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret)
-  return ret
-
-pm_range_migration = PatternMatcher([
-  # STAGE on shape () is nothing
-  (UPat(Ops.STAGE, src=(UPat.var('x'),)), lambda x: x if x.shape == () else None),
-  # reshape of a single element shaped value to scalar is an index
-  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(0) if x.marg == () and x.src[0].shape == (1,) else None),
-  # handle movement ops on INDEX
-  (UPat(GroupOp.Movement, name="r").index(name="idx", allow_any_len=True), _mop_index),
-  (UPat(Ops.STACK, name="stack").index(name="idx", allow_any_len=True), index_on_stack),
-  # move movement ops and INDEX after AFTER
-  (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
-   lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)),
-  # block bitcast that changes shape
-  (UPat(Ops.BITCAST, name="b").index(allow_any_len=True),
-   lambda b: panic(RuntimeError, "shape changing bitcast not allowed in rangeify") if b.src[0].shape != b.shape else None),
-  # pass index through elementwise
-  (UPat(GroupOp.Elementwise, name="b").index(name="idx", allow_any_len=True),
-   lambda b,idx: b.replace(src=tuple(s.index(*idx.src[1:]) for s in b.src))),
-  # INDEX without src is nothing (must be at the bottom)
-  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
-])
-
-# *** split into kernels ***
-
-@dataclass
-class SplitCtx:
-  call_args:list[UOp] = field(default_factory=list)
-  buffers:dict[UOp, int] = field(default_factory=dict)
-  range_number:int = -1
-  addrspace:AddrSpace = AddrSpace.GLOBAL
-
-def _split_graph(ctx:SplitCtx, u:UOp) -> UOp|None:
-  if u.tag is not None: return None
-  if u.addrspace != ctx.addrspace: return None
-  if u.addrspace == AddrSpace.ALU: return u.param_like(-1).rtag().reshape(u.shape)
-
-  # A kernel takes each underlying buffer state once. In particular, AFTER and its buffer must use the same slot, with AFTER kept as the call
-  # argument so its dependencies are preserved.
-  key = u.buf_uop if u.op is Ops.AFTER else u
-  if (slot:=ctx.buffers.get(key)) is None:
-    slot = ctx.buffers[key] = len(ctx.call_args)
-    ctx.call_args.append(u)
-  elif u.op is Ops.AFTER:
-    ctx.call_args[slot] = u
-
-  # Parameters describe the max-sized physical allocation. A symbolic logical shape is a view of that allocation, not part of the PARAM itself.
-  param = u.param_like(slot).rtag().replace(src=(shape_to_shape_arg((u.max_numel(),)),))
-  return param.reshape(u.max_shape).shrink_to(u.shape)
-
-def _renumber_range(ctx:SplitCtx, u:UOp) -> UOp|None:
-  if u.tag is not None: return None
-  ctx.range_number += 1
-  return u.replace(arg=(ctx.range_number, u.arg[-1])).rtag()
-
-pm_split_graph = pm_range_migration+PatternMatcher([
-  (UPat((Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK), name="u"), _split_graph),
-  (UPat(Ops.RANGE, name="u"), _renumber_range),
-])
-
-def _is_fully_invalid_state(x:UOp) -> bool:
-  while x.op in GroupOp.Movement|{Ops.INDEX}: x = x.src[0]
-  if x.op is not Ops.AFTER or len(x.src) != 2: return False
-  end = x.src[1]
-  st = end.src[0] if end.op is Ops.END else end
-  if st.op is not Ops.STORE or not st.src[1].base.is_invalid: return False
-  if end.op is not Ops.END: return st.src[0].max_numel() == x.max_numel()
-  covered = 1
-  for r in end.src[1:]:
-    if r.op is not Ops.RANGE or r.src[0].op is not Ops.CONST or not isinstance(r.src[0].val, int): return False
-    covered *= r.src[0].val
-  return covered == x.max_numel()
-
-def split_store(x:UOp) -> UOp|None:
-  st = x.src[0] if x.op is Ops.END else x
-  if st.op is Ops.STORE and st.src[0].is_variable: return None
-  if st.op is Ops.STORE and st.src[0] is st.src[1]: return UOp(Ops.NOOP)
-  # A directly-invalid value makes this store a no-op. An AFTER carrying an
-  # invalid partial store is still a valid buffer state: uncovered elements
-  # must continue to read from the previous state.
-  if st.op is Ops.STORE and (st.src[1].base.is_invalid or _is_fully_invalid_state(st.src[1])): return UOp(Ops.NOOP)
-  ret = graph_rewrite(x, pm_split_graph, ctx:=SplitCtx(), name="split kernel", bottom_up=True)
-  # TODO: params and args should be able to be in any order
-  ctx.addrspace = AddrSpace.ALU
-  ret = graph_rewrite(ret, pm_split_graph, ctx, name="split kernel (vars)", bottom_up=True)
-  ret = graph_rewrite(ret, remove_all_tags, name="remove split tags", bottom_up=True)
-  return ret.sink(arg=KernelInfo()).call(*ctx.call_args)
-
-split_kernels = PatternMatcher([
-  (UPat((Ops.STORE, Ops.END), name="x"), split_store),
-])
-
-# cleanups
-
-def strip_zero_offset_shrink(x:UOp) -> UOp:
-  return x.src[0] if x.op is Ops.SHRINK and all(resolve(start == 0, False) for start,_ in x.marg) else x
-
-def no_indexing_calls(u:UOp):
-  new_srcs = []
-  for x in u.src:
-    if x.op is Ops.INDEX:
-      # sometimes if call srcs have children the call will get an INDEX. we remove it here.
-      # TODO: we should add safety checks here for contiguous
-      new_srcs.append(x.src[0])
-    elif x.op is Ops.SHRINK:
-      # SHRINK with offset 0 is fine
-      new_srcs.append(strip_zero_offset_shrink(x))
-    elif x.op is Ops.MSTACK:
-      new_srcs.append(x.replace(src=tuple(strip_zero_offset_shrink(s) for s in x.src)))
-    else:
-      # everything else we pass through
-      new_srcs.append(x)
-  return u.replace(src=tuple(new_srcs))
-
-pm_no_indexing_calls = PatternMatcher([
-  (UPat(Ops.CALL, name="u"), no_indexing_calls),
-  (UPat(Ops.AFTER, name="u"), lambda u: u.replace(src=tuple(s for s in u.src if s.op is not Ops.NOOP))),
-])
-
-# *** main rangeify ***
-
-debug_tag_factor = PatternMatcher([
-  (UPat(GroupOp.All, name="x"), lambda ctx,x: x.rtag(ctx[0][x] if x not in ctx[1] else 'REAL') if x.tag is None else None),
-])
-
-def remove_stage(ctx, x:UOp) -> UOp:
-  dtype = strong_dtype(x.dtype)
-  buf = UOp.new_buffer(x.arg.device, x.max_numel(), dtype, num=next(ctx))
-  val = x.src[0] if x.src[0].dtype == dtype else x.src[0].cast(dtype)
-  return buf.after(buf.reshape(x.shape).index(*x.src[1:]).store(val).end(*x.src[1:])).reshape(x.shape)
-
-pm_remove_stage = PatternMatcher([
-  (UPat(Ops.STAGE, name="x"), remove_stage),
-])+fix_mselect_mstack
-
-def remove_selected_stage(ctx:set[UOp], stage:UOp, idx:UOp) -> UOp|None:
-  return stage.src[0] if stage in ctx and stage.src[1:] == idx.src[1:] else None
-
-pm_remove_selected_stage = PatternMatcher([
-  (UPat(Ops.STAGE, name="stage").index(name="idx", allow_any_len=True), remove_selected_stage),
-])
-
-def inline_stage_index(stage:UOp, idx:UOp) -> UOp:
-  replacements, cache = dict(zip(stage.src[1:], idx.src[1:])), {}
-  def replace(x:UOp) -> UOp:
-    if x in replacements: return replacements[x]
-    if x.has_buffer_identity(after_ok=True) or x.op is Ops.STAGE: return x
-    if x not in cache: cache[x] = x.replace(src=tuple(replace(s) for s in x.src))
-    return cache[x]
-  return replace(stage.src[0])
-
-MAX_RECOMPUTE = 8
-MAX_SCALAR_RECOMPUTE = 64
-
-def recompute_cost(x:UOp, seen:set[UOp]|None=None) -> int|None:
-  if seen is None: seen = set()
-  if x in seen or x.op is Ops.STAGE or x.has_buffer_identity(after_ok=True): return 0
-  seen.add(x)
-  if x.op is Ops.REDUCE: return None
-  costs = [recompute_cost(s, seen) for s in (x.src[:1] if x.op is Ops.INDEX else x.src)]
-  return None if any(c is None for c in costs) else sum(c for c in costs if c is not None) + (x.op in GroupOp.Elementwise)
-
-def materialize_call_args(c:UOp) -> UOp:
-  srcs:list[UOp] = []
-  for x in c.src[1:]:
-    device = x.device or c.device
-    srcs.append(x if x.op is Ops.STAGE or x.is_bound_var or has_buffer_view(x) or x.shape == () or device is None
-                else x.bufferize(arg=BufferizeOpts(device=device)))
-  return c.replace(src=(c.src[0], *srcs))
-
-def materialize_mselect(m:UOp, x:UOp) -> UOp|None:
-  if x.device is None or x.op is Ops.STAGE or (x.op not in GroupOp.ALU and x.has_buffer_identity(after_ok=True)): return None
-  return m.replace(src=(x.bufferize(arg=BufferizeOpts(device=x.device)),))
-
-pm_materialize_call_args = PatternMatcher([
-  (UPat(Ops.CALL, name="c"), materialize_call_args),
-  (UPat(Ops.MSELECT, src=(UPat(name="x"),), name="m"), materialize_mselect),
-])
-
-@rewrite_group(new_ctx=False)
-def get_kernel_graph(sink:UOp) -> UOp:
-  tsink = graph_rewrite(sink, pm_lil_prepare_graph, bottom_up=True, name="prepare graph")
-  # Calls can only receive buffer states. Materialize lazy constants/computations instead of silently unwrapping them to a nonexistent base buffer.
-  tsink = graph_rewrite(tsink, pm_materialize_call_args, name="materialize call args")
-
-  read_cache:dict[UOp, set[UOp]] = {}
-  def read_buffers(x:UOp) -> set[UOp]:
-    if x not in read_cache:
-      read_cache[x] = {x.buf_uop} if x.has_buffer_identity(after_ok=True) else set().union(*(read_buffers(s) for s in x.src))
-    return read_cache[x]
-  stores = [u for u in tsink.toposort() if u.op is Ops.STORE and not u.src[0].is_variable]
-  dests = [u.src[0].buf_uop for u in stores]
-  reads = [read_buffers(u.src[1]) for u in stores]
-  force_stage:set[UOp] = set()
-  param_writes = [i for i,dest in enumerate(dests) if dest.op is Ops.PARAM]
-  if len(param_writes) == 2:
-    i, j = param_writes
-    if stores[j].src[1].op_in_backward_slice_with_self(Ops.REDUCE) and dests[i] is not dests[j] and \
-       dests[i] in reads[j] and dests[j] in reads[i]: force_stage.add(stores[j].src[1])
-
-  # add safe STAGEs to never duplicate compute
-  # we compute the number of times a buffer is consumed. if > 1, we realize
-  realize = {}
-  consumes = {tsink:0}
-  for u in reversed(tsink.toposort()):
-    assert u in consumes, f"{u.op} not in consumes"
-    if u in force_stage:
-      realize[u] = u.rtag(1).bufferize(arg=BufferizeOpts(device=u.device, removable=False))
-      consumes[u] = 1
-    elif (u.op in GroupOp.ALU or u.op is Ops.REDUCE) and consumes[u] > 1 and u.device is not None:
-      # TODO: rename to stage
-      realize[u] = u.rtag(1).bufferize(arg=BufferizeOpts(device=u.device))
-      consumes[u] = 1
-    if u.op is Ops.STORE: consumes[u] = 1
-    if u.op is Ops.EXPAND: consumes[u] *= u.max_numel() // u.src[0].max_numel()
-    for i,s in enumerate(u.src):
-      if s not in consumes: consumes[s] = 0
-      if u.op is not Ops.STORE or i > 0:
-        consumes[s] += consumes[u]
-  if VIZ:
-    with Context(TRACK_MATCH_STATS=0): ctags = graph_rewrite(tsink, debug_tag_factor, ctx=(consumes, realize), bottom_up=True)
-    graph_rewrite(ctags, PatternMatcher([]), name="View Consumes")
-
-  # add stages
-  tsink = graph_rewrite(tsink.substitute(realize), remove_all_tags, name="untag")
-
-  # simple rangeify
-  tsink = graph_rewrite(tsink, pm_range_creation+pm_range_migration, ctx=itertools.count(0), bottom_up=True, name="simple rangeify")
-
-  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
-
-  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify, name="pre-fusion reduce simplify")
-
-  # remove stage boundaries when this doesn't duplicate expensive compute or nest reductions
-  while 1:
-    staged:dict[UOp, list[UOp]] = {}
-    children:dict[UOp, list[UOp]] = {}
-    for u in tsink.toposort():
-      for s in u.src: children.setdefault(s, []).append(u)
-      if u.op is Ops.INDEX and u.src[0].op is Ops.STAGE: staged.setdefault(u.src[0], []).append(u)
-
-    boundary_cache:dict[UOp, set[UOp]] = {}
-    def boundaries(x:UOp) -> set[UOp]:
-      if x not in boundary_cache:
-        boundary_cache[x] = set().union(*({c} if c.op in {Ops.STAGE, Ops.STORE, Ops.CALL} else boundaries(c)
-                                          for c in children.get(x, [])))
-      return boundary_cache[x]
-
-    reduce_cache:dict[UOp, bool] = {}
-    def feeds_reduce(x:UOp) -> bool:
-      if x not in reduce_cache:
-        reduce_cache[x] = any(c.op is Ops.REDUCE or (c.op not in {Ops.STAGE, Ops.STORE, Ops.CALL} and feeds_reduce(c))
-                              for c in children.get(x, []))
-      return reduce_cache[x]
-
-    def boundary_work(boundary:UOp) -> int|None:
-      value = boundary.src[1] if boundary.op is Ops.STORE else boundary.src[0]
-      if any(r.src[0].op is not Ops.CONST or not isinstance(r.src[0].val, int) for r in value.ranges): return None
-      return prod(r.src[0].val for r in value.ranges)
-
-    def recompute_work(idxs:list[UOp]) -> int|None:
-      works = [boundary_work(next(iter(bs))) for idx in idxs if len(bs:=boundaries(idx)) == 1]
-      return sum(x for x in works if x is not None) if len(works) == len(idxs) and all(x is not None for x in works) else None
-
-    replacements:dict[UOp, UOp] = {}
-    range_replacements:dict[UOp, UOp] = {}
-    selected_stages:set[UOp] = set()
-    for stage,idxs in staged.items():
-      if not stage.arg.removable or children.get(stage) != idxs: continue
-      inlinable = stage.src[0].op in GroupOp.ALU or stage.src[0].op is Ops.REDUCE
-      cost = recompute_cost(stage.src[0])
-
-      # passthrough stages don't duplicate compute when indexing is unchanged
-      if not inlinable:
-        if len(idxs) == 1 and stage.src[1:] == idxs[0].src[1:]:
-          replacements[idxs[0]] = stage.src[0]
-          break
-        continue
-
-      # duplicate cheap elementwise stages; reductions require identical indexing into one output boundary
-      if len(idxs) > 1:
-        work = recompute_work(idxs)
-        if cost is not None and cost <= MAX_RECOMPUTE and work is not None and work <= stage.max_numel() * len(idxs):
-          replacements.update((idx, inline_stage_index(stage, idx)) for idx in idxs)
-        elif cost is None and all(idx.src[1:] == idxs[0].src[1:] for idx in idxs) and not any(feeds_reduce(idx) for idx in idxs):
-          stage_boundaries = set().union(*(boundaries(idx) for idx in idxs))
-          if len(stage_boundaries) == 1 and boundary_work(next(iter(stage_boundaries))) == stage.max_numel():
-            range_replacements.update(zip(stage.src[1:], idxs[0].src[1:]))
-            selected_stages.add(stage)
-        if replacements or range_replacements: break
-        continue
-
-      idx = idxs[0]
-      stage_boundaries = boundaries(idx)
-      work = boundary_work(next(iter(stage_boundaries))) if len(stage_boundaries) == 1 else None
-      scalar = stage.max_numel() == 1 and cost is not None and cost <= MAX_SCALAR_RECOMPUTE and all(r.op is Ops.CONST for r in idx.src[1:])
-      small = cost is not None and cost <= MAX_SCALAR_RECOMPUTE and stage.max_numel() <= 8 and idx.max_numel() <= 8
-      if cost is not None:
-        if stage.src[0].op_in_backward_slice_with_self(Ops.THREEFRY) or scalar or small or \
-           (cost <= MAX_RECOMPUTE and work is not None and work <= stage.max_numel()):
-          replacements[idx] = inline_stage_index(stage, idx)
-      elif len(stage_boundaries) == 1 and not feeds_reduce(idx) and work == stage.max_numel():
-        if all(r.op is Ops.RANGE for r in idx.src[1:]):
-          range_replacements.update(zip(stage.src[1:], idx.src[1:]))
-          selected_stages.add(stage)
-        else: replacements[idx] = inline_stage_index(stage, idx)
-      if replacements or range_replacements: break
-
-    if not replacements and not range_replacements: break
-    tsink = tsink.substitute(replacements).substitute(range_replacements)
-    if selected_stages:
-      selected_stages = {stage.substitute(range_replacements) for stage in selected_stages}
-      tsink = graph_rewrite(tsink, pm_remove_selected_stage, ctx=selected_stages, name="remove selected stage")
-
-  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify, name="reduce simplify")
-
-  # ***** MERGING AND SPLITTING (should be totally optional) *****
-
-  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Merged Rangeify")
-
-  next_buffer_num = itertools.count(1000)
-  tsink = graph_rewrite(tsink, symbolic+pm_remove_stage, ctx=next_buffer_num, bottom_up=True, name="remove stage")
-  tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
-  tsink = graph_rewrite(tsink, pm_no_indexing_calls, name="remove indexing from call args")
-
-  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
-  if SPEC:
-    # validate the kernel graph
-    from tinygrad.uop.spec import type_verify, spec_kernel_graph
-    type_verify(tsink, spec_kernel_graph, enter_calls=False)
-  return tsink
-

tinygrad/tensor.py

@@ -114,8 +114,7 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
   # add the outputs to the call
   srcs = c.src[0].src
   resolved = [c.gettuple(i) for i in range(len(srcs))]
-  # CALL outputs are max-sized physical buffers. Keep symbolic shapes as views so writable arguments and returned buffers stay identical.
-  outs = tuple(r.pad_to(r.max_shape).empty_like() for r in resolved)
+  outs = tuple(r.empty_like() for r in resolved)
   targets = [o.param_like(len(c.src)-1+i).shrink_to(s.shape) for i,(o,s) in enumerate(zip(outs, srcs))]
 
   subs:dict[UOp, UOp] = {}
@@ -134,8 +133,11 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
 
   # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION)
   new_call = UOp(Ops.CALL, src=(fxn, *input_buffers, *outs), arg=c.arg)
-  # NOTE: use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes.
-  rets = tuple(o.after(new_call).shrink_to(rs.shape) for o,rs in zip(outs, resolved))
+  rets = tuple(o.after(new_call) for o in outs)
+
+  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
+  # NOTE: must use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes
+  rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, resolved))
 
   return UOp.maketuple(*rets)
 
@@ -218,32 +220,6 @@ pm_replace_buf = PatternMatcher([
   (UPat(Ops.AFTER, name="b"), lambda ctx,b: replace_input_buffer(ctx, b) if b.is_bound_var else None),
 ])
 
-def _check_state_cycles(sink:UOp):
-  # Track only the first storage state on each path. Combining a raw read with an assigned state is stale, while distinct AFTERs can be valid
-  # independent snapshots (for example RNG state). Memoizing the two state sets avoids repeatedly walking large training graphs.
-  states:dict[UOp, tuple[frozenset[UOp], frozenset[UOp]]] = {}
-  state:tuple[frozenset[UOp], frozenset[UOp]]
-  for u in sink.toposort(enter_calls=False):
-    if u.op is Ops.BUFFER and u.addrspace == AddrSpace.GLOBAL: state = (frozenset((u,)), frozenset())
-    elif u.op is Ops.AFTER and u.addrspace == AddrSpace.GLOBAL:
-      key = u.buf_uop
-      stores = [x for x in u.src[1:] if x.op is Ops.STORE and x.src[0].buf_uop is key]
-      # Ordering dependencies can STORE to another buffer without changing this state. Self-dependent updates continue the existing state lineage;
-      # only a write independent of the old value creates a conflicting state.
-      self_update = any(key in states[x.src[1]][0] or key in states[x.src[1]][1] for x in stores)
-      state = states[u.src[0]] if not stores or self_update else (frozenset(), frozenset((key,)))
-    else:
-      srcs = u.src[1:] if u.op in {Ops.CALL, Ops.FUNCTION} else u.src
-      raw = frozenset().union(*(states[x][0] for x in srcs))
-      assigned = frozenset().union(*(states[x][1] for x in srcs))
-      # CONTIGUOUS snapshots stale raw reads before a pending assignment, but cannot make a post-assignment read happen earlier.
-      state = (frozenset() if u.op is Ops.CONTIGUOUS else raw, assigned)
-    states[u] = state
-    if u.op in GroupOp.ALU:
-      branches = [states[x] for x in u.src]
-      if any((a[0] & b[1]) or (a[1] & b[0]) for i,a in enumerate(branches) for b in branches[i+1:]):
-        raise RuntimeError("cycle detected while combining buffer states")
-
 @rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}")
 def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
@@ -430,9 +406,7 @@ class Tensor(RandMixin):
     # weakness ends where storage begins
     if any(t.dtype in dtypes.weaks and t.uop.device is not None for t in (self,)+lst):
       raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
-    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
-    _check_state_cycles(big_sink)
-    big_sink, becomes_map = transform_to_call(big_sink)
+    big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
     _apply_map_to_tensors(becomes_map, name="buffers")
     return create_linear_with_vars(big_sink)
 

tinygrad/uop/ops.py

@@ -1173,13 +1173,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
     return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
   def param_like(self, slot:int):
-    # if it's a PARAM or BUFFER, we just replace the slot
-    buf = self
-    while buf.op is Ops.AFTER: buf = buf.src[0]
-    if buf.op in {Ops.PARAM, Ops.BUFFER}: return UOp(Ops.PARAM, src=buf.src, arg=replace(buf.arg, slot=slot))
-    # otherwise we create a new param
-    addrspace = buf.addrspace if buf.addrspace is not None else AddrSpace.GLOBAL
-    return UOp.param(slot, buf.dtype, buf.shard_shape if buf.axis is not None else buf._shape, self.device, addrspace=addrspace, axis=buf.axis)
+    # Variables become ALU params in the call body; the stored value (if bound) stays in the call args
+    if self.is_bound_var or self.is_variable:
+      b = self.src[0] if self.op is Ops.AFTER else self
+      return UOp(Ops.PARAM, src=b.src, arg=replace(b.arg, slot=slot, name=f"p{slot}"))
+    addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
+    return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)
 
   @staticmethod
   def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)

tinygrad/uop/spec.py

@@ -257,18 +257,12 @@ spec_kernel_graph = PatternMatcher([
   # const + stack to make vconsts and shape args. a 0-size/bound reduce keeps its const casted
   (UPat(Ops.CONST, src=()), lambda: True),
   (UPat(Ops.CAST, src=(UPat(Ops.CONST, src=()),)), lambda: True),
-  # symbolic shape expressions can remain outside kernels (for example, flattening (n, 4) produces n*4)
-  (UPat(GroupOp.ALU, name="x"), lambda x: x.dtype in dtypes.ints+(dtypes.weakint,) or None),
   (UPat(Ops.STACK, name="s"), lambda s: all(x.op in (Ops.CONST, Ops.PARAM) or x.is_variable or x.is_bound_var for x in s.src) or None),
   # linear for more kernels (TODO: we should enter non sink calls)
   #(UPat(Ops.LINEAR), lambda: True),
   # param is outside buffer, buffer is local buffer
   (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
   (UPat(Ops.BUFFER, name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.GLOBAL, AddrSpace.ALU)),
-  # indexing/movement views on kernel buffers/call results are allowed to carry symbolic shape expressions
-  (UPat({Ops.INDEX}|GroupOp.Movement,
-        src=(UPat({Ops.INDEX}|GroupOp.Movement|{Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST}),),
-        allow_any_len=True), lambda: True),
   # RESHAPE/BITCAST are NOOPs in the kernel graph (do we need them?)
   (UPat((Ops.RESHAPE, Ops.BITCAST)), lambda: True),
   # mstack/mselect