sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

never bufferize_to_store weak input [PR] (#17751)

Browse Source 
github github
This commit is contained in:
chenyu
2026-08-26 13:25:52 -04:00
committed by GitHub
parent 4bdc865131
commit 0abcf09b74
1 changed files with 8 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/schedule/rangeify.py
+8 -13
View File
@@ -1,7 +1,7 @@
from dataclasses import dataclass, field, replace
from typing import cast
import itertools
from tinygrad.dtype import dtypes, AddrSpace, Invalid, strong_dtype
from tinygrad.dtype import dtypes, AddrSpace, Invalid
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group
from tinygrad.uop.symbolic import symbolic
@@ -217,7 +217,7 @@ pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary)
def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
size = prod(x.shape)
dtype = strong_dtype(x.dtype) # a BUFFER is never weak: store at the concrete dtype, the .cast(x.dtype) on the result keeps readers unchanged
if x.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {x.dtype}")
rngs = sorted(idx.ranges, key=lambda x: x.arg)
assert size > 0 and isinstance(size, int), f"no zero sized or symbolic sized buffers {size}"
@@ -238,15 +238,15 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
# NOTE: the local BUFFER needs to be disambiguated here
if x.arg.addrspace == AddrSpace.GLOBAL:
buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(next(ctx), dtype, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
do_store = buf.index(idx).store(x.src[0].cast(dtype)).end(*rngs)
return buf.after(do_store).cast(x.dtype)
buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(next(ctx), x.dtype, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
do_store = buf.index(idx).store(x.src[0]).end(*rngs)
return buf.after(do_store)
if allow_locals:
# handle locals
buf = UOp.placeholder((size,), dtype, next(ctx), AddrSpace.LOCAL)
do_store = buf.index(idx).store(x.src[0].cast(dtype)).end(*rngs)
return buf.after(do_store).cast(x.dtype)
buf = UOp.placeholder((size,), x.dtype, next(ctx), AddrSpace.LOCAL)
do_store = buf.index(idx).store(x.src[0]).end(*rngs)
return buf.after(do_store)
# collapse any BUFFERIZE to single input BUFFERIZE
def flatten_bufferize(x:UOp):
@@ -271,11 +271,6 @@ def remove_noop_afters(x:UOp) -> UOp|None:
pm_add_buffers = pm_mops+pm_flatten_bufferize+PatternMatcher([
(UPat(Ops.STAGE, src=(UPat(), UPat(name="idx")), name="x"), lambda ctx,x,idx: bufferize_to_store(ctx, x, idx, allow_locals=False)),
# INDEX of a buffer through the weak cast added above: index the buffer directly and cast the loaded value instead.
# this must run in the same rewrite that adds the cast, or the expander expands the whole casted buffer into one big VECTORIZE
(UPat(Ops.INDEX, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("buf"),)),), allow_any_len=True, name="u"),
lambda u,buf: u.replace(dtype=None, src=(buf,)+u.src[1:]).cast(u.dtype)),
# 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)),