geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

refactor buffer_view op structure [pr] (#8540)

Browse Source 
* refactor buffer_view op [pr]

* only empty now

* same st

* empty shape is fine
This commit is contained in:
qazal
2025-01-09 03:07:46 -05:00
committed by GitHub
parent 4c5c32ff5f
commit 7595352dfc
2 changed files with 13 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/engine/schedule.py
+2 -3
View File
@@ -56,10 +56,9 @@ tensor_uop_spec = PatternMatcher([
# ** TODO: these UOps need new specs, the current representation relies on hacks
# BUFFER and VIEW specify device and shape for EMPTY and BUFFER_VIEW
(UPat(Ops.VIEW, name="view", src=(UPat(Ops.BUFFER, name="buf"), UPat({Ops.EMPTY, Ops.BUFFER_VIEW}, name="uop"))),
# BUFFER and VIEW specify device and shape for EMPTY
(UPat(Ops.VIEW, name="view", src=(UPat(Ops.BUFFER, name="buf"), UPat(Ops.EMPTY, name="uop"))),
lambda view,buf,uop: view.dtype == buf.dtype == uop.dtype and view.size == buf.size),
# NOTE: EMPTY just ensures the source BUFFER is allocated before children run
# TODO: this should be EMPTY(VIEW(BUFFER))
(UPat(Ops.EMPTY, src=(), arg=None), lambda: True),
tinygrad/ops.py
+11 -10
View File
@@ -284,8 +284,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
if self.op in GroupOp.Buffer: return vsrc[0] if len(vsrc:=[x.st for x in self.src if x.op is Ops.VIEW]) != 0 else None
if not (src_sts := [x.st for x in self.src if x.st is not None]): return None
assert all_same([x.shape for x in src_sts]), f"UOp sources must have the same shape {self} {[x.shape for x in src_sts]}"
if self.op is Ops.BUFFER_VIEW:
shape = src_sts[0].shape
if self.dtype.itemsize != (input_sz:=self.src[0].dtype.itemsize): shape = shape[:-1]+((shape[-1]*input_sz) // self.dtype.itemsize,)
# only reduce ops are allowed to change shape, everything else derives shape from sources
shape = src_sts[0].reduce(self.axis_arg) if self.op in (Ops.REDUCE_AXIS, Ops.WMMA) else src_sts[0].shape
elif self.op in {Ops.REDUCE_AXIS, Ops.WMMA}: shape = src_sts[0].reduce(self.axis_arg)
else: shape = src_sts[0].shape
from tinygrad.shape.shapetracker import ShapeTracker
return ShapeTracker.from_shape(shape)
@@ -363,13 +367,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
return ret
return UOp(Ops.CAST, dtype, (self,))
def bitcast(self, dtype:DType):
if self.can_view() and self.device.startswith("DISK"):
if self.dtype.itemsize == dtype.itemsize: output_shape = self.shape
else:
# https://pytorch.org/docs/stable/generated/torch.Tensor.view.html
if (self.shape[-1]*self.dtype.itemsize) % dtype.itemsize != 0: raise RuntimeError("unsupported size in bitcast")
output_shape = self.shape[:-1]+((self.shape[-1]*self.dtype.itemsize) // dtype.itemsize,)
return UOp.metaop(Ops.BUFFER_VIEW, output_shape, dtype, self.device, None, (self,))
if self.st is not None and self.shape and ((self.shape[-1]*self.dtype.itemsize)%dtype.itemsize != 0):
raise RuntimeError(f"unsupported size in bitcast {dtype}")
# shape changing bitcast can use a subbuffer on DISK
# TODO: this should be moved to realize.py
if self.can_view() and self.device.startswith("DISK"): return UOp(Ops.BUFFER_VIEW, dtype, (self,))
return UOp(Ops.BITCAST, dtype, (self,))
def gep(self, i:Union[tuple[int, ...], int]):
if isinstance(i, int):
@@ -424,8 +426,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self,x))
def contiguous(self, allow_buffer_view=True):
if not unwrap(self.st).contiguous or self.size != self.base.size or self.base.op is Ops.CONST:
if allow_buffer_view and self.can_view(): return self.metaop(Ops.BUFFER_VIEW, self.shape, self.dtype, self.device, None, (self,))
return self.alu(Ops.CONTIGUOUS)
return self.alu(Ops.BUFFER_VIEW if allow_buffer_view and self.can_view() else Ops.CONTIGUOUS)
forced_realize.add(self.base)
return self