diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 1288657e47..242a6535bb 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -48,7 +48,7 @@ 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}: assigned_to = assigned_to.src[0].base + while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.MULTI}: assigned_to = assigned_to.src[0].base if assigned_to.op is not Ops.BUFFER: return src.contiguous(tag=u.tag) def _make_buffer_view(src:UOp) -> UOp|None: @@ -65,7 +65,7 @@ def _make_buffer_view(src:UOp) -> UOp|None: def contiguous_mops_to_view(c:UOp, src:UOp): """CONTIGUOUS(MOPS(BUFFER)) → CONTIGUOUS(SLICE) when movement ops collapse to a contiguous range.""" buf = src.base - if buf.op not in {Ops.BUFFER, Ops.SLICE}: return None + if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.MULTI}: return None if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE}: return None # no symbolic shape @@ -202,7 +202,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]: # uop list is a list in the original_sink graph and we can map to the tags later # here we build buffer map dont_realize = {Ops.CONST, Ops.BUFFER, Ops.BIND, Ops.AFTER} - ctx = AllocCtx(bases=set([x.multibase for x in big_sink.src if x.base.op not in dont_realize and x.base.addrspace is not AddrSpace.ALU])) + ctx = AllocCtx(bases=set([x.base for x in big_sink.src if x.base.op not in dont_realize and x.base.addrspace is not AddrSpace.ALU])) # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up # this is the only one where we have to be careful to not break the tensor graph diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 2c507c53cb..ba0784406c 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -668,13 +668,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass): @property def base(self) -> UOp: - if self.op in GroupOp.Movement: return self.src[0].base - if self.op is Ops.MULTI: return self.src[0].base # MULTI is really a VIEW - if self.op is Ops.DETACH: return self.src[0].base # DETACH can't change base - return self - - @property - def multibase(self) -> UOp: if self.op in GroupOp.Movement: return self.src[0].base if self.op is Ops.DETACH: return self.src[0].base # DETACH can't change base return self @@ -822,7 +815,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): @property def buffer(self) -> Buffer|MultiBuffer: - if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer + if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.MULTI, 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 @@ -864,6 +857,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return ret @property def realized(self) -> Buffer|MultiBuffer|None: + if self.op is Ops.MULTI: return self.src[0].realized # only these can be realized if self.op not in (Ops.BUFFER, Ops.MSTACK): return None # LOCAL/REG scratch buffers are never realized