nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

don't use Ops.UNIQUE in Buffer [PR] (#16754)

Browse Source 
build new_buffer with ParamArg
This commit is contained in:
chenyu
2026-06-25 16:00:59 -04:00
committed by GitHub
parent 7f4ceeb2b9
commit d953eb95d5
6 changed files with 27 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_tensor_uop_mixin.py
+4 -2
View File
@@ -1,9 +1,11 @@
import math, unittest
from dataclasses import replace
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
from tinygrad.uop.ops import ParamArg, UOp, UPat, Ops, PatternMatcher, graph_rewrite
_strip_unique_pm = PatternMatcher([
(UPat((Ops.UNIQUE, Ops.LUNIQUE), name="u"), lambda u: u.replace(arg=0) if u.arg != 0 else None),
(UPat(Ops.LUNIQUE, name="u"), lambda u: u.replace(arg=0) if u.arg != 0 else None),
(UPat(Ops.BUFFER, name="b"), lambda b: b.replace(arg=replace(b.arg, slot=0)) if isinstance(b.arg, ParamArg) and b.arg.slot != 0 else None),
])
def _strip_unique(u: UOp) -> UOp: return graph_rewrite(u, _strip_unique_pm)
test/null/test_uops.py
+1 -1
View File
@@ -5,7 +5,7 @@ from tinygrad.tensor import Tensor
from tinygrad.helpers import Timing, Context, cdiv
from tinygrad.dtype import dtypes, ConstFloat # noqa: F401
from tinygrad.device import Device
from tinygrad.uop.ops import Ops, UOp, UPat, exec_alu
from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, exec_alu # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests
from tinygrad.uop.spec import spec_shared
from tinygrad.uop.symbolic import sym
from test.helpers import eval_uop, to_uops_list
tinygrad/callify.py
+3 -2
View File
@@ -1,6 +1,6 @@
from dataclasses import dataclass, field
from tinygrad.dtype import dtypes, AddrSpace, PtrDType, ImageDType
from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, track_rewrites
from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, ParamArg, graph_rewrite, track_rewrites
from tinygrad.helpers import VIZ, pluralize, all_int
@dataclass
@@ -193,7 +193,8 @@ pm_finalize_call = PatternMatcher([
pm_replace_buf = PatternMatcher([
# replace BUFFER with PARAM for cache key normalization
(UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
(UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
# replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
(UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer),
# strip value from BIND for cache key normalization, so different values hit same cache
tinygrad/uop/ops.py
+6 -11
View File
@@ -724,18 +724,14 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
if len(usrcs) == 0: return UOp(op, self.dtype, (self,), arg)
return UOp(op, self.dtype, (self,)+UOp.sink(*usrcs).simplify().src)
# *** uop UNIQUE ***
# TODO: use this in Buffer
unique_num = itertools.count(0)
@staticmethod
def unique(arg:int|None=None): return UOp(Ops.UNIQUE, arg=next(UOp.unique_num) if arg is None else arg)
# *** uop Buffer stuff ***
unique_num = itertools.count(0)
@staticmethod
def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
return UOp(Ops.BUFFER, dtype, (UOp.unique(num), UOp(Ops.DEVICE, arg=device)), size)
slot = next(UOp.unique_num) if num is None else num
return UOp(Ops.BUFFER, dtype, (shape_to_shape_arg((size,)),), ParamArg(slot, device=device))
@staticmethod
def from_buffer(opaque:Buffer, device:str|tuple[str, ...]|None=None):
if (uop:=UOp.new_buffer(device or opaque.device, opaque.size, opaque.dtype, num=-id(opaque))) not in buffers: buffers[uop] = opaque.ref(1)
@@ -873,7 +869,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
assert all_same([(x.size, x.dtype) for x in ret.bufs]), "multibuffers mismatch buffers"
return ret
assert self.op is Ops.BUFFER, f"must be BUFFER {self.op}"
assert self.src[0].op is Ops.UNIQUE, f"buffer src[0] must be UNIQUE, not {self.src[0].op}"
if (cret:=buffers.get(self)) is not None: return cret
rdtype = self.dtype if isinstance(self.dtype, ImageDType) else self.dtype.base
if isinstance(self.device, tuple): ret = MultiBuffer(self.device, self.max_numel(), rdtype).ref(1)
@@ -884,8 +879,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
def realized(self) -> Buffer|MultiBuffer|None:
# only these can be realized
if self.op not in (Ops.BUFFER, Ops.MSTACK): return None
# ParamArg is LOCAL/REG and never realized
if self.op is Ops.BUFFER and isinstance(self.arg, ParamArg): return None
# LOCAL/REG scratch buffers are never realized
if self.op is Ops.BUFFER and isinstance(self.arg, ParamArg) and self.addrspace in (AddrSpace.LOCAL, AddrSpace.REG): return None
# LUNIQUEs are never realized
if self.op_in_backward_slice_with_self(Ops.LUNIQUE): return None
# NOTE: this is used by the JIT to determine which inputs we capture
tinygrad/uop/render.py
+6 -5
View File
@@ -1,6 +1,6 @@
from tinygrad.dtype import dtypes
from tinygrad.dtype import AddrSpace, dtypes
from tinygrad.uop import Ops, GroupOp
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort
from tinygrad.uop.ops import ParamArg, UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort
from tinygrad.helpers import strip_parens
def pretty_print(x:UOp, cache=None, d=0)->str:
@@ -73,14 +73,15 @@ def render_marg(ctx,x:UOp):
pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg]
return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.UNIQUE, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
pm_pyrender_extra = PatternMatcher([
(UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
(UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
(UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
(UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
f"UOp.new_buffer({repr(d.arg)}, {x.arg}, {x.dtype}, {u.arg})"),
(UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x:
f"UOp.new_buffer({repr(x.arg.device)}, {x.max_numel()}, {x.dtype}, {x.arg.slot})"
if isinstance(x.arg, ParamArg) and x.addrspace is AddrSpace.GLOBAL else None),
(UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
(UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda ctx,x: f"UOp(Ops.CUSTOM_FUNCTION, {x.dtype}, src={srcs(ctx, x.src)}, arg={x.arg!r})"),
(UPat(Ops.REDUCE, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {r.arg[1]})" if len(r.arg[1]) else None),
tinygrad/uop/spec.py
+7 -5
View File
@@ -116,18 +116,20 @@ spec_shared = PatternMatcher([
(UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 8),
])
def is_device(d): return isinstance(d, str) or (isinstance(d, tuple) and all(isinstance(s, str) for s in d))
# these ops can exist in tensor but not programs. example: movement
spec_tensor = PatternMatcher([
# DEVICE
(UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d:
isinstance(d.arg, str) or (isinstance(d.arg, tuple) and all(isinstance(s, str) for s in d.arg))),
(UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d: is_device(d.arg)),
# UNIQUE
(UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True),
(UPat(Ops.LUNIQUE, dtypes.void, ()), lambda: True),
# BUFFER
(UPat(Ops.BUFFER, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE)), name="buf"),
(UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf:
(isinstance(buf.dtype, DType) and buf.src[0].dtype.scalar() == dtypes.weakint and is_device(buf.arg.device))
if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None),
(UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="buf"),
lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, DType)),
# Tensor variable bindings