forked from tinygrad/tinygrad
dtype_from_uop is never None (#17819)
This commit is contained in:
chenyu
@@ -126,7 +126,7 @@ def do_devectorize(b:UOp):
|
||||
if not all(x.shape == b.shape or x.base.is_invalid for x in b.src): return None
|
||||
src = []
|
||||
for idx_c in itertools.product(*[[UOp.const(i) for i in range(x)] for x in b.shape]):
|
||||
src.append(b.replace(dtype=None, src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src)))
|
||||
src.append(b.replace(src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src)))
|
||||
return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
|
||||
|
||||
def do_stack_wmma(u:UOp):
|
||||
@@ -418,7 +418,7 @@ def line_rewrite(lst:list[UOp], pm:PatternMatcher, ctx=None) -> list[UOp]:
|
||||
newlst = []
|
||||
replaced: dict[UOp, UOp] = {}
|
||||
for u in lst:
|
||||
nu = u.replace(dtype=None, src=tuple([replaced.get(x, x) for x in u.src]))
|
||||
nu = u.replace(src=tuple([replaced.get(x, x) for x in u.src]))
|
||||
ret: tuple[UOp, list[UOp]] = pm.rewrite(nu, ctx) or (nu, [nu])
|
||||
replaced[u] = ret[0]
|
||||
newlst.extend(ret[1])
|
||||
|
||||
@@ -9,8 +9,8 @@ from tinygrad.device import Compiler
|
||||
|
||||
# an access takes its dtype from the buffer it indexes, so accessing at another dtype restates the storage on the buffer that owns it
|
||||
def with_storage(x:UOp, dt:DType) -> UOp:
|
||||
if x.op in {Ops.PARAM, Ops.BUFFER}: return x.replace(dtype=None, arg=replace(x.arg, dtype=dt))
|
||||
return x.replace(dtype=None, src=(with_storage(x.src[0], dt),)+x.src[1:])
|
||||
if x.op in {Ops.PARAM, Ops.BUFFER}: return x.replace(arg=replace(x.arg, dtype=dt))
|
||||
return x.replace(src=(with_storage(x.src[0], dt),)+x.src[1:])
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Estimates:
|
||||
|
||||
@@ -298,7 +298,7 @@ def abi(ctx:IselContext, x:UOp) -> UOp|None:
|
||||
dt = dtypes.uint64 if x.op is Ops.PARAM and x.arg.addrspace is AddrSpace.GLOBAL else x.dtype
|
||||
arg = replace(x.arg, dtype=dt) if x.op is Ops.PARAM else x.arg
|
||||
# the shape srcs of a PARAM are not values, tag them so they aren't materialized into registers
|
||||
def _reg_arg(r:Register) -> tuple[UOp, ...]: return (x.replace(dtype=None, arg=arg, src=tuple(s.rtag() for s in x.src), tag=(r,)),)
|
||||
def _reg_arg(r:Register) -> tuple[UOp, ...]: return (x.replace(arg=arg, src=tuple(s.rtag() for s in x.src), tag=(r,)),)
|
||||
def _stack_arg(disp:int):
|
||||
return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=(X86Ops.FRAME_INDEX, dtypes.int32), tag=disp), imm(dtypes.uint8, 8))
|
||||
if sys.platform == "win32": src = _reg_arg((RCX, RDX, GPR[8], GPR[9])[i]) if i < 4 else _stack_arg((i-3)*8+32)
|
||||
@@ -809,13 +809,13 @@ class X86Renderer(ISARenderer):
|
||||
def stack_pointer(self) -> UOp: return def_reg(dtypes.uint64, RSP)
|
||||
# the value of a BUFFER is its address, it moves through registers and the stack as a 64bit int
|
||||
def copy(self, x:UOp, reg:Register):
|
||||
if x.op is Ops.BUFFER: x = x.replace(dtype=None, arg=replace(x.arg, dtype=dtypes.uint64))
|
||||
if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
|
||||
ret = isel_matcher.rewrite(UOp(Ops.COPY, src=(x,), tag=reg))
|
||||
assert ret is not None, f"failed to copy {x}"
|
||||
return ret
|
||||
|
||||
def spill(self, disp:UOp, x:UOp) -> UOp:
|
||||
if x.op is Ops.BUFFER: x = x.replace(dtype=None, arg=replace(x.arg, dtype=dtypes.uint64))
|
||||
if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
|
||||
is_xmm = isinstance(x.tag, tuple) and x.tag[0].cons[0].size == 16
|
||||
op = X86Ops.VMOVUPSm if is_xmm else X86Ops.MOVm
|
||||
return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
|
||||
|
||||
@@ -125,7 +125,7 @@ class NIRRenderer(Renderer):
|
||||
(UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
|
||||
# a bool is one bit in NIR but a byte in memory, so every access to a bool buffer goes through a uint8 view of it
|
||||
(UPat(Ops.LOAD, dtypes.bool, name="x"),
|
||||
lambda x: x.replace(dtype=None, src=(with_storage(x.src[0], dtypes.uint8),)+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())
|
||||
lambda x: x.replace(src=(with_storage(x.src[0], dtypes.uint8),)+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())
|
||||
+x.src[2:]).cast(dtypes.bool)),
|
||||
(UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
|
||||
lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), x.src[1].cast(dtypes.uint8))+x.src[2:])),
|
||||
|
||||
@@ -47,7 +47,7 @@ ptx_matcher = PatternMatcher([
|
||||
lambda x: (UOp(x.op, src=tuple(vv.cast(dtypes.float32) for vv in x.src), arg=x.arg).cast(dtypes.half))),
|
||||
# a bool is a predicate register in PTX but a byte in memory, so a bool buffer is accessed through a uint8 view of it
|
||||
(UPat(Ops.LOAD, dtypes.bool, src=(UPat(name="idx"),), name="x", allow_any_len=True),
|
||||
lambda x,idx: x.replace(dtype=None, src=(with_storage(idx, dtypes.uint8),) + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ())
|
||||
lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8),) + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ())
|
||||
+ x.src[2:]).cast(dtypes.bool) if idx.addrspace != AddrSpace.REG else None),
|
||||
(UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
|
||||
lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), x.src[1].cast(dtypes.uint8))+x.src[2:]) if idx.addrspace != AddrSpace.REG else None),
|
||||
|
||||
+12
-21
@@ -115,11 +115,11 @@ def promo_dtype(src:tuple[UOp,...]) -> DType:
|
||||
dts = [x.dtype for x in src]
|
||||
return dts[0] if all_same(dts) else least_upper_dtype(*dts)
|
||||
|
||||
def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
|
||||
# here are the dtype production rules, eventually this will go in UOp as a recursive property
|
||||
def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType:
|
||||
# here are the dtype production rules, total over all Ops
|
||||
match op:
|
||||
case Ops.STORE | Ops.LINEAR | Ops.SINK | Ops.PROGRAM | Ops.SOURCE | \
|
||||
Ops.END | Ops.BARRIER | Ops.GROUP | Ops.IF | Ops.ENDIF | \
|
||||
Ops.END | Ops.BARRIER | Ops.GROUP | Ops.IF | Ops.ENDIF | Ops.NOOP | \
|
||||
Ops.TUPLE | Ops.FUNCTION | Ops.CUSTOM_FUNCTION | Ops.REWRITE_ERROR | Ops.PYLITERAL:
|
||||
# always void
|
||||
return dtypes.void
|
||||
@@ -131,10 +131,8 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
|
||||
return arg[1]
|
||||
case Ops.INS:
|
||||
# arg is (instruction, dtype), a queue command or an asm line is void
|
||||
assert isinstance(arg, tuple) and len(arg) == 2 and isinstance(arg[1], DType), f"INS arg must be (instruction, DType), got {arg}"
|
||||
return arg[1]
|
||||
case Ops.NOOP:
|
||||
# NOOP can be void or carry any dtype (e.g. x.f(Ops.NOOP) or substitute base with NOOP)
|
||||
return None
|
||||
case Ops.INDEX:
|
||||
# an image access is always float, no matter the storage dtype
|
||||
# TODO: should there be a CAST so src[0].dtype just work?
|
||||
@@ -197,9 +195,9 @@ class UOpMetaClass(type):
|
||||
ucache:dict[tuple, weakref.ReferenceType[UOp]] = {}
|
||||
def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
|
||||
metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None):
|
||||
if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void
|
||||
if dtype is None: dtype = dtype_from_uop(op, src, arg)
|
||||
# TODO: delete this once the dtype field is removed, for now it just re-implements spec.py
|
||||
if SPEC == 2 and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
|
||||
elif SPEC == 2 and (expected_dtype:=dtype_from_uop(op, src, arg)) != dtype:
|
||||
raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}")
|
||||
if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret
|
||||
UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key))
|
||||
@@ -257,11 +255,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
if self.op is Ops.BUFFER and self.realized is not None: args.append(self.realized)
|
||||
return UOp, tuple(args)
|
||||
def replace(self, **kwargs) -> UOp:
|
||||
new_args = (kwargs.pop("op", self.op), kwargs.pop("dtype", self.dtype), kwargs.pop("src", self.src),
|
||||
kwargs.pop("arg", self.arg), kwargs.pop("tag", self.tag))
|
||||
new_args = (kwargs.pop("op", self.op), kwargs.pop("src", self.src), kwargs.pop("arg", self.arg), kwargs.pop("tag", self.tag))
|
||||
assert len(kwargs) == 0, f"unused kwargs in replace {list(kwargs)}"
|
||||
if (self.op, self.dtype, self.src, self.arg, self.tag) == new_args: return self
|
||||
return UOp(*new_args)
|
||||
if (self.op, self.src, self.arg, self.tag) == new_args: return self
|
||||
return UOp(new_args[0], src=new_args[1], arg=new_args[2], tag=new_args[3])
|
||||
def rtag(self, tag=True): return self.replace(tag=tag)
|
||||
@property
|
||||
def val(self):
|
||||
@@ -282,8 +279,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
return repr(self.arg)
|
||||
def tagstr(self): return f", tag={self.tag}" if self.tag is not None else ""
|
||||
|
||||
def f(self, op, **kwargs): return UOp(op, dtype=kwargs.pop("dtype", self.dtype), src=(self,), **kwargs)
|
||||
|
||||
@functools.cached_property
|
||||
def backward_slice(self:UOp) -> dict[UOp, None]:
|
||||
res: dict[UOp, None] = self.toposort()
|
||||
@@ -807,7 +802,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
case Ops.PERMUTE | Ops.FLIP: src_args = []
|
||||
case Ops.STACK:
|
||||
srcs = (self,)+tuple(arg)
|
||||
dtype = cast(DType, dtype_from_uop(Ops.STACK, srcs, None))
|
||||
dtype = dtype_from_uop(Ops.STACK, srcs, None)
|
||||
return UOp(Ops.STACK, src=tuple(u if u.base.is_invalid else UOp.const(u.val, dtype) if u.op is Ops.CONST else u.cast(dtype) for u in srcs))
|
||||
case _: raise RuntimeError(f"{op} is not a MovementOp")
|
||||
usrcs = [shape_to_shape_arg(arg) for arg in src_args]
|
||||
@@ -1694,7 +1689,7 @@ class RewriteContext:
|
||||
else:
|
||||
# rebuild node with rewritten srcs
|
||||
new_src = tuple(self.replace.get(x, x) for x in n.src)
|
||||
new_n = UOp(n.op, _rebuild_dtype(n, new_src), new_src, n.arg, n.tag) if new_src != n.src else n
|
||||
new_n = UOp(n.op, src=new_src, arg=n.arg, tag=n.tag) if new_src != n.src else n
|
||||
# top-down: try pm on rebuilt node, use result as-is (no re-traversal)
|
||||
if self.pm is not None and (rewritten:=self.pm_rewrite(new_n)) is not None: new_n = rewritten
|
||||
self.replace[n] = new_n
|
||||
@@ -1753,7 +1748,7 @@ class RewriteContext:
|
||||
continue
|
||||
else:
|
||||
# if srcs changed from rewrites, construct a new UOp with the new srcs
|
||||
new_src_n = UOp(new_n.op, _rebuild_dtype(new_n, new_src), new_src, new_n.arg, new_n.tag)
|
||||
new_src_n = UOp(new_n.op, src=new_src, arg=new_n.arg, tag=new_n.tag)
|
||||
# trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n
|
||||
stack.append((n, 2, new_src_n))
|
||||
stack.append((new_src_n, 0, new_src_n))
|
||||
@@ -1773,10 +1768,6 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N
|
||||
rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls)
|
||||
return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)
|
||||
|
||||
def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType:
|
||||
# TODO: delete this once the dtype field is removed, every rebuild will re-derive
|
||||
if all(a.dtype is b.dtype for a,b in zip(n.src, new_src)): return n.dtype
|
||||
return dtype_from_uop(n.op, new_src, n.arg) or n.dtype
|
||||
|
||||
def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(x, dtype)
|
||||
def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape)
|
||||
|
||||
@@ -150,7 +150,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
|
||||
(UPat(GroupOp.ALU-{Ops.THREEFRY}, src=bare_const, name="a"), fold_const_alu),
|
||||
(UPat(GroupOp.ALU-{Ops.THREEFRY}, src=casted_const, name="a"), fold_const_alu),
|
||||
(UPat(GroupOp.Binary-{Ops.THREEFRY}, src=[casted_const, bare_const], name="a"), lambda a:
|
||||
a.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in a.src))
|
||||
a.replace(src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in a.src))
|
||||
if (dt:=promo_dtype(a.src)) not in dtypes.weaks else None),
|
||||
# bool MUL is AND, ADD/MAX is OR. prevents other rules to rewrite bool ADD/MUL incorrectly
|
||||
(UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), lambda x,y: x&y),
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from dataclasses import replace
|
||||
from tinygrad.dtype import dtypes, DType, AddrSpace, Invalid, least_upper_dtype, strong_dtype, weak_dtype
|
||||
from tinygrad.helpers import unwrap
|
||||
|
||||
from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, GroupOp, dtype_from_uop, promo_dtype
|
||||
|
||||
def default_dtype(u:UOp):
|
||||
@@ -18,14 +18,14 @@ def commit_weak_consts(u:UOp, dt:DType|None) -> UOp|None:
|
||||
# the concrete dtypes u commits its srcs at: the operands' meet and u's own derived dtype, None if either is weak
|
||||
def derived_dtypes(u:UOp, src:tuple[UOp, ...]) -> tuple[DType, DType]|None:
|
||||
if u.op not in GroupOp.Broadcastable or (meet:=promo_dtype(src)) in dtypes.weaks \
|
||||
or (result:=unwrap(dtype_from_uop(u.op, src, u.arg))) in dtypes.weaks: return None
|
||||
or (result:=dtype_from_uop(u.op, src, u.arg)) in dtypes.weaks: return None
|
||||
return meet, result
|
||||
|
||||
def commit_srcs_at(u:UOp, dt:DType) -> UOp|None:
|
||||
# the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too
|
||||
dts = derived_dtypes(u, u.src)
|
||||
ret = u.replace(dtype=None, src=tuple(UOp.const(dt.const(s.val)) if s.op is Ops.CONST and s.dtype in dtypes.weaks and dts is not None else
|
||||
commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
|
||||
ret = u.replace(src=tuple(UOp.const(dt.const(s.val)) if s.op is Ops.CONST and s.dtype in dtypes.weaks and dts is not None else
|
||||
commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
|
||||
return None if ret is u else ret
|
||||
|
||||
def commit_weak_srcs(u:UOp) -> UOp|None:
|
||||
@@ -60,10 +60,10 @@ def lower_weak_node(u:UOp) -> UOp|None:
|
||||
# resolve whole once every weak expression lowered: a Binary widens from its own bounds too, derivable consts wait
|
||||
if u.op in _lower_weak_ops and src != u.src and not any(s.dtype in dtypes.weaks and s.op is not Ops.CONST for s in src[start:]):
|
||||
dt = strong_dtype(least_upper_dtype(default_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
|
||||
else unwrap(dtype_from_uop(u.op, src, u.arg)))
|
||||
return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid or s.dtype in dtypes.weaks else commit_weak(s, dt)
|
||||
else dtype_from_uop(u.op, src, u.arg))
|
||||
return u.replace(src=src[:start]+tuple(s if s.base.is_invalid or s.dtype in dtypes.weaks else commit_weak(s, dt)
|
||||
for s in src[start:])).cast(u.dtype)
|
||||
return None if src == u.src else u.replace(dtype=None, src=src)
|
||||
return None if src == u.src else u.replace(src=src)
|
||||
|
||||
pm_lower_weak = PatternMatcher([
|
||||
# a gated long index into a small buffer narrows; its out-of-gate value is discarded
|
||||
@@ -74,7 +74,7 @@ pm_lower_weak = PatternMatcher([
|
||||
(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),
|
||||
lambda u,x: x.cast(default_dtype(u.src[0])).cast(default_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
|
||||
(UPat((Ops.PARAM, Ops.BUFFER), dtype=dtypes.weakint, name="u"),
|
||||
lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=default_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
|
||||
lambda u: u.replace(arg=replace(u.arg, dtype=default_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
|
||||
(UPat(GroupOp.All, name="u"), lower_weak_node),
|
||||
])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user