From 0c3c765ce060154e08ea207d30f3bc24c0cfac32 Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Sun, 2 Aug 2026 21:33:48 -0400 Subject: [PATCH 1/4] render casted const [pr] --- test/mockgpu/amd/pcode.py | 46 ++++++++++++++++++++---------------- tinygrad/renderer/cstyle.py | 37 +++++++++++++++++++++++------ tinygrad/renderer/isa/x86.py | 11 ++++++--- tinygrad/renderer/llvmir.py | 3 +++ tinygrad/renderer/nir.py | 4 +++- tinygrad/renderer/ptx.py | 7 +++++- tinygrad/renderer/wgsl.py | 6 +++++ 7 files changed, 81 insertions(+), 33 deletions(-) diff --git a/test/mockgpu/amd/pcode.py b/test/mockgpu/amd/pcode.py index bb66ffe03f1f2..6a5ac6cdc2bad 100644 --- a/test/mockgpu/amd/pcode.py +++ b/test/mockgpu/amd/pcode.py @@ -7,6 +7,10 @@ # Type alias for vars dict: stores UOps and tuples for lambda definitions VarVal = UOp | tuple[str, list[str], str] +# a committed constant reaches the emulator in either honest form: the typed CONST or a CAST over one. +# the inner CONST is bare for a tinygrad pair and typed for the emulator's own IR, so neither dtype is required +def _cc(u): return u.op is Ops.CONST or (u.op is Ops.CAST and u.src[0].op is Ops.CONST) +def _cv(u): return u.src[0].val if u.op is Ops.CAST else u.val def _const(dt, v): return UOp.const(v, dt) def _u32(v): return _const(dtypes.uint32, v) def _u64(v): return _const(dtypes.uint64, v) @@ -55,8 +59,8 @@ def _expr_bits(v: UOp) -> int: if v.op in (Ops.AND, Ops.XOR): widths: list[int] = [] for src in v.src: - if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0: - widths.append(src.val.bit_length()) + if _cc(src) and isinstance(_cv(src), int) and _cv(src) > 0 and (_cv(src) & (_cv(src) + 1)) == 0: + widths.append(_cv(src).bit_length()) if widths: return max(widths) return v.dtype.bitsize @@ -144,9 +148,9 @@ def minmax(a: UOp, b: UOp) -> UOp: def _find_two_pi_mul(x): if x.op != Ops.MUL or len(x.src) != 2: return None for i, s in enumerate(x.src): - if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586) + if _cc(s) and abs(_cv(s) - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586) if s.op == Ops.MUL and len(s.src) == 2: - vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST] + vals = [_cv(ss) for ss in s.src if _cc(ss)] if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1]) return None @@ -163,7 +167,7 @@ def _trig_reduce(x, phase=0.0): def _signext(val: UOp) -> UOp: for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]: - if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8: + if (val.op == Ops.AND and len(val.src) == 2 and _cc(val.src[1]) and _cv(val.src[1]) == mask) or val.dtype.itemsize == bits // 8: v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val sb = (v32 >> _u32(bits - 1)) & _u32(1) return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int) @@ -497,7 +501,7 @@ def _apply_binop(self, left, right, op): if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32) return (left >> right) if op == '>>' else (left << right) case '+' | '-': - if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val) + if op == '-' and _cc(left) and _cc(right): return _const(left.dtype, _cv(left) - _cv(right)) return (left + right) if op == '+' else (left - right) case '*' | '/': # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow @@ -507,7 +511,7 @@ def _apply_binop(self, left, right, op): left, right = left.cast(pdt), right.cast(pdt) if op == '*': return left * right return (left // right) if dtypes.is_int(left.dtype) else (left / right) - case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left + case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if _cc(left) and _cv(left) == 2.0 else left _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)] @@ -529,8 +533,8 @@ def unary(self) -> UOp: return inner.eq(_const(inner.dtype, 0)) if self.try_eat_val('-', 'OP'): inner = self.unary() - if inner.op == Ops.CONST: - return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val) + if _cc(inner): + return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -_cv(inner)) return inner.neg() if self.try_eat_val('+', 'OP'): return self.unary() return self.postfix() @@ -669,15 +673,15 @@ def _handle_bracket_rest(self, first: UOp, base: UOp, var_name: str | None = Non self.eat('OP') width = self.parse() self.eat('RBRACKET') - if width.op == Ops.CONST: - w = int(width.val) + if _cc(width): + w = int(_cv(width)) return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1) return base if self.try_eat('COLON'): second = self.parse() self.eat('RBRACKET') - if first.op == Ops.CONST and second.op == Ops.CONST: - a, b = int(first.val), int(second.val) + if _cc(first) and _cc(second): + a, b = int(_cv(first)), int(_cv(second)) if a < b: return _bitreverse(base, b - a + 1) hi, lo = a, b if lo >= base.dtype.itemsize * 8: @@ -698,8 +702,8 @@ def _handle_bracket_rest(self, first: UOp, base: UOp, var_name: str | None = Non dt_suffix = DTYPES.get(self.eat('IDENT').val, dtypes.uint32) if var_name is None: var_name = self._find_var_name(base) - if first.op == Ops.CONST: - idx = int(first.val) + if _cc(first): + idx = int(_cv(first)) # Check for array element (var@idx) if var_name and f'{var_name}@{idx}' in self.vars: v = self.vars[f'{var_name}@{idx}'] @@ -872,7 +876,7 @@ def mindex(idx:UOp): return mem.index(idx.valid(active) if active is not None el def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]: if l.dtype != r.dtype: - if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int) + if r.dtype == dtypes.int and _cc(r) and _cv(r) < 0: l = l.cast(dtypes.int) else: r = r.cast(l.dtype) return l, r @@ -969,8 +973,8 @@ def parse_bound(): p.eat('QUOTE') if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl')) expr = p.parse().simplify() - assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}" - return int(expr.val) + assert _cc(expr), f"loop bound must be constant, got {expr}" + return int(_cv(expr)) start_val = parse_bound() p.eat('COLON') end_val = parse_bound() @@ -1258,7 +1262,7 @@ def parse_bound(): def parse_cond(s, kw): ll = s.lower() return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), env, funcs)) - def is_const(c, v): return c.op == Ops.CONST and c.val is v + def is_const(c, v): return _cc(c) and _cv(c) is v cond = parse_cond(line, 'if') conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else [] branch_assigns: list[tuple[UOp, list]] = [] # (cond, assigns_list) for side-effect merging @@ -1341,9 +1345,9 @@ def _cond_side_effect(cnd, dest, val): # Build combined condition: each branch fires when its cond is true AND no earlier cond was true remaining = UOp.const(True) for bc, bse in branch_assigns: - effective = remaining & bc if remaining.op != Ops.CONST else bc + effective = remaining & bc if not _cc(remaining) else bc for dest, val in bse: assigns.append(_cond_side_effect(effective, dest, val)) - remaining = remaining & bc.logical_not() if remaining.op != Ops.CONST else bc.logical_not() + remaining = remaining & bc.logical_not() if not _cc(remaining) else bc.logical_not() for dest, val in else_side_effects: assigns.append(_cond_side_effect(remaining, dest, val)) continue diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index e25ab0c605fc3..1e3872c819f5b 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -23,13 +23,22 @@ # casting (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" \ if x.max_numel() > 1 and x.addrspace is AddrSpace.REG else None), - (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx[x.src[0]])})"), - (UPat(Ops.BITCAST, name="x"), lambda ctx,x: ctx[x.src[0]] if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None), - (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"__builtin_bit_cast({ctx.render_type(x)}, ({ctx.render_type(x.src[0])})({ctx[x.src[0]]}))"), - - # GPU stuff - (UPat(Ops.BARRIER), lambda ctx: ctx.barrier), - (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0]](x.arg[-1])}; /* {(x.src[0]).render()} */"), + # the pair CAST(dt, CONST(v)) is the typed constant it commits to: width from the CAST, value from the bare CONST. + # it sits above the generic CAST rule so a committed constant renders as a literal, not as a cast of one + (UPat(Ops.CAST, dtypes.floats, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), + lambda ctx,x,c: f"({ctx.render_cast(x, ctx.infinity if c.val > 0 else f'-{ctx.infinity}')})" if math.isinf(c.val) else + f"({ctx.render_cast(x, ctx.nan)})" if math.isnan(c.val) else None), + (UPat(Ops.CAST, dtypes.float, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),)), lambda ctx,c: f"{c.val}f"), + (UPat(Ops.CAST, dtypes.int64, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),)), lambda ctx,c: f"{c.val}l"), + (UPat(Ops.CAST, (dtypes.uint64, dtypes.uint32), src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), + lambda ctx,x,c: f"{truncate[x.dtype](c.val)}{'ul' if x.dtype is dtypes.uint64 else 'u'}"), + (UPat(Ops.CAST, dtypes.bool, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),)), lambda ctx,c: "1" if c.val else "0"), + (UPat(Ops.CAST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), + lambda ctx,x,c: f"({ctx.render_cast(x, f'{c.val}f')})"), + (UPat(Ops.CAST, (dtypes.uint8, dtypes.uint16), src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), + lambda ctx,x,c: f"({ctx.render_cast(x, f'{c.val}u')})"), + (UPat(Ops.CAST, (dtypes.int8, dtypes.int16), src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), + lambda ctx,x,c: f"({ctx.render_cast(x, str(c.val))})"), # const (UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, ctx.infinity)})"), @@ -47,6 +56,14 @@ # default const render (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.val)), + (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx[x.src[0]])})"), + (UPat(Ops.BITCAST, name="x"), lambda ctx,x: ctx[x.src[0]] if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None), + (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"__builtin_bit_cast({ctx.render_type(x)}, ({ctx.render_type(x.src[0])})({ctx[x.src[0]]}))"), + + # GPU stuff + (UPat(Ops.BARRIER), lambda ctx: ctx.barrier), + (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0]](x.arg[-1])}; /* {(x.src[0]).render()} */"), + # SHRINK/INDEX (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat.cvar()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), @@ -321,6 +338,8 @@ class OpenCLRenderer(CStyleLanguage): # bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort (UPat(Ops.CONST, dtypes.bfloat16, name="x"), lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.val)))[0] >> 16)}u"), + (UPat(Ops.CAST, dtypes.bfloat16, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),)), + lambda ctx,c: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(c.val)))[0] >> 16)}u"), # load/store image (OpenCL) (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"), (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))), @@ -498,6 +517,10 @@ def __init__(self, target:Target, use_hipcc=False): # gfx942 => MI300, gfx1100 = (UPat(Ops.CONST, dtypes.fp8s, arg=math.inf, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.infinity}, {fp8_index(x.dtype)})"), (UPat(Ops.CONST, dtypes.fp8s, arg=-math.inf, name="x"), lambda ctx,x: f"f32_to_fp8(-{ctx.infinity}, {fp8_index(x.dtype)})"), (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.val}f, {fp8_index(x.dtype)})"), + (UPat(Ops.CAST, dtypes.fp8s, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), + lambda ctx,x,c: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(c.val) else + f"f32_to_fp8({'-' if c.val < 0 else ''}{ctx.infinity}, {fp8_index(x.dtype)})" if math.isinf(c.val) else + f"f32_to_fp8({c.val}f, {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",), lambda ctx,x: f"f32_to_fp8({ctx[x.src[0]]}, {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",), diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 1346a026785ab..3f70746c2ace2 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -225,7 +225,11 @@ def lane(x:UOp, i:int) -> int: return s.src[1].val if (s:=x.src[i]).op is Ops.IN def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt] def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, arg=X86Ops.DEFINE, tag=None if reg is None else (reg,)) def imm(dt:DType, v:int) -> UOp: return UOp.const(truncate[dt](v), dt).rtag() +# a committed constant reaches isel in either honest form; the pair states the same number at the same width +def is_pair(u:UOp) -> bool: return u.op is Ops.CAST and u.src[0].op is Ops.CONST and u.src[0].dtype in dtypes.weaks +def unpair(u:UOp) -> UOp: return UOp.const(u.src[0].val, u.dtype) if is_pair(u) else u def to_imm(c:UOp) -> UOp|None: + c = unpair(c) if c.op is not Ops.CONST: return None if c.dtype is dtypes.int64: return imm(dtypes.int32, c.val) if not c.overflows(dtypes.int32) else None if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.val) if not c.overflows(dtypes.uint32) else None @@ -647,6 +651,7 @@ def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|No # 0b10 -- signals memory access with 32bit displacement # 0b11 -- signals no memory access if disp_uop is not None: + disp_uop = unpair(disp_uop) assert disp_uop.op is Ops.CONST, "displacement must be a constant" assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int" # rbp/r13 always require a displacement @@ -667,7 +672,7 @@ def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|No inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.val) # IMM byte if imm_uop is not None: - if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.val) + if (c:=unpair(imm_uop)).op is Ops.CONST: inst += struct.pack(unwrap(c.dtype.fmt), c.val) elif isinstance(greg(imm_uop), Register): inst += bytes([(greg(imm_uop).index & 0b1111) << 4 | 0b0000]) return inst @@ -677,13 +682,13 @@ def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|No if x.arg in X86GroupOp.WriteMem: if len(x.src) > 4: address, rest = x.src[:4], x.src[4:] else: address, rest = (x, None, None, None), x.src - imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,) + imm_uop = rest[:1] if rest and (rest[0].op is Ops.CONST or is_pair(rest[0])) else (None,) return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *imm_uop)) if x.arg in X86GroupOp.Rm1st: if len(x.src) > 3: address, rest = x.src[:4], x.src[4:] else: address, rest = (x.src[0], None, None, None), x.src[1:] - imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,) + imm_uop = rest[:1] if rest and (rest[0].op is Ops.CONST or is_pair(rest[0])) else (None,) return _encode(x, *address, *(None, *imm_uop)) if reg is None else _encode(None, *address, *(x if sel else None, *imm_uop)) if x.arg in X86GroupOp.Rm2nd: diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 56b3b38027a7f..5f99606caff44 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -170,7 +170,10 @@ def _render_kernel(self, uops: list[UOp], prefix:list[str]|None=None) -> tuple[t kernel.append(f" {r[u]} = addrspacecast [{size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{size} x {ldt(u.dtype)}]*") else: kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype)}], align 16") + # a bare weak CONST has no width: it is only the value half of a pair, rendered by its CAST parent + elif u.op is Ops.CONST and u.dtype in dtypes.weaks: continue elif u.op is Ops.CONST: r[u] = lconst(u.val, u.dtype) + elif u.op is Ops.CAST and u.src[0].op is Ops.CONST and u.src[0].dtype in dtypes.weaks: r[u] = lconst(u.src[0].val, u.dtype) elif u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype): r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast else: diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index b0d3a38f2ce63..2a7d04a9c80c2 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -145,6 +145,8 @@ class NIRRenderer(Renderer): def_rewrite = PatternMatcher([ (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.val, x.dtype)), + # the pair reads as the typed constant it commits to + (UPat(Ops.CAST, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), lambda ctx,x,c: nimm(ctx.b, c.val, x.dtype)), (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, x.dtype.itemsize if x.addrspace is AddrSpace.ALU else 8)), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))), (UPat(Ops.STORE, src=(UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val"))), @@ -191,7 +193,7 @@ def render(self, uops:list[UOp]): ranges: list[mesa.nir_def|None] = [] for u in uops: - if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.STACK and len(u.src) == 0): pass + if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.STACK and len(u.src) == 0) or (u.op is Ops.CONST and u.dtype in dtypes.weaks): pass elif u.op in {Ops.INDEX, Ops.SHRINK}: # INDEX on a register value picks the element, memory INDEX is handled in the LOAD/STORE patterns if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].val) diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index 1b6b77859aebb..c1215aa97e600 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -80,6 +80,11 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i string_rewrite = PatternMatcher([ (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.val, x.dtype)}, 0;"), + # the pair reads as the typed constant it commits to + (UPat(Ops.CAST, dtypes.bool, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), + lambda ctx,x,c: f"setp.ne.s16 {ctx.r[x]}, {render_val(c.val, x.dtype)}, 0;"), + (UPat(Ops.CAST, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), + lambda ctx,x,c: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(c.val, x.dtype)};"), (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.val, x.dtype)};"), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"), (UPat(Ops.PARAM, name="x"), lambda ctx, x: @@ -186,7 +191,7 @@ def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str: name = "test" for u in uops: - if u.op in {Ops.NOOP, Ops.GROUP}: continue + if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.CONST and u.dtype in dtypes.weaks): continue if u.op is Ops.AFTER: self.r[u] = self.r[u.src[0]] continue diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index e9e2bfdcc3584..a95c690bd8015 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -72,6 +72,12 @@ class WGSLRenderer(CStyleLanguage): (UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"), lambda x: f"bitcast({x.val})" if x.val < 0 else f"{x.val&0xFFFFFFFF}u"), (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}"), + # the pair reads as the typed constant it commits to + (UPat(Ops.CAST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), src=(UPat(Ops.CONST, dtypes.weaks, name="c"),)), + lambda c: f"bitcast({c.val})" if c.val < 0 else f"{c.val&0xFFFFFFFF}u"), + (UPat(Ops.CAST, dtype=dtypes.int32, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), + lambda ctx,x,c: f"{truncate[x.dtype](c.val)}"), + (UPat(Ops.CAST, dtype=dtypes.bool, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),)), lambda c: "true" if c.val else "false"), (UPat(Ops.BUFFER, name="x"), lambda ctx,x: f"var{'' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"), (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)), From 5ed6761db346885d0968b4de851acf58e2144edd Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Sun, 2 Aug 2026 21:54:22 -0400 Subject: [PATCH 2/4] BOUNDARY --- tinygrad/codegen/__init__.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 3f37a4c1ddd50..063730b75c075 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -341,7 +341,10 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # extra symbolic before decomp. crashes without this? sink = graph_rewrite(sink, sym, name="extra symbolic") - # lower index dtype + # ***** THE BOUNDARY: weakness ends here ***** + # above this line a value may be widthless; below it every width is explicit and a rule that mints a const states its width. + # pm_commit_weak runs before lower_weak_srcs inside pm_lower_index_dtype and that order is load-bearing (the reverse costs 60 null tests). + # pm_lower_weak must never be merged into a matcher that also folds CAST-of-CONST: the marker it mints is folded straight back, a 2-cycle. # NOTE: we need indexing_simplify to remove the cast to long using the Invalid sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes") @@ -357,7 +360,10 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops) sink = graph_rewrite(sink, pm_decomp, name="early decompositions") - # late decomps + move gates from unrenderable INVALID where + # below the boundary weakness re-enters only from minting rounds, and each is closed by the commit round that follows it: + # "early decompositions" mints -> "decomp dtypes" commits (pm_commit_weak) + # "late decompositions" mints -> "final rewrite" commits (pm_commit_weak, which is why it heads that sum) + # a new minting round below here needs a commit round after it, or its weak nodes reach spec_program and are rejected sink = graph_rewrite(sink, pm_dtype_decomps+pm_commit_weak, ctx=(set(), ren), name="decomp dtypes") pm_decomp = pm_decomp+\ get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\ @@ -367,6 +373,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # final rules for the renderer (without sym) extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([]) + # pm_commit_weak heads this sum: renderer matchers in extra_matcher may mint weak, and this is the last round that can commit them pm_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite") From 631d7a4932285a6b9f90405e3c62a91fa7a01923 Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Mon, 3 Aug 2026 09:40:40 -0400 Subject: [PATCH 3/4] machine_const? --- test/backend/test_encodings.py | 63 +++++++++++++------ test/backend/test_isel.py | 19 ++++-- tinygrad/codegen/late/regalloc.py | 9 ++- tinygrad/renderer/isa/__init__.py | 4 ++ tinygrad/renderer/isa/x86.py | 100 ++++++++++++++++-------------- 5 files changed, 120 insertions(+), 75 deletions(-) diff --git a/test/backend/test_encodings.py b/test/backend/test_encodings.py index 38e1373beb60c..cb62b27747eb8 100644 --- a/test/backend/test_encodings.py +++ b/test/backend/test_encodings.py @@ -2,7 +2,7 @@ from tinygrad import Device from tinygrad.uop.ops import UOp, Ops from tinygrad.dtype import dtypes -from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg +from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, machine_const, def_reg def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag) @@ -14,51 +14,51 @@ def encode(self, u:UOp): return Device[Device.DEFAULT].renderer.render([u]) # displacement of 0 isn't emitted def test_base_address(self): - load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RDI) + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), UOp(Ops.NOOP), machine_const(0), machine_const(4)), RDI) # mov edi, dword ptr [rdi] self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 3F")) # rsp/r12 require a sib byte when used as base memory address def test_rsp_base_address(self): - load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RSP) + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), machine_const(0), machine_const(4)), RSP) # mov esp, dword ptr [rsp] self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 24 24")) # rbp/r13 require a displacement when used as base memory address def test_rbp_base_address(self): - load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RBP), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RBP) + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RBP), UOp(Ops.NOOP), machine_const(0), machine_const(4)), RBP) # mov ebp, dword ptr [rbp + 0] self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 6D 00")) # test [base + index*scale] def test_base_index_address(self): - load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RDX), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RAX) + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RDX), machine_const(0), machine_const(4)), RAX) # mov eax, dword ptr [rax + rdx*4] self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 04 90")) # rsp as index means no index def test_rsp_index_address(self): - load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RSP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RAX) + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, RSP), machine_const(0), machine_const(4)), RAX) # mov eax, dword ptr [rax] self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 00")) # however r12 is a valid index def test_r12_index_address(self): load = ins(X86Ops.MOV, dtypes.int32, - (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, GPR[12]), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), RAX) + (def_reg(dtypes.uint64, RAX), def_reg(dtypes.int32, GPR[12]), machine_const(0), machine_const(4)), RAX) # mov eax, dword ptr [rax + r12*4] self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("42 8B 04 A0")) # test [base + index*scale + 8bit disp] def test_complex_address_8bit_disp(self): - load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4)), RDI) + load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10), machine_const(4)), RDI) # mov edi, dword ptr [rdi + rsi*4 + 0xa] self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 7C B7 0A")) # test [base + index*scale + 32bit disp] def test_complex_address_32bit_disp(self): load = ins(X86Ops.MOV, dtypes.int32, - (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10000), imm(dtypes.uint8, 4)), RDI) + (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10000), machine_const(4)), RDI) # mov edi, dword ptr [rdi + rsi*4 + 0x2710] self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B BC B7 10 27 00 00")) @@ -116,30 +116,57 @@ def test_reg_in_imm_field(self): # when writting to mem the uop takes the store form where dtype is void and there's no definition def test_write_mem(self): - address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4)) + address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10), machine_const(4)) xmm0 = def_reg(dtypes.float32, XMM[0]) - extr = ins(X86Ops.VPEXTRD, dtypes.void, address + (xmm0, imm(dtypes.uint8, 0))) + extr = ins(X86Ops.VPEXTRD, dtypes.void, address + (xmm0, machine_const(0))) # vpextrd dword ptr [rdi + rsi*4 + 0xa], xmm0, 0 self.assertEqual(bytes.fromhex(self.encode(extr)), bytes.fromhex("C4 E3 79 16 44 B7 0A 00")) # test two address instruction with fused load works def test_two_address_load(self): - address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4)) + address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10), machine_const(4)) cmove = ins(X86Ops.CMOVE, dtypes.int32, address, RAX) # cmove eax, dword ptr [rdi + rsi*4 + 0xa] self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 44 B7 0A")) # test instruction where displacement and imm have the same value def test_disp_imm_same_value(self): - address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int8, RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 1)) - mov = ins(X86Ops.MOVi, dtypes.void, address + (imm(dtypes.int8, 10),)) + address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int8, RSI), machine_const(10), machine_const(1)) + mov = ins(X86Ops.MOVi, dtypes.void, address + (machine_const(10),)) # mov byte ptr [rdi + rsi + 0xa], 0xa self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex("40 C6 44 37 0A 0A")) - address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10), imm(dtypes.uint8, 4)) - imul = ins(X86Ops.IMULi, dtypes.int32, address + (imm(dtypes.int32, 10),), RDI) + address = (def_reg(dtypes.uint64, RDI), def_reg(dtypes.int32, RSI), machine_const(10), machine_const(4)) + imul = ins(X86Ops.IMULi, dtypes.int32, address + (machine_const(10),), RDI) # imul edi, dword ptr [rdi + rsi*4 + 0xa], 0xa - self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 BC B7 0A 00 00 00 0A 00 00 00")) + self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 7C B7 0A 0A 00 00 00")) + + # a displacement takes the smallest width that holds it wherever it was minted, so a stack argument offset is a disp8 like any other + def test_disp_width(self): + def load(d): return ins(X86Ops.MOV, dtypes.uint64, (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), machine_const(d), machine_const(8)), RAX) + # mov rax, qword ptr [rsp + 8] / [rsp + 127] + self.assertEqual(bytes.fromhex(self.encode(load(8))), bytes.fromhex("48 8B 44 24 08")) + self.assertEqual(bytes.fromhex(self.encode(load(127))), bytes.fromhex("48 8B 44 24 7F")) + # the rule is on the magnitude, so 128 and -128 both take 4 bytes + self.assertEqual(bytes.fromhex(self.encode(load(128))), bytes.fromhex("48 8B 84 24 80 00 00 00")) + self.assertEqual(bytes.fromhex(self.encode(load(-128))), bytes.fromhex("48 8B 84 24 80 FF FF FF")) + + # a machine const is a number with no width, the field it lands in writes its low bytes + def test_imm_low_bytes(self): + # movabs rax, 0xfffffffffffffffe, the value is wider than the register + movabs = ins(X86Ops.MOVABS, dtypes.uint64, (machine_const(2**65-2),), RAX) + self.assertEqual(bytes.fromhex(self.encode(movabs)), bytes.fromhex("48 B8 FE FF FF FF FF FF FF FF")) + for dt,v,b in ((dtypes.int8, 300, "2C"), (dtypes.int8, -200, "38"), (dtypes.uint8, 300, "2C")): + # mov al, + self.assertEqual(bytes.fromhex(self.encode(ins(X86Ops.MOVi, dt, (machine_const(v),), RAX))), bytes.fromhex("40 C6 C0 " + b)) + # mov ax, 0x1170 + self.assertEqual(bytes.fromhex(self.encode(ins(X86Ops.MOVi, dtypes.int16, (machine_const(70000),), RAX))), bytes.fromhex("66 40 C7 C0 70 11")) + + # a constant in an immediate slot takes the width of the slot, not of the value, in either honest form + def test_pair_imm_width(self): + # add rax, 5 with an imm32, the operand being 64 bit doesn't widen the immediate + for c in (machine_const(5), UOp.const(5).cast(dtypes.int64)): + self.assertEqual(bytes.fromhex(self.encode(ins(X86Ops.ADDi, dtypes.int64, (c,), RAX))), bytes.fromhex("48 81 C0 05 00 00 00")) # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding def test_cmove_ignore_cmp(self): @@ -148,4 +175,4 @@ def test_cmove_ignore_cmp(self): self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0")) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/test/backend/test_isel.py b/test/backend/test_isel.py index 6965a5db1711a..5326523dd0918 100644 --- a/test/backend/test_isel.py +++ b/test/backend/test_isel.py @@ -1,10 +1,9 @@ import unittest from typing import cast from tinygrad import Device -from tinygrad.uop import Ops from tinygrad.uop.ops import UOp, dtypes, graph_rewrite -from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops -from tinygrad.renderer.isa import IselContext +from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops, to_imm +from tinygrad.renderer.isa import IselContext, machine_const # INDEX on a register value with a constant index extracts a single element (the old GEP) def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar()) @@ -48,8 +47,18 @@ def test_complex_address(self): a = UOp.variable("a", 0, 0, dtypes.int32) load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load() n = self.isel_rewrite(load) - # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32 - self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4) + # displacement is the constant in "a" scaled to the buffer element size; encoding chooses its width + self.assertIs(n.src[2], machine_const(4)) + + # a constant reaches isel as a strong CONST or as a cast weak one, both select the same machine const + def test_to_imm_forms(self): + for dt in (dtypes.int8, dtypes.uint8, dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64): + for v in (0, 1, dt.min, dt.max): self.assertIs(to_imm(UOp.const(v, dt)), to_imm(UOp.const(v).cast(dt)), f"{dt} {v}") + + # a machine const states the number the hardware sees, and only fits an immediate if it fits 4 bytes + def test_to_imm_value(self): + self.assertIs(to_imm(UOp.const(-1, dtypes.uint32)), machine_const(2**32-1)) + for dt,v in ((dtypes.int64, 2**31), (dtypes.uint64, 2**32), (dtypes.weakint, 2**31)): self.assertIsNone(to_imm(UOp.const(v, dt))) if __name__ == "__main__": unittest.main() diff --git a/tinygrad/codegen/late/regalloc.py b/tinygrad/codegen/late/regalloc.py index 0d675e62ee019..c136abc3a781d 100644 --- a/tinygrad/codegen/late/regalloc.py +++ b/tinygrad/codegen/late/regalloc.py @@ -1,8 +1,7 @@ import itertools from tinygrad.helpers import dedup from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat -from tinygrad.renderer.isa import ISARenderer, Register, greg -from tinygrad.dtype import dtypes +from tinygrad.renderer.isa import ISARenderer, Register, greg, machine_const PSEUDO_OPS = {Ops.CONST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK} @@ -52,7 +51,7 @@ def fill(v:Register, i:int, cons:tuple[Register, ...]|None=None) -> Register: # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize) offset = self.stack_size + (sz - self.stack_size % sz) % sz - self.spills[v] = UOp.const(offset, dtypes.int32) + self.spills[v] = machine_const(offset) self.stack_size = offset + sz r = alloc(cons if cons is not None else v.cons, i) self.insert_before.setdefault(i, []).append((v, r)) @@ -84,7 +83,7 @@ def fill(v:Register, i:int, cons:tuple[Register, ...]|None=None) -> Register: # allocate stack array if u.op is Ops.BUFFER: - self.locals[u] = UOp.const(self.stack_size, dtypes.int32) + self.locals[u] = machine_const(self.stack_size) self.stack_size += u.max_numel() * u.dtype.itemsize # loop prologue, avoid loading inside the loop @@ -125,7 +124,7 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp): # alloc/dealloc stack if ctx.stack_size > 0: sp = ctx.ren.stack_pointer() - offset = UOp.const(ctx.stack_size, sp.dtype) + offset = machine_const(ctx.stack_size) if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))] diff --git a/tinygrad/renderer/isa/__init__.py b/tinygrad/renderer/isa/__init__.py index fb74f1569def1..fb97cc457c462 100644 --- a/tinygrad/renderer/isa/__init__.py +++ b/tinygrad/renderer/isa/__init__.py @@ -4,6 +4,10 @@ from tinygrad.renderer import Renderer from tinygrad.uop.ops import PatternMatcher, UOp, Ops, consumer_map_from_toposort +# a constant an instruction reads directly: a number with no width of its own, the field it lands in decides how many bytes are written. +# the tag marks it as already selected so isel doesn't materialize it into a register +def machine_const(v:int) -> UOp: return UOp.const(v).rtag() + @dataclass(frozen=True) class Register: name: str diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 3f70746c2ace2..6a9adbb018b92 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -5,8 +5,8 @@ from tinygrad.dtype import dtypes, DType, truncate, AddrSpace from tinygrad.uop import FastEnum, auto, Ops, GroupOp from tinygrad.uop.ops import UOp, UPat, PatternMatcher -from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg -from tinygrad.helpers import getenv, NUM_CPU_THREADS, unwrap, Target +from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg, machine_const +from tinygrad.helpers import getenv, NUM_CPU_THREADS, Target # ***** X86 Ops ***** @@ -224,23 +224,23 @@ def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX def lane(x:UOp, i:int) -> int: return s.src[1].val if (s:=x.src[i]).op is Ops.INDEX else 0 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt] def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, arg=X86Ops.DEFINE, tag=None if reg is None else (reg,)) -def imm(dt:DType, v:int) -> UOp: return UOp.const(truncate[dt](v), dt).rtag() # a committed constant reaches isel in either honest form; the pair states the same number at the same width def is_pair(u:UOp) -> bool: return u.op is Ops.CAST and u.src[0].op is Ops.CONST and u.src[0].dtype in dtypes.weaks def unpair(u:UOp) -> UOp: return UOp.const(u.src[0].val, u.dtype) if is_pair(u) else u +# the number a const denotes: CONST keeps the value it was built with, a machine const has to carry the one the hardware will see +def cval(c:UOp) -> int: return truncate[c.dtype](c.val) if c.dtype in dtypes.ints else c.val +# an immediate field is at most 4 bytes wide, a value that doesn't fit it has to go through a register. a machine const is already sized def to_imm(c:UOp) -> UOp|None: c = unpair(c) if c.op is not Ops.CONST: return None - if c.dtype is dtypes.int64: return imm(dtypes.int32, c.val) if not c.overflows(dtypes.int32) else None - if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.val) if not c.overflows(dtypes.uint32) else None - if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.val) - return None + if c.dtype in (dtypes.int64, dtypes.weakint) and c.overflows(dtypes.int32) or c.dtype is dtypes.uint64 and c.overflows(dtypes.uint32): return None + return machine_const(cval(c)) if c.dtype in dtypes.ints+(dtypes.bool, dtypes.weakint) else None def cmp(x:UOp) -> UOp: if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void) if x.src[0].dtype is dtypes.float64: return x.ins(X86Ops.VUCOMISD, dtype=dtypes.void) return x.ins(X86Ops.CMP, dtype=dtypes.void) if (i:=to_imm(x.src[1])) is None else x.ins(X86Ops.CMPi, dtype=dtypes.void, src=(x.src[0], i)) def vcmp(x:UOp) -> UOp: - v = imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op]) + v = machine_const({Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op]) if x.dtype.scalar() is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, src=x.src + (v,)) return x.ins(X86Ops.VCMPSD if x.max_numel() == 1 else X86Ops.VCMPPD, src=x.src + (v,)) @@ -250,22 +250,22 @@ def vcmp(x:UOp) -> UOp: def vinsertps(x:UOp) -> UOp: def _insert(ret:UOp, i:int) -> UOp: s, v = base(x, i), lane(x, i) - return x.ins(X86Ops.VINSERTPS, src=(ret, s, imm(dtypes.uint8, v << 6 | i << 4))) + return x.ins(X86Ops.VINSERTPS, src=(ret, s, machine_const(v << 6 | i << 4))) return functools.reduce(_insert, range(len(x.src)), def_reg(x.dtype)) # vpinsq xmm2, xmm0, rax, imm # inserts element in rax into any position in xmm0, result is written to xmm2 according to imm def vpins(x:UOp) -> UOp: op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.scalar().itemsize] - return functools.reduce(lambda ret,i: x.ins(op, src=(ret, x.src[i], imm(dtypes.uint8, i))), range(len(x.src)), def_reg(x.dtype)) + return functools.reduce(lambda ret,i: x.ins(op, src=(ret, x.src[i], machine_const(i))), range(len(x.src)), def_reg(x.dtype)) # we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg def idiv(ctx:IselContext, x:UOp) -> UOp: op = X86Ops.DIV if x.dtype in dtypes.uints else X86Ops.IDIV # for >8bit need to zero/sign extend rax to rdx if x.dtype in dtypes.int8s: ext = [] - elif x.dtype in dtypes.uints: ext = [x.ins(X86Ops.MOVi, src=(imm(min(dtypes.uint32, x.dtype), 0),), tag=(RDX,))] - else: ext = [x.ins(X86Ops.SARi, src=(x.src[0], imm(dtypes.uint8, x.dtype.itemsize * 8 - 1)), tag=(RDX,))] + elif x.dtype in dtypes.uints: ext = [x.ins(X86Ops.MOVi, src=(machine_const(0),), tag=(RDX,))] + else: ext = [x.ins(X86Ops.SARi, src=(x.src[0], machine_const(x.dtype.itemsize * 8 - 1)), tag=(RDX,))] # for 8bit need to zero/sign extend al to ah if x.dtype is dtypes.uint8: dividend = UOp(Ops.INS, arg=X86Ops.MOVZX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,)) elif x.dtype is dtypes.int8: dividend = UOp(Ops.INS, arg=X86Ops.MOVSX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,)) @@ -286,16 +286,15 @@ def shift(x:UOp, op:X86Ops) -> UOp: # a memory address operand is (base, index, displacement, size). size is the element size, it scales the index and is the memory operand width. # it is materialized as an immediate so the address stays correct if the base register is ever spilled and refilled def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]: - def _disp(v:int) -> UOp: return imm(dtypes.int32 if abs(v) > dtypes.int8.max else dtypes.int8, v) def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v - if x.op not in {Ops.INDEX, Ops.SHRINK}: return (x, UOp(Ops.NOOP), _disp(0), imm(dtypes.uint8, x.dtype.itemsize)) + if x.op not in {Ops.INDEX, Ops.SHRINK}: return (x, UOp(Ops.NOOP), machine_const(0), machine_const(x.dtype.itemsize)) base, idx = x.src[0], x.src[1] # buffers are indexed by element, everything else (the stack pointer) by byte scale = base.dtype.itemsize if base.op in {Ops.PARAM, Ops.BUFFER, Ops.AFTER} else 1 - sz = imm(dtypes.uint8, base.dtype.itemsize) - if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].val * scale), sz) - if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.val * scale), sz) - return (base, _cast(idx), _disp(0), sz) + sz = machine_const(base.dtype.itemsize) + if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), machine_const(idx.src[1].val * scale), sz) + if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), machine_const(idx.val * scale), sz) + return (base, _cast(idx), machine_const(0), sz) def abi(ctx:IselContext, x:UOp) -> UOp|None: if isinstance(x.tag, tuple): return None @@ -305,7 +304,7 @@ def abi(ctx:IselContext, x:UOp) -> UOp|None: # 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=dt, 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, dtype=dtypes.int32, tag=disp), imm(dtypes.uint8, 8)) + return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=X86Ops.FRAME_INDEX, dtype=dtypes.int32, tag=disp), machine_const(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) else: src = _reg_arg((RDI, RSI, RDX, RCX, GPR[8], GPR[9])[i]) if i < 6 else _stack_arg((i-5)*8) # this move "cleanses" the abi register constraint @@ -357,7 +356,7 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: # cast of void is a noop (UPat.var("y").cast(name="x"), lambda y,x: y if y.dtype == dtypes.void else None), # range is lowered to acc, cmp, jmp after regalloc - (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.val),) + x.src[1:])), + (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(machine_const(c.val),) + x.src[1:])), (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None), # really all a backedge END is is an IF with a tag referencing the RANGE start label (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"), @@ -371,8 +370,8 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: # function abi constraints (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi), # constants that can't be immediates, move them to registers - (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.val),)) if not x.tag else None), - (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.val),)) if not x.tag else None), + (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(machine_const(cval(x)),)) if not x.tag else None), + (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(machine_const(cval(x)),)) if not x.tag else None), (UPat.cvar("x", dtypes.floats), lambda x: UOp.const(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, x.val))[0], dt).bitcast(x.dtype) if not x.tag else None), # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists @@ -415,9 +414,9 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPS)), (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPD)), (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x: - x.ins(X86Ops.VROUNDSS, src=(y, y, imm(dtypes.uint8, 3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPS, src=(y, imm(dtypes.uint8, 3)))), + x.ins(X86Ops.VROUNDSS, src=(y, y, machine_const(3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPS, src=(y, machine_const(3)))), (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x: - x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPD, src=(y, imm(dtypes.uint8, 3)))), + x.ins(X86Ops.VROUNDSD, src=(y, y, machine_const(3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPD, src=(y, machine_const(3)))), # for float16 we route the srcs through gprs, this is suboptimal for values in xmms, in that case we want vpunpcklwd (UPat(Ops.STACK, dtypes.float16, name="x"), lambda x: vpins(x.replace(src=tuple(s.bitcast(dtypes.int16) for s in x.src)))), @@ -425,15 +424,15 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), vpins), # INDEX on a vector register value extracts a single element (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, machine_const(c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.int16s).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, machine_const(c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.int32s).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, machine_const(c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.int64s).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, machine_const(c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.floats).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.val * x.dtype.itemsize))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, machine_const(c.val * x.dtype.itemsize))) if _is_vec_xmm(y) else None), # packed bitwise ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.max_numel() > 1 else None), ((UPat() | UPat()).named("x"), lambda x: x.ins(X86Ops.VPOR) if x.max_numel() > 1 else None), @@ -457,9 +456,9 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: # scalar int binary ((UPat(dtype=dtypes.ints).alu(Ops.CDIV, UPat())).named("x"), idiv), # scalar int binary with immediate - (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.val)))), - (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))), - (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))), + (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, machine_const(c.val)))), + (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, machine_const(c.val)))), + (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, machine_const(c.val)))), (UPat.var("a", dtypes.ints) + UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None), (UPat.var("a", dtypes.ints) * UPat.cvar("c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None), (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None), @@ -492,7 +491,7 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: (UPat(dtype=dtypes.float64).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPD2DQ) if x.max_numel() > 1 else None), (UPat(dtype=dtypes.float32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PD) if x.max_numel() > 1 else None), (UPat(dtype=dtypes.float64).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPD2PS) if x.max_numel() > 1 else None), - (UPat(dtype=dtypes.float32).cast(dtypes.float16, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PH, src=x.src + (imm(dtypes.uint8, 4),))), + (UPat(dtype=dtypes.float32).cast(dtypes.float16, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PH, src=x.src + (machine_const(4),))), (UPat(dtype=dtypes.float16).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPH2PS)), (UPat(dtype=dtypes.float32).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSS2SI)), (UPat(dtype=dtypes.float64).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSD2SI)), @@ -518,7 +517,7 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: (UPat(dtype=dtypes.int16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWQ)), (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXDQ)), # bitcasts between scalar floats and ints - (UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, 0)))), + (UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, machine_const(0)))), (UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), vpins), (UPat(dtype=dtypes.int32s).bitcast(dtypes.float32).named("x"), lambda x: x.ins(X86Ops.VMOVD)), (UPat(dtype=dtypes.int64s).bitcast(dtypes.float64).named("x"), lambda x: x.ins(X86Ops.VMOVQ)), @@ -533,17 +532,17 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: (UPat(Ops.COPY, dtypes.floats, name="x"), lambda x: x.ins(_xmm_sz(x))), (UPat(Ops.COPY, dtypes.ints+(dtypes.bool,), name="x"), lambda x: x.ins(X86Ops.MOV) if x.max_numel() == 1 else x.ins(_xmm_sz(x))), (UPat(Ops.LOAD, dtypes.floats, src=(UPat(name="a"),), name="x"), lambda x,a: - x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else + x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (machine_const(0),)) if x.max_numel() * x.dtype.itemsize == 2 else x.ins(_xmm_sz(x), src=fold_address(a))), (UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.MOV, src=fold_address(a)) if x.max_numel() == 1 else - x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else + x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (machine_const(0),)) if x.max_numel() * x.dtype.itemsize == 2 else x.ins(_xmm_sz(x), src=fold_address(a))), (UPat.var("a").store(UPat.var("b", dtypes.floats), name="x"), lambda a,b,x: - x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, imm(dtypes.uint8, 0))) if b.max_numel() * b.dtype.itemsize == 2 else + x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, machine_const(0))) if b.max_numel() * b.dtype.itemsize == 2 else x.ins(_xmm_sz_m(b), src=fold_address(a) + (b,))), (UPat.var("a").store(UPat.var("b", dtypes.ints+(dtypes.bool,)), name="x"), lambda a,b,x: - x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, imm(dtypes.uint8, 0))) if b.max_numel() > 1 and b.max_numel() * b.dtype.itemsize == 2 else + x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, machine_const(0))) if b.max_numel() > 1 and b.max_numel() * b.dtype.itemsize == 2 else x.ins(_xmm_sz_m(b), src=fold_address(a) + (b,)) if b.max_numel() > 1 else x.ins(X86Ops.MOVm, src=fold_address(a) + (b,)) if (i:=to_imm(b)) is None else x.ins(X86Ops.MOVi, src=fold_address(a) + (i,))), # allocate virtual registers @@ -575,7 +574,7 @@ def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]: # loop, cmp on backedge all we need is a jmp tag if x.dtype is dtypes.void: return (label, [label]) else: - acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:]) + acc = x.ins(X86Ops.MOVi, src=(machine_const(0),) + x.src[1:]) cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0])) jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}") ctx.loop_label[acc] = loop_label @@ -584,7 +583,7 @@ def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]: def lower_end(ctx, x:UOp) -> tuple[UOp, list[UOp]]: end_label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}") jmp = UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}") - inc = x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),)) + inc = x.src[1].ins(X86Ops.ADDi, src=(machine_const(1),)) return (inc, [inc, jmp, end_label]) def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]: @@ -595,7 +594,7 @@ def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]: # final rewrite to match the isa spec post_regalloc_matcher = PatternMatcher([ # rewrite FRAME_INDEX to IMM now that the stack size is known - (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=x.const_like(ctx.stack_size + x.tag), [nx])), + (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=machine_const(ctx.stack_size + x.tag), [nx])), # expand the cmp here so we can preserve rng src edge to get label from ctx (UPat(Ops.INS, arg=X86Ops.LOOP_CMP, name="x"), lower_loop), # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound @@ -609,6 +608,9 @@ def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]: # ***** X86 instruction encoding ***** +# a machine const states a number, the field it lands in states the width. this writes the low bytes of that number, the way the hardware reads them +def low_bytes(v:int, sz:int) -> bytes: return (v % (1 << 8*sz)).to_bytes(sz, 'little') + def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> bytes|None: def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|None=None, sz_uop:UOp|None=None, vvvv_uop:UOp|None=None, imm_uop:UOp|None=None) -> bytes: @@ -653,9 +655,10 @@ def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|No if disp_uop is not None: disp_uop = unpair(disp_uop) assert disp_uop.op is Ops.CONST, "displacement must be a constant" - assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int" + # a displacement is 1 or 4 bytes, always the smallest that holds it + disp_sz = 4 if abs(disp_uop.val) > dtypes.int8.max else 1 # rbp/r13 always require a displacement - if disp_uop.val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10 + if disp_uop.val != 0 or rm == 0b101: mod = 0b01 if disp_sz == 1 else 0b10 else: mod = 0b00 else: mod = 0b11 # x 0b0 and idx 0b100 means rsp which means no index exists @@ -669,10 +672,13 @@ def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|No # DISP byte if mod == 0b01 or mod == 0b10: assert disp_uop is not None - inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.val) + inst += low_bytes(disp_uop.val, disp_sz) # IMM byte if imm_uop is not None: - if (c:=unpair(imm_uop)).op is Ops.CONST: inst += struct.pack(unwrap(c.dtype.fmt), c.val) + if (c:=unpair(imm_uop)).op is Ops.CONST: + # vex encoded immediates and shift counts are always 1 byte, everything else is the operand width capped at the widest immediate + imm_sz = 1 if sel or x.arg in {X86Ops.SHLi, X86Ops.SHRi, X86Ops.SARi} else min(sz, 4) + inst += low_bytes(c.val, imm_sz) elif isinstance(greg(imm_uop), Register): inst += bytes([(greg(imm_uop).index & 0b1111) << 4 | 0b0000]) return inst @@ -704,9 +710,9 @@ def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|No # prefix field: None -> 0 | 66 -> 1 | F3 -> 2 | F2 -> 3 # opcode map select: 0F -> 1 | 0F38 -> 2 | 0F3A -> 3 encodings = { - # moves + # moves NOTE: the movabs immediate is the only one as wide as the register, its width is the width of the definition X86Ops.MOVABS: lambda x: - bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].val), + bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + low_bytes(x.src[0].val, x.dtype.itemsize), X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0), X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D), X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1), From c4d2ec154dc9a77cab6b7fe0b52d726b42248ea5 Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Mon, 3 Aug 2026 15:58:24 -0400 Subject: [PATCH 4/4] done? --- test/null/test_gpudims.py | 2 +- test/null/test_graph_rewrite.py | 35 +++++----- test/null/test_simplify_valid_idx.py | 23 +++--- test/null/test_uop_graph.py | 12 ++-- test/null/test_uop_symbolic.py | 2 +- test/null/test_uop_vmin_vmax.py | 2 +- test/null/test_uops.py | 14 ++-- test/null/test_validate_oob.py | 6 +- test/unit/test_dtype_weak.py | 44 ++++++++++-- tinygrad/codegen/__init__.py | 59 +++++++--------- tinygrad/codegen/decomp/dtype.py | 18 +++-- tinygrad/codegen/decomp/op.py | 27 +++---- tinygrad/codegen/late/coalesce.py | 2 +- tinygrad/codegen/simplify.py | 4 +- tinygrad/renderer/cstyle.py | 21 ++++-- tinygrad/renderer/llvmir.py | 2 +- tinygrad/renderer/nir.py | 5 +- tinygrad/renderer/ptx.py | 15 ++-- tinygrad/schedule/rangeify.py | 2 +- tinygrad/uop/ops.py | 101 ++++++++++++--------------- tinygrad/uop/render.py | 7 +- tinygrad/uop/spec.py | 14 ++-- tinygrad/uop/symbolic.py | 39 +++++++---- 23 files changed, 254 insertions(+), 202 deletions(-) diff --git a/test/null/test_gpudims.py b/test/null/test_gpudims.py index 8d73d907bc1a0..c4ad4b332acf4 100644 --- a/test/null/test_gpudims.py +++ b/test/null/test_gpudims.py @@ -27,7 +27,7 @@ def _verify_indices_z3(self, idxs, dims): flat = UOp.const(0) for i, idx in enumerate(idxs): flat = flat + idx * int(math.prod(dims[i+1:])) - flat_p = flat.substitute({s: UOp(Ops.SPECIAL, src=s.src, arg=s.arg+"_p") for s in specials}) + flat_p = flat.substitute({s: UOp(Ops.SPECIAL, s.dtype, src=s.src, arg=s.arg+"_p") for s in specials}) solver = z3.Solver() [z3_flat, z3_flat_p] = uops_to_z3(solver, flat, flat_p) # bounds diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index a10d13ea8bc33..a3b33bede277b 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -22,6 +22,8 @@ def apply_rewrite_values(expr): def evaluate_uop(uop, variables): if uop.op == Ops.CONST: return uop.val + elif uop.op == Ops.CAST: + return uop.dtype.const(evaluate_uop(uop.src[0], variables)) elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU: return variables[uop.expr] elif uop.op in GroupOp.ALU: @@ -33,31 +35,27 @@ def evaluate_uop(uop, variables): class TestArithmeticSimplifications(unittest.TestCase): def test_full_graph_rewrite_division_by_zero(self): optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0)) - self.assertEqual(optimized_div_uop.op, Ops.CONST) - self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val)) + self.assertEqual((optimized_div_uop.op, optimized_div_uop.dtype, optimized_div_uop.src[0].op), (Ops.CAST, dtypes.float, Ops.CONST)) + self.assertTrue(math.isinf(optimized_div_uop.src[0].val) or math.isnan(optimized_div_uop.src[0].val)) def test_full_graph_rewrite_redundant_operations(self): optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0)) - self.assertEqual(optimized_uop.op, Ops.CONST) - self.assertEqual(optimized_uop.val, 10.0) + self.assertIs(optimized_uop, UOp.const(10.0).cast(dtypes.float)) def test_full_graph_rewrite_large_graph(self): prev_uop = UOp.const(0) for i in range(1, 101): prev_uop += UOp.const(i) optimized_uop = apply_rewrite(prev_uop) - self.assertEqual(optimized_uop.op, Ops.CONST) - self.assertEqual(optimized_uop.val, sum(range(1, 101))) + self.assertIs(optimized_uop, UOp.const(sum(range(1, 101))).cast(dtypes.int)) def test_full_graph_rewrite_division_by_one(self): optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0)) - self.assertEqual(optimized_uop.op, Ops.CONST) - self.assertEqual(optimized_uop.val, 42.0) + self.assertIs(optimized_uop, UOp.const(42.0).cast(dtypes.float)) def test_full_graph_rewrite_modulo_by_one(self): optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1)) - self.assertEqual(optimized_uop.op, Ops.CONST) - self.assertEqual(optimized_uop.val, 0) + self.assertIs(optimized_uop, UOp.const(0).cast(dtypes.int)) class TestFoldingAndReduction(unittest.TestCase): @@ -109,26 +107,24 @@ def test_full_graph_rewrite_modulo_folding_with_define_var(self): # index dtype because div-mod rules only work on index x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint) optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4) - self.assertEqual(optimized_mod_uop.op, Ops.CONST) - self.assertEqual(optimized_mod_uop.val, 2) + self.assertIs(optimized_mod_uop, UOp.const(2).cast(dtypes.int)) def test_full_graph_rewrite_division_folding_with_define_var(self): # index dtype because div-mod rules only work on index n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint) optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3) self.assertEqual(optimized_div_uop.op, Ops.MUL) - self.assertEqual(optimized_div_uop.src[1].val, 2) + self.assertIs(optimized_div_uop.src[1], UOp.const(2).cast(dtypes.int)) def test_full_graph_rewrite_complex_mod_div_folding(self): # index dtype because div-mod rules only work on index k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint) optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2) - self.assertEqual(optimized_div_uop.op, Ops.CONST) - self.assertEqual(optimized_div_uop.val, 1) + self.assertIs(optimized_div_uop, UOp.const(1).cast(dtypes.int)) def test_graph_rewrite_div_folding_bug(self): lhs = UOp(Ops.ADD, src=( - UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4), + UOp(Ops.STACK, arg=None, src=(UOp.special(32, 'lidx0'),)*4), UOp.const((0, 256, 512, 768)))) rhs = UOp.const((2,)*4) unopt = lhs 0, - f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}") + self.assertEqual((optimized_log2_neg.op, optimized_recip_zero.op), (Ops.CAST, Ops.CAST)) + self.assertTrue(math.isnan(optimized_log2_neg.src[0].val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.src[0].val}") + self.assertTrue(math.isinf(optimized_recip_zero.src[0].val) and optimized_recip_zero.src[0].val > 0, + f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.src[0].val}") @unittest.skip("broken") def test_full_graph_rewrite_modulo_negative_dividend(self): diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 15a9ce9171201..ed2e7f57c9871 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -2,7 +2,7 @@ from tinygrad.codegen.late.coalesce import indexing_simplify from tinygrad.dtype import dtypes -from tinygrad.uop.ops import UOp, Ops, graph_rewrite, pm_lower_index_dtype +from tinygrad.uop.ops import UOp, Ops, graph_rewrite, pm_address_demand from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load from tinygrad.helpers import Context from test.helpers import full_rewrite @@ -23,7 +23,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)), )) -def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr) +def Special(expr, nmax): return UOp.special(nmax, expr) def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax) def Range(n, nmax): return UOp.range(nmax, n) @@ -493,13 +493,12 @@ def test_drop_non_monotonic_window(self): self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0") def test_drop_gate_committed_in_the_index_pass(self): - # the fused index pass runs without symbolic, so committing a weak src must not leave a CAST that - # symbolic later folds inside the index only: the gate's copy of the expression has to stay the same node + # the last indexing value round precedes its commit, so the gate and index keep the same expression f = UOp.variable("f", 0.0, 9.0, dtypes.float) idx_y = (f + UOp.const(1.0)).cast(dtypes.int) load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)), (Special("gidx0", 10), idx_y)) - off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0] + off = graph_rewrite(graph_rewrite(load.sink(), indexing_simplify, ctx={}), pm_address_demand, bottom_up=True).src[0].src[0] self.assertEqual(off.src[1].get_valid(), UOp.const(True)) class TestDropTrueGate(unittest.TestCase): @@ -529,7 +528,7 @@ def test_range_shrink_single_guard(self): load = get_gated_load_uop(r < UOp.const(4), r) ranges = self.get_ranges(load.sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].val, 4) + self.assertIs(ranges[0].src[0], UOp.const(4).cast(dtypes.int)) def test_range_shrink_picks_max_guard(self): # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8 @@ -538,7 +537,7 @@ def test_range_shrink_picks_max_guard(self): load2 = get_gated_load_uop(r < UOp.const(8), r) ranges = self.get_ranges(UOp.sink(load1, load2)) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].val, 8) + self.assertIs(ranges[0].src[0], UOp.const(8).cast(dtypes.int)) def test_range_no_shrink_guard_ge_max(self): # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain) @@ -546,7 +545,7 @@ def test_range_no_shrink_guard_ge_max(self): load = get_gated_load_uop(r < UOp.const(300), r) ranges = self.get_ranges(load.sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].val, 204) + self.assertIs(ranges[0].src[0], UOp.const(204).cast(dtypes.int)) def test_range_no_shrink_when_unguarded_elsewhere(self): # one load guards r < 4, but another load uses r without a gate -> no shrink @@ -555,7 +554,7 @@ def test_range_no_shrink_when_unguarded_elsewhere(self): load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),)) ranges = self.get_ranges(UOp.sink(load1, load2)) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].val, 204) + self.assertIs(ranges[0].src[0], UOp.const(204).cast(dtypes.int)) def test_range_no_shrink_when_used_in_reduce(self): # range used in both a gated load AND directly in the reduce expression -> no shrink @@ -564,7 +563,7 @@ def test_range_no_shrink_when_used_in_reduce(self): red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD) ranges = self.get_ranges(red.sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].val, 204) + self.assertIs(ranges[0].src[0], UOp.const(204).cast(dtypes.int)) def test_range_shrink_to_single_iteration(self): # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely @@ -580,7 +579,7 @@ def test_range_shrink_store_where_invalid(self): x = (r < 4).where(UOp.const(1.0), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].val, 4) + self.assertIs(ranges[0].src[0], UOp.const(4).cast(dtypes.int)) def test_range_shrink_store_where_invalid_flipped(self): # above, but flipped @@ -589,7 +588,7 @@ def test_range_shrink_store_where_invalid_flipped(self): x = (r < 4).where(UOp.const(1.0), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].val, 4) + self.assertIs(ranges[0].src[0], UOp.const(4).cast(dtypes.int)) if __name__ == '__main__': unittest.main() diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index e889df8611bca..ac2a946d9588c 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -325,7 +325,7 @@ def test_cast_alu_fold(self): alu = (ld<1).cast(dtypes.bool) out = d0.index(idx).store(alu) uops = to_uops_list([out]) - self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0) + self.assertTrue(all(x.src[0].op is Ops.CONST for x in uops if x.op is Ops.CAST)) def test_double_cast_fold(self): d0 = UOp.param(0, dtypes.float, (1,)) @@ -335,7 +335,7 @@ def test_double_cast_fold(self): alu = ld.cast(dtypes.float).cast(dtypes.float) out = d0.index(idx).store(alu) uops = to_uops_list([out]) - self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1) + self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 1) def test_depth_2_const_fold(self): v = UOp.variable("tmp", 0, 1, dtypes.int) @@ -372,7 +372,7 @@ def test_where_on_gated_load_fold(self): uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE - if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5 + if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1] is UOp.const(5).cast(dtypes.long) def test_where_on_gated_load_folds_swapped_branches(self): ridx0 = UOp.range(100, 0) @@ -382,7 +382,7 @@ def test_where_on_gated_load_folds_swapped_branches(self): uops = to_uops_list([w]) for u in uops: assert u.op is not Ops.WHERE - if u.op is Ops.LOAD: assert u.src[1].val==5 + if u.op is Ops.LOAD: assert u.src[1] is UOp.const(5).cast(dtypes.long) def test_where_on_gated_load_with_cast(self): ridx0 = UOp.range(100, 0) @@ -394,7 +394,7 @@ def test_where_on_gated_load_with_cast(self): uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE - if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 5 + if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1] is UOp.const(5, dtypes.int) def test_where_on_casted_gated_load_extra_cond(self): ridx0 = UOp.range(100, 0) @@ -426,7 +426,7 @@ def test_where_in_store_becomes_gate(self): uops = to_uops_list([st]) for u in uops: assert u.op is not Ops.WHERE - if u.op is Ops.STORE: assert u.src[1].val==5 + if u.op is Ops.STORE: assert u.src[1] is UOp.const(5).cast(dtypes.long) def test_load_idx_becomes_int(self): # mnist indexing with split reduceop diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 36629e025a650..d6ddf5792c42f 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -1021,7 +1021,7 @@ def test_where_cast(self): # the vars are now scalar PARAMs pvar = {u.expr: u for u in rewritten_uop.toposort() if u.op is Ops.PARAM} - self.assertEqual(rewritten_uop, (pvar['s']>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29)) - self.assertNotIn(Ops.CAST, ops) + self.assertTrue(all(u.src[0].op is Ops.CONST for u in uops if u.op is Ops.CAST)) @unittest.expectedFailure def test_fast_idiv_overflow(self): @@ -395,7 +395,7 @@ def test_uop_variables(self): self.assertEqual(list(var_vals)[0], a.expr) def test_const_factor(self): - gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(8),), arg='gidx0') + gidx0 = UOp.special(8, 'gidx0') self.assertEqual(UOp.const(17).const_factor(), 17) self.assertEqual(gidx0.const_factor(), 1) self.assertEqual((gidx0*3).const_factor(), 3) diff --git a/test/null/test_validate_oob.py b/test/null/test_validate_oob.py index 2d52f8e70e044..25b6693cbd07f 100644 --- a/test/null/test_validate_oob.py +++ b/test/null/test_validate_oob.py @@ -150,7 +150,7 @@ def test_load_bool_as_mask(self): with Context(CHECK_OOB=1, SPEC=2): buf_bool = UOp.param(0, dtypes.bool, (16,)) buf_int = UOp.param(1, dtypes.int, (8,)) - gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0") + gidx = UOp.special(16, "gidx0") ld_bool = buf_bool.index(gidx).load() with self.assertRaises(RuntimeError): to_uops_list([buf_int.index(gidx.valid(ld_bool)).load()]) # gidx 0..15, buf_int size 8 @@ -164,8 +164,8 @@ def test_in_bounds_access_gated_local(self): sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL) # Define indices, valids and barrier - gidx = UOp(Ops.SPECIAL, src=(UOp.const(416),), arg="gidx0") - lidx = UOp(Ops.SPECIAL, src=(UOp.const(10),), arg="lidx0") + gidx = UOp.special(416, "gidx0") + lidx = UOp.special(10, "lidx0") gate = (gidx<400) & (lidx<8) diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 30ce55fe656e2..1223be1d54d87 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -2,10 +2,14 @@ from tinygrad import Tensor, dtypes, TinyJit from tinygrad.helpers import Context -from tinygrad.dtype import least_upper_float -from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite, pm_lower_index_dtype, pm_commit_weak +from tinygrad.dtype import least_upper_float, AddrSpace, Invalid +from tinygrad.helpers import Target +from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite, pm_lower_index_dtype, pm_boundary_demand, pm_commit_weak from tinygrad.uop.symbolic import symbolic_simple -from tinygrad.uop.spec import spec_shared, type_verify +from tinygrad.uop.spec import spec_shared, spec_tensor, spec_program, type_verify +from tinygrad.codegen import to_program +from tinygrad.renderer.cstyle import CStyleLanguage +from tinygrad.renderer.wgsl import WGSLRenderer from tinygrad.engine.jit import JitError @@ -83,8 +87,7 @@ def test_store_weak_value_uses_destination_dtype(self): dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32)) gate = UOp.const(True) out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_lower_index_dtype, ctx={}) - # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it - self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0, dtypes.bfloat16), gate)) + self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0).cast(dtypes.bfloat16), gate)) def test_weak_srcs_commit_only_at_a_concrete_lub(self): weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0))) @@ -96,7 +99,36 @@ def test_weak_srcs_commit_only_at_a_concrete_lub(self): def test_weak_shift_lhs_commits_the_node(self): # a shift derives its lhs's dtype, so committing the lhs restates the root (WGSL's packed store writes `mask << shift_am`) shl = graph_rewrite(UOp.const(0xFFFF) << UOp.variable("x", 0, 16, dtypes.uint), symbolic_simple+pm_commit_weak) - self.assertEqual((shl.dtype, shl.src[0]), (dtypes.uint, UOp.const(0xFFFF, dtypes.uint))) + self.assertEqual((shl.dtype, shl.src[0]), (dtypes.uint, UOp.const(0xFFFF).cast(dtypes.uint))) + + def test_program_allows_only_pair_value_weakness(self): + type_verify(UOp.const(1), spec_program) + type_verify(UOp.const(1).cast(dtypes.int).sink(), spec_program) + type_verify(UOp.param(0, dtypes.float, (3,)), spec_program) # PARAM shape is widthless metadata + type_verify(graph_rewrite(UOp.special(3, "gidx0"), pm_boundary_demand), spec_program) # SPECIAL size is widthless metadata + with self.assertRaises(RuntimeError): type_verify(UOp.const(1).sink(), spec_program) + with self.assertRaises(RuntimeError): type_verify((UOp.const(1)+UOp.const(2)).sink(), spec_program) + + def test_special_source_dtype_agrees(self): + # under SPEC=2 the construction itself is rejected, so build inside the context that expects the raise + with self.assertRaises(RuntimeError): + type_verify(UOp(Ops.SPECIAL, dtypes.int64, src=(UOp.const(8, dtypes.int8),), arg="gidx0"), spec_tensor) + type_verify(UOp(Ops.SPECIAL, dtypes.int64, src=(UOp.const(Invalid),), arg="gidx0"), spec_tensor) + + def test_cstyle_reads_int_pair_and_lane(self): + ren, pair = CStyleLanguage(Target()), UOp.const(2).cast(dtypes.int) + self.assertEqual(ren.string_rewrite.rewrite(pair, ctx=ren), "2") + buf = UOp.param(0, dtypes.float, (4,), addrspace=AddrSpace.ALU) + ren.r = {buf: "vec", pair: "2"} + self.assertEqual(ren.render_index(buf.index(pair), buf, pair), "vec.z") + + def test_pre_boundary_decomps_keep_shifts(self): + # SPEC=1: wgsl's packed_load builds a uint LOAD over a ushort INDEX, which SPEC=2 rejects (master does this too) + with Context(EMULATED_DTYPES="bfloat16", SPEC=1): + out = (Tensor.empty(8, dtype=dtypes.bfloat16, device="NULL")+Tensor.empty(8, dtype=dtypes.bfloat16, device="NULL")).contiguous() + src = to_program(out.schedule_linear().src[-1].src[0], WGSLRenderer(Target("WEBGPU"))).src[2].arg + for shift in ("<<16u", ">>16u", ">>7u"): self.assertIn(shift, src) + for arithmetic in ("*65536u", "/65536u", "/128u"): self.assertNotIn(arithmetic, src) @unittest.expectedFailure # TODO: a weak const defers to its consumer (JAX): these dtypes change once python scalars are weak consts def test_changed_rows(self): diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 063730b75c075..295bc57bf6d65 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -3,7 +3,7 @@ from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp -from tinygrad.uop.ops import AxisType, pm_commit_weak, pm_cast_weak +from tinygrad.uop.ops import AxisType, pm_address_demand, pm_boundary_demand, pm_commit_weak, pm_lower_weak from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program from tinygrad.renderer import Renderer, Estimates @@ -12,7 +12,8 @@ # import all pattern matchers here from tinygrad.codegen.gpudims import pm_add_gpudims -from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid +from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic +from tinygrad.uop.symbolic import pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid from tinygrad.uop.movement import mop_cleanup from tinygrad.codegen.decomp.dtype import pm_dtype_decomps from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns @@ -30,10 +31,9 @@ from tinygrad.uop.ops import _broadcast_shape, identity_element from tinygrad.schedule.rangeify import BufferizeOpts -def do_number_param(ctx:list[int], x:UOp): +def do_number_param(ctx:itertools.count, x:UOp): if x.arg.slot != -1: return None - ctx[0] += 1 - return x.replace(arg=replace(x.arg, slot=ctx[0]-1)) + return x.replace(arg=replace(x.arg, slot=next(ctx))) pm_number_params = PatternMatcher([ (UPat(Ops.PARAM, name="x"), do_number_param), @@ -343,39 +343,31 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # ***** THE BOUNDARY: weakness ends here ***** # above this line a value may be widthless; below it every width is explicit and a rule that mints a const states its width. - # pm_commit_weak runs before lower_weak_srcs inside pm_lower_index_dtype and that order is load-bearing (the reverse costs 60 null tests). - # pm_lower_weak must never be merged into a matcher that also folds CAST-of-CONST: the marker it mints is folded straight back, a 2-cycle. + # the decomps are backend legalization, so they run below: a const-keyed rule down here reads the pair, it does not re-widen it. # NOTE: we need indexing_simplify to remove the cast to long using the Invalid - sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes") + sink = graph_rewrite(sink, indexing_simplify, name="final indexing simplification") + sink = graph_rewrite(sink, pm_address_demand, name="commit addresses", bottom_up=True) + # a consumer's demand decides an index's width first; ROLE is the fallback for the indices no consumer resolved + sink = graph_rewrite(sink, pm_lower_index_dtype, name="lower all index dtypes") + sink = graph_rewrite(sink, pm_boundary_demand, name="commit remaining roles") - # final symbolic before decomp - sink = graph_rewrite(sink, symbolic, name="final symbolic") + sink = graph_rewrite(sink, symbolic+pm_cast_float_alu, name="final symbolic") - sink = graph_rewrite(sink, pm_cast_float_alu, name="cast float alu operands") - - # **** decomps **** - - # floordiv+mod / dtype decomp (early) + # below the boundary weakness re-enters only from minting rounds, so every decomp round carries the commit that closes it supported_ops = tuple(ren.code_for_op.keys()) - pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops) - sink = graph_rewrite(sink, pm_decomp, name="early decompositions") - - # below the boundary weakness re-enters only from minting rounds, and each is closed by the commit round that follows it: - # "early decompositions" mints -> "decomp dtypes" commits (pm_commit_weak) - # "late decompositions" mints -> "final rewrite" commits (pm_commit_weak, which is why it heads that sum) - # a new minting round below here needs a commit round after it, or its weak nodes reach spec_program and are rejected - sink = graph_rewrite(sink, pm_dtype_decomps+pm_commit_weak, ctx=(set(), ren), name="decomp dtypes") - pm_decomp = pm_decomp+\ - get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\ - get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2) - sink = graph_rewrite(sink, pm_decomp, ctx=ren, name="late decompositions") + pm_commit = pm_commit_weak+pm_lower_weak + pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops)+get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV)) + sink = graph_rewrite(sink, pm_decomp+pm_commit, ctx=ren, name="early decompositions") + + sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes") + sink = graph_rewrite(sink, pm_commit, name="commit decomp dtypes") + pm_decomp = pm_decomp+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2) + sink = graph_rewrite(sink, pm_decomp+pm_commit, ctx=ren, name="late decompositions") sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index") - # final rules for the renderer (without sym) - extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([]) - # pm_commit_weak heads this sum: renderer matchers in extra_matcher may mint weak, and this is the last round that can commit them - pm_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends - sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite") + # renderer matchers in extra_matcher may mint weak too, so the last round commits as well + sink = graph_rewrite(sink, pm_decomp+pm_commit+(ren.extra_matcher or PatternMatcher([]))+pm_split_ends+pm_remove_invalid, + ctx=ren, name="final rewrite") # add implicit barriers (stores/loads through LOCAL memory ordered by AFTER or across loop iterations need workgroup barriers) sink = graph_rewrite(sink, pm_implicit_barriers, name="add implicit barriers") @@ -385,7 +377,8 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # put unnumbered variable PARAMs in slots num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1]) - sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True) + sink = graph_rewrite(sink, pm_number_params+pm_commit_weak+pm_lower_weak, ctx=itertools.count(num_params), + name="number params and final commit", walk=True) if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST") if SPEC: type_verify(sink, spec_program) diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index cbb324ace05c1..e118d639becb6 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -11,10 +11,13 @@ l2i_dt = {dtypes.long: dtypes.int, dtypes.ulong: dtypes.uint} def unpack32(v:UOp) -> tuple[UOp, UOp]: return v.bitcast(dtypes.uint) & 0xFFFF, shr(v.bitcast(dtypes.uint), 16) def reindex(idx:UOp, off:int, mul=2) -> UOp: + # the word this becomes is stated identity-free: a unit scale and a zero offset are not written down + new_idx = idx.src[1]*mul if mul != 1 else idx.src[1] + if off: new_idx = new_idx+off if idx.op is Ops.SHRINK: assert mul == 1, "can't reindex SHRINK with mul != 1" - return idx.replace(op=Ops.INDEX, src=(idx.src[0], idx.src[1]+off)) - return idx.replace(src=(idx.src[0], idx.src[1]*mul+off, *idx.src[2:])) + return idx.replace(op=Ops.INDEX, src=(idx.src[0], new_idx)) + return idx.replace(src=(idx.src[0], new_idx, *idx.src[2:])) # 4.3.1 is the relevant section in TAOCP def l2i(op: Ops, dt: DType, *uops:UOp): @@ -131,7 +134,15 @@ def f2f_store(st, idx, val, fr:DType, to:DType): return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.index(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n))) # tag is the 32-bit word this node becomes - (0 for the low word, 1 for the high, the dtype the consumer wants) +word_tags = {(w, dt) for w in (0, 1) for dt in l2i_dt.values()} +def const_word(x:UOp) -> UOp: + # the committed pair states a long constant's width on its CAST, so the value is read through it + v, (word, dt) = int((x.src[0] if x.op is Ops.CAST else x).val), x.tag + return UOp.const(truncate[dt]((v >> 32) if word == 1 else (v & 0xFFFFFFFF)), dt) + pm_long_decomp = PatternMatcher([ + # a long constant is a constant however it is spelled: take its word directly, never decompose the CAST that commits it + (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat(Ops.CONST),), tag=word_tags, name='x'), const_word), (UPat(GroupOp.Defines, src=(UPat.var("sz"),), name="x"), lambda x,sz: x.replace(dtype=l2i_dt[x.dtype], arg=replace(x.arg, dtype=l2i_dt[x.dtype]), src=(sz*2,)) if x.dtype in l2i_dt else None), (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: @@ -157,8 +168,7 @@ def f2f_store(st, idx, val, fr:DType, to:DType): if x.tag is not None else None), (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx: x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None), - (UPat(Ops.CONST, tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), lambda x: - UOp.const(truncate[x.tag[1]]((x.val >> 32) if x.tag[0] == 1 else (x.val & 0xFFFFFFFF)), x.tag[1])) + (UPat(Ops.CONST, tag=word_tags, name='x'), const_word) ]) # float decomposition patterns - ctx is (fr, to) tuple diff --git a/tinygrad/codegen/decomp/op.py b/tinygrad/codegen/decomp/op.py index 6a48cdca53cc0..513a36671100b 100644 --- a/tinygrad/codegen/decomp/op.py +++ b/tinygrad/codegen/decomp/op.py @@ -1,7 +1,7 @@ from typing import Callable import functools from tinygrad.dtype import dtypes -from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher +from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, cvar from tinygrad.renderer import Renderer # *** integer division *** @@ -77,7 +77,8 @@ def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher: # these are rewrites that make things simpler pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)] # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod - if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None)) + if Ops.AND in ops: + pat.append((UPat.var("x", dtypes.ints)%cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None)) pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod)) # no real hardware supports THREEFRY, but NullRenderer does if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32)) @@ -91,41 +92,41 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(), lambda x,y: (x | y).logical_not())] # rewrite MUL/CDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y) - if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.val, 0)) else None)] + if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.val, 0)) else None)] if Ops.SHR in ops: # uint CDIV by 2**v -> x >> v (FLOORDIV is lowered to CDIV by the rule above before reaching here) - pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.uints), UPat.cvar("c"))), + pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.uints), cvar("c"))), lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None)] # signed CDIV (trunc) by 2**v -> (x + (x<0 ? c-1 : 0)) >> v - pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("c"))), + pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), cvar("c"))), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v if (v:=powers_of_two.get(c.val, 0)) else None)] if not disable_fast_idiv: # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0) - pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))), + pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), cvar("d"))), lambda ctx, x, d: fast_idiv(ctx, x, d.val) if x.vmin >= 0 or x.dtype in dtypes.uints else None)] # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs; # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail) pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))), lambda x, d: x - d * x.alu(Ops.CDIV, d) if x.vmin >= 0 or x.dtype in dtypes.uints else None)] if Ops.NEG in ops: - pat += [(UPat.var('x')*-1, lambda ctx,x: x.alu(Ops.NEG))] + pat += [(UPat.var('x')*cvar(arg=-1), lambda ctx,x: x.alu(Ops.NEG))] if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda ctx,x,y: x.alu(Ops.SUB, y))] if Ops.CMPLT in ops: # These are late rewrites because simplex expects equalities to be a certain format pat += [ - ((UPat.var("x", dtypes.sints) < UPat.cvar("c")).logical_not(), lambda x,c: c-1 x==c ] if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))] if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))] # also fuse (x << n) + c → MULACC(x, 2^n, c) since MUL→SHL may run first - if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1< a/b if Ops.FDIV in ops: pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))] diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 47ec328017ac6..523a863963ded 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -141,7 +141,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])] for full_grp in grouped_offsets: while len(full_grp): - offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(full_grp[0]) + offset = (base+full_grp[0] if full_grp[0] else base) if isinstance(base, UOp) else UOp.const(full_grp[0]) length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0] grp = full_grp[:length] # NOTE: we apply the valid again after we determine the length diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 4aadeb6d6feef..604f2c0f25266 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -123,10 +123,10 @@ def reduce_unparented(red:UOp) -> UOp|None: pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([ # lift x+y out of reduce on ne - ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None), + ((UPat.var("x")+UPat.var("y")).or_casted()!=UPat.var("c"), lambda x,y,c: x.cast(c.dtype)!=c-y.cast(c.dtype) if all(map(no_range,(y,c))) else None), # reduce on gated load becomes can substitute the range and remove the reduce ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD), - lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)), + lambda r,idx,expr: (v:=(idx >= 0) & (idx < r.src[0])).where(expr.substitute({r:idx.valid(v)}),0)), ]) def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp|None: diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 1e3872c819f5b..72d4c881e51e7 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -39,6 +39,7 @@ lambda ctx,x,c: f"({ctx.render_cast(x, f'{c.val}u')})"), (UPat(Ops.CAST, (dtypes.int8, dtypes.int16), src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, str(c.val))})"), + (UPat(Ops.CAST, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),)), lambda ctx,c: str(c.val)), # const (UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, ctx.infinity)})"), @@ -62,11 +63,11 @@ # GPU stuff (UPat(Ops.BARRIER), lambda ctx: ctx.barrier), - (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0]](x.arg[-1])}; /* {(x.src[0]).render()} */"), + (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0]](x.arg[-1])}; /* {x.src[0].ssimplify()} */"), # SHRINK/INDEX (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), - (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat.cvar()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), + (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat(Ops.CONST).or_casted()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), (UPat(Ops.STACK, name="x"), lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_type(x))}" + \ f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"), @@ -133,6 +134,7 @@ def wmma_args(uops:list[UOp]): for uop in uops if uop.op is Ops.WMMA) class CStyleLanguage(Renderer): + inline_pairs = True abi: str = "" kernel_typedef: str = "void" buffer_prefix: str = "" @@ -180,8 +182,9 @@ def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str def render_index(self, x:UOp, buf:UOp, idx:UOp): if buf.addrspace == AddrSpace.ALU: # this is lane access in C - if idx.op is not Ops.CONST: return f"({self[buf]})[{self[idx]}]" - return self[buf]+(f"[{idx.val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.val]}") + const_idx = idx.src[0] if idx.op is Ops.CAST and idx.src[0].op is Ops.CONST and idx.src[0].dtype in dtypes.weaks else idx + if const_idx.op is not Ops.CONST: return f"({self[buf]})[{self[idx]}]" + return self[buf]+(f"[{const_idx.val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[const_idx.val]}") return f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})" def render_buffer(self, x:UOp): @@ -227,7 +230,8 @@ def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[ c: defaultdict[str, int] = defaultdict(int) name = "test" for u in uops: - if u.op in {Ops.NOOP, Ops.GROUP}: continue + if u.op is Ops.CONST and u.dtype in dtypes.weaks: r[u] = str(u.val) + if u.op in {Ops.NOOP, Ops.GROUP} or u.op is Ops.CONST and u.dtype in dtypes.weaks: continue if u.op == Ops.STACK and len(u.src) == 0: continue if u.op is Ops.AFTER: r[u] = r[u.src[0]] @@ -253,10 +257,13 @@ def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[ assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}" if u.op in {Ops.ENDIF, Ops.END}: depth -= 1 - if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \ + pair = u.op is Ops.CAST and u.src[0].op is Ops.CONST and u.src[0].dtype in dtypes.weaks + can_inline = u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \ (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \ (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \ - (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))): + (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA")) or \ + (self.inline_pairs and pair) + if (u.op is not Ops.CAST or u.max_numel() == 1) and can_inline: r[u] = l else: if u.op not in {Ops.RANGE, Ops.STORE, Ops.BUFFER} and u.dtype != dtypes.void: diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 5f99606caff44..07f84708a8577 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -66,7 +66,7 @@ def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str: (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat((Ops.BUFFER, Ops.PARAM, Ops.AFTER)),), allow_any_len=True, name="x"), lambda ctx,x: f" {ctx[x]} = getelementptr inbounds {ldt(x.dtype)}, {ldt(x.dtype, ptr=True)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"), # register index - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("idx")), name="x"), lambda ctx,buf,idx,x: + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("idx").or_casted()), name="x"), lambda ctx,buf,idx,x: f" {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.val}" if buf.addrspace == AddrSpace.ALU else None), # load/store diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 2a7d04a9c80c2..1de08352d6d55 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -187,7 +187,8 @@ def postrender(self, uops:list[UOp]): pass def render(self, uops:list[UOp]): self.prerender(uops) - for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].val + for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: + self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].ssimplify() self.r: dict[UOp, Any] = {} self.param_idx = 0 ranges: list[mesa.nir_def|None] = [] @@ -196,7 +197,7 @@ def render(self, uops:list[UOp]): if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.STACK and len(u.src) == 0) or (u.op is Ops.CONST and u.dtype in dtypes.weaks): pass elif u.op in {Ops.INDEX, Ops.SHRINK}: # INDEX on a register value picks the element, memory INDEX is handled in the LOAD/STORE patterns - if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].val) + if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].ssimplify()) elif u.op is Ops.AFTER: self.r[u] = self.r[u.src[0]] elif u.op == Ops.SINK: diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index c1215aa97e600..420c86c83fc50 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -83,8 +83,6 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i # the pair reads as the typed constant it commits to (UPat(Ops.CAST, dtypes.bool, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), lambda ctx,x,c: f"setp.ne.s16 {ctx.r[x]}, {render_val(c.val, x.dtype)}, 0;"), - (UPat(Ops.CAST, src=(UPat(Ops.CONST, dtypes.weaks, name="c"),), name="x"), - lambda ctx,x,c: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(c.val, x.dtype)};"), (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.val, x.dtype)};"), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"), (UPat(Ops.PARAM, name="x"), lambda ctx, x: @@ -206,9 +204,14 @@ def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str: continue if u.op in {Ops.INDEX, Ops.SHRINK, Ops.LOAD} and u.src[0].addrspace in (AddrSpace.REG, AddrSpace.ALU): # on REG, INDEX/SHRINK pick the register (must be CONST) and LOAD is a noop - if u.op is not Ops.LOAD and u.src[1].op is not Ops.CONST: + if u.op is not Ops.LOAD and (idx:=u.src[1].src[0] if u.src[1].op is Ops.CAST else u.src[1]).op is not Ops.CONST: raise RuntimeError(f"PTX does not support dynamic register indexing: {u}") - r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].val] + r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][idx.val] + continue + if u.op is Ops.CONST and u.dtype in dtypes.weaks: continue + pair = u.op is Ops.CAST and (wc:=u.src[0]).op is Ops.CONST and wc.dtype in dtypes.weaks + if pair and u.dtype is not dtypes.bool: + r[u] = render_val(wc.val, u.dtype) continue if u.op is Ops.SPECIAL: r[u] = "%" + u.arg elif u.op is Ops.LOAD: @@ -220,9 +223,11 @@ def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str: [ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[1]]), 4 // u.src[0].dtype.scalar().itemsize)], [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.scalar().itemsize)]] r[u] = [ssa("wmma", dtype=self.types[u.dtype.scalar()]) for _ in range(u.max_numel())] + prefix_op = Ops.CONST if pair else u.op prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.END: ("pred", "pred"), Ops.RANGE: ("ridx", None), Ops.CONST: ("const", None), Ops.BUFFER: ("local", "u64"), Ops.INDEX: ("bidx", "u64"), Ops.SHRINK: ("bidx", "u64"), - Ops.PARAM: ("dat", "u64" if u.addrspace is AddrSpace.GLOBAL else None), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None)) + Ops.PARAM: ("dat", "u64" if u.addrspace is AddrSpace.GLOBAL else None), + **{op: ("alu", None) for op in GroupOp.ALU}}.get(prefix_op, (None, None)) if u.op is Ops.RANGE and u.dtype == dtypes.void: prefix = None # loop headers don't have a register if prefix: r[u] = ssa(prefix, u, dtype) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 3c4e3750aada5..db59f959a0520 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -361,7 +361,7 @@ def visitor(u:UOp) -> frozenset[UOp]: def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True): size = prod(x.shape) - dtype = strong_dtype(x.dtype) # a BUFFER is never weak: store at the concrete dtype, the .cast(x.dtype) on the result keeps readers unchanged + dtype = strong_dtype(x.dtype) # a BUFFER is never weak: store at the concrete dtype, the result cast keeps readers unchanged rngs = sorted(idx.ranges, key=lambda x: x.arg) assert size > 0 and isinstance(size, int), f"no zero sized or symbolic sized buffers {size}" diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 0b20d84b245c9..98215279c7344 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -131,9 +131,12 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None: 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.SPECIAL: + # SPECIAL's dtype is stated explicitly: src[0] is widthless launch metadata + return None case Ops.LOAD | Ops.INDEX | Ops.UNSHARD | Ops.REDUCE | Ops.AFTER | Ops.RANGE | \ Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.COPY | Ops.STAGE | Ops.DETACH | \ - Ops.MSTACK | Ops.MSELECT | Ops.ALLREDUCE | Ops.SPECIAL: + Ops.MSTACK | Ops.MSELECT | Ops.ALLREDUCE: # pass through first return src[0].dtype case Ops.CMPLT | Ops.CMPNE | Ops.CMPEQ: @@ -526,7 +529,11 @@ def simplify(self, tracked=False): from tinygrad.uop.symbolic import symbolic with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value): return graph_rewrite(self, symbolic, name="simplify") - def ssimplify(self) -> UOp|ConstType: return ret.val if (ret:=self.simplify()).op is Ops.CONST else ret + def ssimplify(self) -> UOp|ConstType: + ret = self.simplify() + if ret.op is Ops.CONST: return ret.val + if ret.op is Ops.CAST and ret.src[0].op is Ops.CONST: return ret.src[0].val + return ret def _eval(self, dtype, expected_type:Type[T]) -> T: assert self.dtype in dtype, f"eval with wrong dtype {self}" vmin, vmax = (simple_self:=self.simplify())._min_max @@ -628,7 +635,7 @@ def range(end:sint, axis_id, axis_type=AxisType.WEAK, *arg, dtype=dtypes.weakint @staticmethod def loop(axis_id:int, *arg): return UOp(Ops.RANGE, src=(UOp(Ops.NOOP),), arg=(axis_id, AxisType.WEAK)+arg) @staticmethod - def special(end:sint, name:str, dtype=dtypes.weakint): return UOp(Ops.SPECIAL, src=(sint_to_uop(end, dtype),), arg=name) + def special(end:sint, name:str, dtype=dtypes.weakint): return UOp(Ops.SPECIAL, dtype, src=(sint_to_uop(end),), arg=name) @staticmethod def wmma(a:UOp, b:UOp, acc:UOp, dims:tuple[int, int, int], device:str, threads:int, tc_upcast_axes=None): # dtype_in is stored in the arg (not derived from src[0].dtype) because bitcast rewrites change src dtypes @@ -1417,6 +1424,9 @@ def match(self:UPat, uop:UOp, store:dict[str, UOp]) -> list[dict[str, UOp]]: res.extend(stores) return res +# below the boundary a constant may be committed: the pair CAST(dt, weak CONST) denotes the same number as the bare const +def cvar(name:str|None=None, arg=None) -> UPat: return UPat.cvar(name, arg=arg).or_casted() + def deconstruct_function(fxn:Callable) -> tuple: new_globals = {k:v for k,v in fxn.__globals__.items() if k in fxn.__code__.co_names} for co in fxn.__code__.co_consts: @@ -1757,63 +1767,42 @@ def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x. def select_dtype(u:UOp): if u.dtype is dtypes.weakfloat: return dtypes.default_float return dtypes.long if u.overflows(dtypes.int32) else dtypes.int -def lower_weak_node(u:UOp) -> UOp|None: - start, src = (1 if u.op is Ops.WHERE else 0), tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src) - if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None - dt = strong_dtype(least_upper_dtype(select_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.cast(dt) for s in src[start:])).cast(u.dtype) +def lower_weak_srcs(u:UOp) -> UOp|None: + if u.op in (Ops.PARAM, Ops.SPECIAL) or u.op is Ops.CAST and u.src[0].op is Ops.CONST: return None + start = 1 if u.op in (Ops.WHERE, Ops.STORE) else 0 + demand: DType|None + if u.op is Ops.STORE: demand = u.src[0].dtype + elif u.op in GroupOp.Comparison: + demand = least_upper_dtype(*([s.dtype for s in u.src if s.dtype not in dtypes.weaks] or [select_dtype(s) for s in u.src])) + elif u.dtype in dtypes.weaks: demand = least_upper_dtype(*(s.dtype for s in u.src[start:])) + elif u.op in GroupOp.Broadcastable|{Ops.CAST}: demand = u.dtype + else: demand = None + src = u.src[:start]+tuple(graph_rewrite(s.cast(demand or select_dtype(s)), pm_lower_weak) + if s.dtype in dtypes.weaks else s for s in u.src[start:]) + return None if demand in dtypes.weaks or src == u.src else u.replace(dtype=None if u.dtype in dtypes.weaks else u.dtype, src=src) + +pm_commit_weak = PatternMatcher([(UPat(GroupOp.Broadcastable|{Ops.STORE, Ops.BUFFER}, name="u"), lower_weak_srcs)]) + +def lower_weak(c:UOp, u:UOp) -> UOp|None: + if c.dtype in dtypes.weaks: return None + if u.op is Ops.CAST: return u.src[0].cast(c.dtype if weak_dtype(c.dtype) is u.dtype else select_dtype(u)).cast(c.dtype) + if u.op is Ops.PARAM: return u.replace(dtype=None,arg=replace(u.arg,dtype=select_dtype(u))).cast(c.dtype) if u.addrspace == AddrSpace.ALU else None + dt = select_dtype(u) if weak_dtype(c.dtype) is not u.dtype or u.op is Ops.RANGE else c.dtype + return u.replace(dtype=None, src=tuple(s.cast(dt) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype) pm_lower_weak = PatternMatcher([ - (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)), - # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default. - # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs) - (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"), - lambda u,x: x.cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None), - # Binary can widen from the bounds, all other nodes derive from the lowered sources. - # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition - (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node), - (UPat(Ops.PARAM, dtype=dtypes.weakint, name="u"), - lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=select_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None), -]) -def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: - if ctx is None: ctx = {} - def lower(s:UOp) -> UOp: - if (r:=ctx.get(s)) is None: - r = graph_rewrite(s, pm_lower_weak) - # the consumer absorbs the cast on its own edge - ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype in dtypes.weaks else r - return r - # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands - ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src)) - return None if ret is u else ret - -def commit_weak(s:UOp, dt:DType) -> UOp: - # a bare weak CONST commits directly (its number must fit), a weak non-const src takes the demand cast - return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt) - -def commit_weak_srcs(u:UOp) -> UOp|None: - if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None - # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too - return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)) - -# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer -pm_commit_weak = PatternMatcher([ - (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs), - # demand from the destination: a STORE's weak value commits at the destination's dtype - (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"), - lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))), + (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.Broadcastable|GroupOp.Unary|{Ops.CAST, Ops.RANGE, Ops.STACK, Ops.PARAM}, + dtype=dtypes.weaks, name="u"),)), lower_weak), ]) -# push cast to weak src -pm_cast_weak = PatternMatcher([ - (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.Broadcastable, dtype=dtypes.weaks, name="u"),)), - lambda c,u: u.replace(dtype=None, src=tuple(commit_weak(s, c.dtype) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype) - if c.dtype not in dtypes.weaks else None), -]) +pm_address_demand = PatternMatcher([(UPat((Ops.INDEX, Ops.SHRINK), name="u"), lower_weak_srcs)]) # ADDRESS commits index operands +pm_boundary_demand = PatternMatcher([ + (UPat((Ops.RANGE, Ops.SPECIAL), dtype=dtypes.weaks, name="u"), lambda u: u.replace(dtype=(dt:=select_dtype(u)), + src=(u.src[0].cast(dt),)+u.src[1:] if u.op is Ops.RANGE else u.src)), # ROLE: RANGE size is machine input, SPECIAL size is metadata + (UPat(Ops.PARAM, dtype=dtypes.weaks, name="u"), + lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=select_dtype(u))) if u.addrspace == AddrSpace.ALU else None), + (UPat(GroupOp.All, name="u"), lambda u: lower_weak_srcs(u) if u.dtype not in dtypes.weaks else None)]) -pm_lower_index_dtype = pm_commit_weak+PatternMatcher([ - (UPat(GroupOp.All, name="u"), - lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None), +pm_lower_index_dtype = pm_commit_weak+pm_lower_weak+PatternMatcher([ # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded) # TODO: more generic (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))), diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index c8467ae654dd3..2f1120198fb2a 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -99,12 +99,11 @@ def render_marg(ctx,x:UOp): # TODO: movement ops simplify stuff, this can break SPEC=2 #(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"), # NOTE: CMPNE doesn't work cause there's no __rne__ - # explicit trunc ops: `//` and `%` parse as FLOORDIV/FLOORMOD, so render CDIV/CMOD via .alu() - (UPat(Ops.CDIV, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CDIV, {ctx[x.src[1]]})"), - (UPat(Ops.CMOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CMOD, {ctx[x.src[1]]})"), + # explicit trunc ops and shifts: Python syntax changes their op or re-promotes operands, so render them via .alu() + (UPat((Ops.CDIV, Ops.CMOD, Ops.SHL, Ops.SHR), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.{x.op.name}, {ctx[x.src[1]]})"), # `.where` re-promotes its operands, so render WHERE via .alu() too (UPat(Ops.WHERE, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.WHERE, {ctx[x.src[1]]}, {ctx[x.src[2]]})"), - (UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, name="x"), lambda ctx,x: + (UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD, Ops.SHL, Ops.SHR}, name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")), (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"), (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \ diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index a6af821b61e0e..723e0979ebca6 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -45,6 +45,9 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): # ***** new specs ***** def matches_dtype(x:UOp, dtype:DType) -> bool: return x.dtype == dtype or x.base.is_invalid # Invalid matches any dtype +def valid_special(s:UOp, x:UOp) -> bool: + # weakint is widthless launch metadata; a concrete source must agree with SPECIAL's dtype (Invalid agrees with every dtype) + return s.dtype in dtypes.ints+(dtypes.weakint,) and (x.dtype is dtypes.weakint or matches_dtype(x, s.dtype)) and isinstance(s.arg, str) # these ops can be used in the tensor graph and programs spec_shared = PatternMatcher([ # NOTE: for testing, we let sinks be anything @@ -158,7 +161,7 @@ def valid_gettuple(g:UOp, t:UOp): return isinstance(g.arg, int) and 0 <= g.arg < (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), valid_gettuple), # SPECIAL is index before index lowering. custom_kernel currently has this - (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.weakint),), name="s"), lambda s,x: matches_dtype(x, s.dtype) and isinstance(s.arg, str)), + (UPat(Ops.SPECIAL, src=(UPat.var("x"),), name="s"), valid_special), # movement ops (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat())), lambda: True), @@ -201,11 +204,12 @@ def valid_gettuple(g:UOp, t:UOp): return isinstance(g.arg, int) and 0 <= g.arg < # these ops can exist in programs but not the tensor spec. example: LOAD spec_program = PatternMatcher([ - # index and weak dtypes are not allowed in programs - (UPat(GroupOp.All, (dtypes.weakint, dtypes.weakfloat)), lambda: False), + # weak values are bare constants consumed only by CAST, except PARAM shapes and SPECIAL sizes + (UPat(GroupOp.All-{Ops.CONST}, dtypes.weaks), lambda: False), + (UPat(GroupOp.All-{Ops.CAST, Ops.PARAM, Ops.SPECIAL}, name="u"), lambda u: False if any(s.dtype in dtypes.weaks for s in u.src) else None), # allow special SHRINK - (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST))), lambda: True), + (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST).or_casted())), lambda: True), # movement ops are not allowed in programs (UPat(GroupOp.Movement), lambda: False), @@ -221,7 +225,7 @@ def valid_gettuple(g:UOp, t:UOp): return isinstance(g.arg, int) and 0 <= g.arg < (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True), # SPECIAL is int32 after index lowering - (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.int32),), name="s"), lambda s,x: matches_dtype(x, s.dtype) and isinstance(s.arg, str)), + (UPat(Ops.SPECIAL, src=(UPat.var("x"),), name="s"), valid_special), ])+spec_shared spec_hcq = PatternMatcher([ diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 76c600f90d872..6b8dca64f5926 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -1,8 +1,8 @@ # all of symbolic lives here now import math, struct from collections import defaultdict -from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu -from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid +from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, cvar, exec_alu +from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, least_upper_dtype, weak_dtype, Invalid from tinygrad.helpers import partition, all_same, prod, flatten, unwrap, IMAGE, dedup from tinygrad.uop.divandmod import div_and_mod_symbolic from tinygrad.uop.movement import mop_cleanup @@ -26,6 +26,8 @@ def convert(v:ConstType) -> ConstType: return struct.unpack(to_fmt, struct.pack( return root.const_like(convert(c.val)) def const_arg(u:UOp) -> ConstType|tuple[ConstType, ...]|None: + # a committed constant is the pair CAST(dt, weak CONST): it denotes the same number as the bare const + if u.op is Ops.CAST and u.src[0].op is Ops.CONST and u.src[0].dtype in dtypes.weaks: u = u.src[0] if u.op is Ops.CONST: return u.val if u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src): return tuple(s.val for s in u.src) return None @@ -134,9 +136,9 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"), lambda x: x.const_like(False, dtypes.bool)), # x != x -> False (only ints) # ** constant folding ** - (UPat(GroupOp.Unary, src=(UPat((Ops.CONST, Ops.STACK)),), name="a"), fold_const_alu), + (UPat(GroupOp.Unary, src=(UPat((Ops.CONST, Ops.STACK)).or_casted(),), name="a"), fold_const_alu), # NOTE: THREEFRY(const,const) folds via its decomposition - (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=(UPat((Ops.CONST, Ops.STACK)),)*2, name="a"), fold_const_alu), + (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=(UPat((Ops.CONST, Ops.STACK)).or_casted(),)*2, name="a"), fold_const_alu), (UPat(GroupOp.Ternary, src=(UPat((Ops.CONST, Ops.STACK)),)*3, name="a"), fold_const_alu), # 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), @@ -151,11 +153,12 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: # x*0 -> 0 or 0*x -> 0 # if x is nan or inf it should render the nan value. # NOTE: this can be wrong for loaded NaN - (UPat.var("x") * 0, lambda x: x.const_like(float("nan") if x.op is Ops.CONST - and isinstance(x.val, float) and (math.isnan(x.val) or math.isinf(x.val)) else 0)), + (UPat.var("x") * 0, lambda x: x.const_like(float("nan") if + (c:=x.src[0] if x.op is Ops.CAST and x.src[0].op is Ops.CONST else x).op is Ops.CONST and + isinstance(c.val, float) and not math.isfinite(c.val) else 0)), # *** cast/bitcast *** - # TODO: delete this once CONST has no dtype - (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val)), + (UPat(Ops.CAST, name="root", src=(UPat.cvar("c").or_casted(),)), lambda root,c: root.const_like(c.val) if + weak_dtype(root.dtype) != weak_dtype(c.dtype) or root.dtype not in dtypes.weaks and c.dtype not in dtypes.weaks else None), (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None), (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast), # b.cast(a).cast(b) -> b if a preserves all values in b @@ -167,9 +170,9 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: (UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow), # positive const ** x (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.val == 1 else (x*math.log2(c.val)).exp2() if c.val > 0 else None), - # unpack a uint64 packed from two uint32 (threefry) - (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), - (((UPat.var('x', dtypes.uint32).cast(dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, + # unpack a uint64 packed from two uint32 (threefry). cvar reads the committed pair below the boundary + (((UPat.var(None, dtypes.uint64)<>cvar(arg=32), lambda x: x.cast(dtypes.uint64)), # ** simple where folding ** # a conditional with the same results either way is a noop, also fold const conditionals @@ -216,7 +219,19 @@ def fold_where_closure(cond:UOp, t:UOp, f:UOp) -> UOp|None: if any(u.op_in_backward_slice_with_self(Ops.INDEX) for u in (cond, t, f)): return None return cond.where(t.substitute({cond: cond.const_like(True)}), f.substitute({cond: cond.const_like(False)})) -symbolic = symbolic_simple+commutative+PatternMatcher([ +def relax_committed_const(u:UOp) -> UOp|None: + carriers = [s.dtype for s in u.src if s.op is not Ops.CONST and not (s.op is Ops.CAST and s.src[0].op is Ops.CONST)] + dt = least_upper_dtype(*carriers) if carriers else u.dtype + def relax(s:UOp) -> UOp: + if s.op is not Ops.CAST or (c:=s.src[0]).op is not Ops.CONST or c.dtype not in dtypes.weaks: return s + if isinstance(c.val, float) and not math.isfinite(c.val): return c + return c if s.dtype == dt and least_upper_dtype(dt, c.dtype) == dt else s + if (src:=tuple(relax(s) for s in u.src)) == u.src: return None + return ret.cast(u.dtype) if (ret:=u.replace(dtype=None, src=src)).dtype in dtypes.weaks and not weak_dtype(u.dtype) else ret + +pm_relax_committed_const = PatternMatcher([(UPat(GroupOp.ALU-{Ops.TRUNC}, name="u"), relax_committed_const)]) + +symbolic = symbolic_simple+commutative+pm_relax_committed_const+PatternMatcher([ # ** boolean algebra ** # TODO: make a more general or folder like simplify_valid (UPat.var("x", dtype=dtypes.bool) | UPat.var("x", dtype=dtypes.bool).logical_not(), lambda x: x.const_like(True)), # x|!x -> True