geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

slice is just shrink (#17483)

Browse Source
This commit is contained in:
sirhcm
2026-08-10 23:37:17 -04:00
committed by GitHub
parent ad2fdeae69
commit 0c6a2c7dd6
5 changed files with 52 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/engine/realize.py
+1 -1
View File
@@ -140,7 +140,7 @@ class ExecContext:
cache: bool = True
def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
if b.op in (Ops.SLICE, Ops.MSELECT, Ops.SHRINK) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
return inputs[b.arg.slot] if b.op is Ops.PARAM else b
def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
tinygrad/runtime/graph/metal.py
+1 -1
View File
@@ -113,5 +113,5 @@ class MetalGraph(GraphRunner):
@staticmethod
def supports_uop(batch_devs, new_call:UOp) -> bool:
# Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
if any(b.op is Ops.SLICE and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
if any(b.op in {Ops.SLICE, Ops.SHRINK} and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
return GraphRunner.supports_uop(batch_devs, new_call)
tinygrad/runtime/support/hcq2.py
+1 -1
View File
@@ -383,7 +383,7 @@ def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64)
pm_early_simplify = PatternMatcher([
(UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
(UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat((Ops.SLICE, Ops.SHRINK), name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
(UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"),), allow_any_len=True, name="x"),
lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))),
])
tinygrad/tensor.py
+26 -24
View File
@@ -23,6 +23,7 @@ class AllocCtx:
bases: set[UOp] = field(default_factory=set)
assigns: list[UOp] = field(default_factory=list)
replacements: list[UOp] = field(default_factory=list)
views: set[UOp] = field(default_factory=set)
def tag_uop(ctx:AllocCtx, x:UOp):
if x.tag is not None: return None
@@ -63,40 +64,37 @@ def replace_contig_with_store_after(u:UOp):
def replace_store_after_with_contig(u:UOp, src:UOp):
assigned_to = u
while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: assigned_to = assigned_to.src[0].base
if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag)
if assigned_to.op is not Ops.BUFFER: return src.contiguous(tag=u.tag)
def _make_buffer_view(src:UOp) -> UOp|None:
"""If movement ops on src collapse to a contiguous range, return SLICE. Otherwise None."""
if (offset := src.contiguous_view_offset()) is None: return None
buf = src.base
if buf.op is Ops.SLICE:
byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize
buf = buf.src[0]
if byte_offset % buf.dtype.itemsize != 0: return None
offset = byte_offset // buf.dtype.itemsize
return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel())
if (cv := src.contiguous_view()) is None: return None
(buf, offset), size = cv, src.max_numel() * src.element_size() // cv[0].element_size()
if buf.op is not Ops.BUFFER: return None
# NB: make offset a UOp.variable here to do the offset computation in the kernels
return buf[offset:offset+size].bitcast(src.dtype)
def contiguous_mops_to_view(c:UOp, src:UOp):
"""MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""
def contiguous_mops_to_view(ctx:AllocCtx, c:UOp, src:UOp):
"""MOPS(BUFFER) → SHRINK when movement ops collapse to a contiguous range."""
buf = src.base
if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.UNSHARD}: return None
if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None
if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None
while buf.op is Ops.BITCAST: buf = buf.src[0].base
if buf.op not in {Ops.BUFFER, Ops.UNSHARD}: return None
# no symbolic shape
if not all_int(c.shape): return None
if buf.op is not Ops.UNSHARD and (view := _make_buffer_view(src)) is not None:
view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
return c.replace(src=(view,)) if c.op is Ops.COPY else view
ctx.views.add(view)
view = view.reshape(c.shape)
return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view
# for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
# for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SHRINK on the resolved result
if not isinstance(c.device, str):
from tinygrad.schedule.multi import multi_pm
resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view")
if resolved.op is not Ops.UNSHARD: return None
if (view := _make_buffer_view(resolved.src[0])) is None: return None
return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag)
ctx.views.add(view)
return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:])
return None
@@ -151,8 +149,9 @@ pm_early_transform_tensor_graph = PatternMatcher([
# resolve TUPLE+GETTUPLE (for precompiled calls)
(UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
# fold MOPS+BITCAST over BUFFER/SLICE into SLICE when movement ops collapse to contiguous range
(UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), contiguous_mops_to_view),
# fold MOPS+BITCAST over BUFFER into SHRINK when movement ops collapse to contiguous range
(UPat((Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), name="c"), contiguous_mops_to_view),
(UPat(Ops.STORE, src=(UPat(Ops.BITCAST, name="src"), UPat()), name="c", allow_any_len=True), contiguous_mops_to_view),
# remove contiguous on movement ops before a copy on disk
(UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy:
@@ -201,6 +200,8 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
def replace_input_view(ctx:AllocCtx, b:UOp): return replace_input_buffer(ctx, b) if b in ctx.views else None
pm_finalize_call = PatternMatcher([
(UPat(Ops.AFTER, name="x"), finalize_after),
(UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
@@ -210,8 +211,9 @@ pm_replace_buf = PatternMatcher([
# replace BUFFER with PARAM for cache key normalization
(UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
# replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
(UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer),
# replace SHRINK with PARAM
(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), name="b", allow_any_len=True), replace_input_view),
(UPat(Ops.BITCAST, src=(UPat.any(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), allow_any_len=True), UPat(Ops.BUFFER)),), name="b"), replace_input_view),
# strip value from BIND for cache key normalization, so different values hit same cache
(UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer),
])
@@ -229,7 +231,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops")
# here we can break the tensor graph. this is the only place you need to maintain numbered tags
big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx=ctx, name="early transform tensor graph")
# here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys
graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")
tinygrad/uop/ops.py
+23 -19
View File
@@ -822,7 +822,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
unique_num = itertools.count(0)
def getaddr(self, device=None) -> UOp:
if self.without_after.op not in {Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
if self.without_after.op not in {Ops.BUFFER, Ops.SLICE, Ops.SHRINK, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
return UOp(Ops.GETADDR, src=(self,), arg=device or to_tuple(self.device)[0])
@staticmethod
def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
@@ -901,8 +901,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
while len(s.src) and s.op not in {Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.MSTACK}: s = s.src[0]
return s
def contiguous_view_offset(self) -> int|None:
"""If movement ops on a BUFFER collapse to a contiguous range, return `offset` in elements. Otherwise None."""
def contiguous_view(self) -> tuple[UOp, int]|None:
from tinygrad.schedule.rangeify import pm_mops
from tinygrad.uop.symbolic import symbolic
@@ -915,7 +914,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
idx = self.flatten().index(UOp.range(self.numel(), 0))
out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset")
return out.val if out.op is Ops.CONST and isinstance(out.val, int) else None
if out.op is not Ops.INDEX or not (b:=out.src[0]).tag or (c:=out.src[1]).op is not Ops.CONST or not isinstance(c.val, int): return None
return b.rtag(None), c.val
def contiguous_view_offset(self) -> int|None: return None if (view := self.contiguous_view()) is None else view[1]
def has_buffer_identity(self, after_ok=False):
"""Check if this UOp has a concrete buffer identity in the graph (RESHAPE/UNSHARD -> BUFFER chain)."""
@@ -932,18 +934,16 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
@property
def buffer(self) -> Buffer|MultiBuffer:
if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
if self.op in {Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
# this buffer can process disk tensors and simple movement ops
if self is not self.base:
buf = self.base.buffer
assert isinstance(buf, Buffer), "must be a Buffer for movement ops"
offset = self.contiguous_view_offset()
if offset is None: raise RuntimeError(f"non-contiguous view is not supported for {buf.device} buffer")
return buf.view(prod(self.max_shape), self.dtype, offset*self.dtype.itemsize)
if self.op is Ops.BITCAST:
buf = self.src[0].buffer
assert isinstance(buf, Buffer), "must be a Buffer for BITCAST"
return buf.view(prod(self.max_shape), self.dtype, 0)
if self is not self.base or self.op is Ops.BITCAST:
if (cv := self.contiguous_view()) is None: raise RuntimeError(f"non-contiguous view is not supported for {self.device} buffer")
buf, offset = (b:=cv[0]).base.buffer, cv[1]
if isinstance(buf, MultiBuffer):
mbuf = MultiBuffer.__new__(MultiBuffer)
mbuf.bufs = [x.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize) for x in buf.bufs]
return mbuf
return buf.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize)
if self.op is Ops.SLICE:
if (cret:=buffers.get(self)) is not None: return cret
buf = self.src[0].buffer
@@ -1775,10 +1775,14 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])
# ctx is source UOp for which we are finding a contiguous view for. used in contiguous_view_offset
pm_contiguous_view_offset = PatternMatcher([
(UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(0)),
(UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(0)),
(UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda c: c),
(UPat(Ops.INDEX, src=(UPat(), UPat.cvar('c'))), lambda ctx, c: c if resolve(ctx.numel() == 1, False) else None),
# normalize to 1d bitcasts
(UPat(Ops.BITCAST, name="b"), lambda b: b.src[0].flatten().bitcast(b.dtype).reshape(b.shape) if len(b.shape) != 1 else None),
(UPat(Ops.BITCAST, name="b").index(UPat.cvar("c")), lambda ctx, b, c:
b.src[0].flatten().index(UOp.range(ctx.numel() * (osz:=b.element_size())//(isz:=b.src[0].element_size()), 0) + (c * osz//isz)) if b.tag else None),
(UPat(Ops.INDEX, src=(UPat.var("b"),)), lambda b: b.rtag().index(0)),
(UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE))), lambda b: b.rtag().index(0)),
(UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c)),
(UPat(Ops.INDEX, src=(UPat.var("b"), UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c) if resolve(ctx.numel() == 1, False) else None),
])
# *** what was symbolic.py ***