bufferize, don't use contig tag

This commit is contained in:
2025-08-15 09:00:18 -07:00
parent 0575c95389
commit 551b34bb0f
4 changed files with 5 additions and 3 deletions
+1 -1
View File
@@ -119,7 +119,7 @@ def map_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp|None=None):
ctx.idx += 1
else:
ranges.append(UOp.const(dtypes.int, 0))
ret = x.src[0].index(*ranges).contiguous(*new_ranges, arg=x.arg)
ret = x.src[0].index(*ranges).bufferize(*new_ranges)
# if there's no open ranges, set arg to None so this uses a DEFINE_GLOBAL
if len(ret.ranges) == 0: ret = ret.replace(arg=None)
ret = ret.index(*passthrough_idx) if len(passthrough_idx) else ret
+1
View File
@@ -19,6 +19,7 @@ class Ops(FastEnum):
# ops that adjust the behavior of the scheduler
CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702
BUFFERIZE = auto()
# blocks in linearizer (only used there)
BLOCK = auto(); BLOCKSTART = auto(); BLOCKEND = auto(); BLOCKFINAL = auto() # noqa: E702
+2 -1
View File
@@ -188,7 +188,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
@functools.cached_property
def ranges(self) -> dict[UOp, None]:
if self.op is Ops.RANGE: return {self:None}
if self.op in {Ops.CONTIGUOUS, Ops.REDUCE, Ops.STORE}:
if self.op in {Ops.BUFFERIZE, Ops.REDUCE, Ops.STORE}:
ret = self.src[0].ranges.copy()
for s in self.src[1:]:
if s in ret: del ret[s]
@@ -295,6 +295,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
def contiguous(self, *args, **kwargs): return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
def bufferize(self, *args, **kwargs): return UOp(Ops.BUFFERIZE, dtype=self.dtype, src=(self,)+args, **kwargs)
def fuse(self): return self.alu(Ops.FUSE)
def allreduce(self, op, device:str|tuple[str, ...]|UOp):
assert isinstance(self.device, tuple), f"allreduce must be on tuple {self.device} isn't"
+1 -1
View File
@@ -19,7 +19,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
**{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF",
Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500",
Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
Ops.CHILD: "#80fff0"}
Ops.CHILD: "#80fff0", Ops.BUFFERIZE: "#FF991C"}
# VIZ API