mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 11:56:08 +00:00
bufferize, don't use contig tag
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user