github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

simplify late_buffer_view [pr] (#14478)

Browse Source 
check the only allowed Ops in the chain, and offset cannot be negative
This commit is contained in:
chenyu
2026-01-31 22:38:40 -05:00
committed by GitHub
parent b4f96301e0
commit da500dbe06
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tinygrad/schedule/rangeify.py
View File

@@ -275,14 +275,15 @@ def late_buffer_view(t:UOp, b:UOp):
size = prod(shape)
# walk up for the INDEX
# NOTE: even though we allow RESHAPE and SHRINK, they can combine to form non-contiguous access patterns (e.g. t[::2])
x = t
while not any(u.op is Ops.INDEX for u in x.src):
assert x.op not in GroupOp.Elementwise, "can't buffer view elementwise"
while x.op is not Ops.INDEX:
assert x.op in {Ops.BITCAST, Ops.CONTIGUOUS, Ops.SHRINK, Ops.RESHAPE}, f"unexpected op {x.op} in buffer view walk"
x = x.src[0]
x = next(u for u in x.src if u.op is Ops.INDEX)
if len(shape) == 0: offset = x.src[1].arg
else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0)
else: offset = sum(idx.vmin for idx in x.src[1:])
if offset < 0: raise RuntimeError(f"negative offset {offset} in buffer view")
return b.replace(src=(UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (size, offset), tag=t.tag), b.src[1]))