diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 5be1e1b2b1..4f01a8eca5 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -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]) diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 86cd044a06..40682e7ccf 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -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: diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index fa7c1ab1f9..3bef96b2fa 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -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) diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index bcf0fdf2e0..71c00de8a1 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -125,7 +125,7 @@ class NIRRenderer(Renderer): (UPat.var('x', dtype=dtypes.bool)=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:])), diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index dcf763a89d..0b895fbe27 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -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), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 07ffb8983c..ab4be971b8 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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) diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 3ed54afa79..bf7b40db88 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -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), diff --git a/tinygrad/uop/weak.py b/tinygrad/uop/weak.py index ccf51c644d..8f946f33f0 100644 --- a/tinygrad/uop/weak.py +++ b/tinygrad/uop/weak.py @@ -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), ])