nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

remove slice from rangeify [PR] (#16981)

Browse Source
This commit is contained in:
sirhcm
2026-07-11 02:01:38 -04:00
committed by GitHub
parent 43ad225d36
commit a7b74ee593
1 changed files with 4 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/schedule/rangeify.py
+4 -8
View File
@@ -282,7 +282,7 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
return src.substitute(replaced, extra_pm=pm_gate_substitute)
def remove_noop_bufferize(idx,b2):
if idx.src[1:] != b2.src[1:] or idx.src[0].op is Ops.SLICE: return None
if idx.src[1:] != b2.src[1:]: return None
return idx.src[0].shrink(tuple((0, s) for s in b2.shape)) if b2.shape else idx.src[0]
def after_all_invalid(after:UOp):
@@ -378,11 +378,7 @@ 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), x.dtype, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
if x.src[0].op is Ops.SLICE:
# no INDEX on SLICE, this could be cleaner
do_store = buf.store(x.src[0]).end(*rngs)
else:
do_store = buf.index(idx).store(x.src[0]).end(*rngs)
do_store = buf.index(idx).store(x.src[0]).end(*rngs)
return buf.after(do_store)
if allow_locals:
@@ -519,11 +515,11 @@ def split_store(x:UOp) -> UOp|None:
lctx = LocalAddBufferContext()
ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True)
# SINK requires all buffers on the same device, but COPY/SLICE are cross-device or special hardware ops
# SINK requires all buffers on the same device, but COPY is cross-device
if ret.op is Ops.STORE: stored = ret.src[1]
elif ret.op is Ops.END and ret.src[0].op is Ops.STORE: stored = ret.src[0].src[1]
else: raise RuntimeError(f"unknown kernel type {ret.op}")
if stored.op in {Ops.COPY, Ops.SLICE}: ret = stored.replace(src=stored.src + ret.ended_ranges)
if stored.op is Ops.COPY: ret = stored.replace(src=stored.src + ret.ended_ranges)
else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
kernel = ret.call(*lctx.map.values(), *lctx.vars.keys())