diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index bf4a5e6e49..f1772f9840 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -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())