From 7175a34d60d11713591cfb5c32959be25dc3f48d Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Wed, 19 Aug 2026 11:59:45 -0400 Subject: [PATCH] consts are weak 2 [pr] --- test/mockgpu/amd/emu.py | 2 +- test/mockgpu/amd/pcode.py | 41 +++++----- test/null/test_const_folding.py | 7 +- test/null/test_simplify_valid_idx.py | 4 +- test/null/test_uop_symbolic.py | 5 +- test/null/test_uops.py | 22 ++++-- test/unit/test_dtype_weak.py | 15 ++-- tinygrad/codegen/__init__.py | 25 +++---- tinygrad/codegen/decomp/dtype.py | 10 ++- tinygrad/codegen/decomp/op.py | 5 +- tinygrad/engine/realize.py | 2 +- tinygrad/renderer/cstyle.py | 5 +- tinygrad/renderer/isa/x86.py | 7 +- tinygrad/renderer/nir.py | 2 +- tinygrad/renderer/wgsl.py | 3 +- tinygrad/runtime/support/hcq2.py | 9 ++- tinygrad/schedule/rangeify.py | 4 +- tinygrad/uop/ops.py | 9 ++- tinygrad/uop/render.py | 11 +-- tinygrad/uop/spec.py | 10 +-- tinygrad/uop/symbolic.py | 54 ++++++++------ tinygrad/uop/weak.py | 107 +++++++++++++++++---------- 22 files changed, 207 insertions(+), 152 deletions(-) diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index 7c6fe9038d..6b66778006 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -709,7 +709,7 @@ class _Ctx: # VGPR bit-slice assignment: VGPR[lane][reg][hi:lo] = (vgpr_idx, rhs_val, hi, lo[, cond]) -> read-modify-write if dest.startswith('VGPR[') and re.search(r'\[\d+:\d+\]', dest): # VGPR bit-slice: (vgpr_idx, rhs_val, hi_bit, lo_bit) - hi/lo are UOp constants - hi_bit, lo_bit = int(val[2].val), int(val[3].val) + hi_bit, lo_bit = int(val[2].src[0].val), int(val[3].src[0].val) width = hi_bit - lo_bit + 1 old = self.vgpr.index(val[0]).load() new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32) diff --git a/test/mockgpu/amd/pcode.py b/test/mockgpu/amd/pcode.py index bb66ffe03f..c6b22da9ef 100644 --- a/test/mockgpu/amd/pcode.py +++ b/test/mockgpu/amd/pcode.py @@ -8,6 +8,9 @@ from tinygrad.codegen.decomp.dtype import f2f VarVal = UOp | tuple[str, list[str], str] def _const(dt, v): return UOp.const(v, dt) +# parser literals are pairs; restating one re-mints it instead of stacking CASTs +def _lit(u:UOp): return u.src[0].val if u.op is Ops.CAST and u.src[0].op is Ops.CONST else None +def _restate(u:UOp, dt): return _const(dt, v) if (v:=_lit(u)) is not None else u.cast(dt) def _u32(v): return _const(dtypes.uint32, v) def _u64(v): return _const(dtypes.uint64, v) def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32) @@ -55,8 +58,7 @@ 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 isinstance(sv:=_lit(src), int) and sv > 0 and (sv & (sv + 1)) == 0: widths.append(sv.bit_length()) if widths: return max(widths) return v.dtype.bitsize @@ -163,7 +165,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 _lit(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) @@ -482,7 +484,7 @@ class Parser: def _apply_binop(self, left, right, op): if op in ('||', '&&', '|', '^', '&'): left, right = self._coerce_bitwise(left, right) elif op in ('>=', '<=', '>', '<', '==', '!=', '<>', '>>', '<<'): left, right = self._coerce_cmp(left, right) - elif left.dtype != right.dtype: right = right.cast(left.dtype) + elif left.dtype != right.dtype: right = _restate(right, left.dtype) match op: case '||' | '|': return left | right case '&&' | '&': return left & right @@ -497,7 +499,7 @@ class Parser: 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 (lv:=_lit(left)) is not None and (rv:=_lit(right)) is not None: return _const(left.dtype, lv - rv) 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 +509,7 @@ class Parser: 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 _lit(left) == 2.0 else left _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)] @@ -529,8 +531,8 @@ class Parser: 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 (v:=_lit(inner)) is not None: + return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -v) return inner.neg() if self.try_eat_val('+', 'OP'): return self.unary() return self.postfix() @@ -669,15 +671,15 @@ class Parser: self.eat('OP') width = self.parse() self.eat('RBRACKET') - if width.op == Ops.CONST: - w = int(width.val) + if (wv:=_lit(width)) is not None: + w = int(wv) 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 (fv:=_lit(first)) is not None and (sv:=_lit(second)) is not None: + a, b = int(fv), int(sv) if a < b: return _bitreverse(base, b - a + 1) hi, lo = a, b if lo >= base.dtype.itemsize * 8: @@ -698,8 +700,8 @@ class Parser: 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 (fv:=_lit(first)) is not None: + idx = int(fv) # Check for array element (var@idx) if var_name and f'{var_name}@{idx}' in self.vars: v = self.vars[f'{var_name}@{idx}'] @@ -758,7 +760,7 @@ class Parser: if type_char == 'F' and inner.dtype in (dtypes.uint32, dtypes.uint64, dtypes.ulong, dtypes.int, dtypes.int64): if inner.dtype.itemsize != dt.itemsize: inner = inner.cast(dtypes.uint32 if dt.itemsize == 4 else dtypes.uint64) return inner.bitcast(dt) - return inner.cast(dt) + return _restate(inner, dt) if self.at('IDENT'): ident = self.peek().val fmt = ident[0].lower() @@ -872,8 +874,8 @@ class Parser: 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) - else: r = r.cast(l.dtype) + if r.dtype == dtypes.int and (rv:=_lit(r)) is not None and rv < 0: l = l.cast(dtypes.int) + else: r = _restate(r, l.dtype) return l, r def _coerce_bitwise(self, l: UOp, r: UOp) -> tuple[UOp, UOp]: @@ -969,8 +971,9 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic 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) + v = _lit(expr) + assert v is not None, f"loop bound must be constant, got {expr}" + return int(v) start_val = parse_bound() p.eat('COLON') end_val = parse_bound() diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index e173114a57..e7072dae05 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -3,7 +3,6 @@ from tinygrad import dtypes, Context from tinygrad.dtype import DType, ConstType from tinygrad.uop.ops import Ops, UOp from test.helpers import full_rewrite -import numpy as np class TestWeakConstFolding(unittest.TestCase): def test_weakint_math(self): @@ -27,16 +26,14 @@ class TestBitcastConstFolding(unittest.TestCase): for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000), (70000, dtypes.int16, dtypes.uint16, 4464), (-5, dtypes.uint32, dtypes.int32, -5)): - self.assertEqual(UOp.const(val, src_dt).bitcast(dst_dt).simplify().val, bits) + self.assertIs(UOp.const(val, src_dt).bitcast(dst_dt).simplify(), UOp.const(bits, dst_dt)) def test_scalar_bitcast(self): def t(cases: dict[DType, ConstType]): for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()): if not math.isnan(from_v): r = UOp.const(from_v, from_dt).bitcast(to_dt).simplify() - self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})") - self.assertEqual(r.dtype, to_dt, msg) - np.testing.assert_equal(r.val, to_v, msg) + self.assertIs(r, UOp.const(to_v, to_dt), f"{from_dt} -> {to_dt} ({from_v} -> {to_v})") t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False}) t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True}) diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 2d278eb75f..2166e69e8a 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -3,7 +3,7 @@ import unittest, itertools from tinygrad.codegen.late.coalesce import indexing_simplify from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp, Ops, graph_rewrite -from tinygrad.uop.weak import pm_lower_index_dtype +from tinygrad.uop.weak import pm_commit_weak from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load from tinygrad.helpers import Context from test.helpers import full_rewrite @@ -496,7 +496,7 @@ class TestImageSimplification(unittest.TestCase): 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(load.sink(), pm_commit_weak+indexing_simplify).src[0].src[0] self.assertEqual(off.src[1].get_valid(), UOp.const(True)) class TestDropTrueGate(unittest.TestCase): diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index db3bd00ed8..079422c1d4 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -6,7 +6,6 @@ from tinygrad.dtype import dtypes, ConstType, DType, Invalid from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer from tinygrad.uop.spec import spec_shared, type_verify from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load -from tinygrad.uop.weak import pm_cast_weak from tinygrad.uop.validate import uops_to_z3 def check_uop_against_string(self, v:UOp, s:str): @@ -36,7 +35,7 @@ class TestSymbolic(unittest.TestCase): self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original") def helper_test_variable(self, v, n, m, s, test_z3:bool=True): - v_simplified = graph_rewrite(v, sym+pm_cast_weak, name="simplify symbolic uop") + v_simplified = graph_rewrite(v, sym, name="simplify symbolic uop") if test_z3: self.check_equal_z3(v, v_simplified) nmin, nmax = v_simplified.vmin, v_simplified.vmax check_uop_against_string(self, v_simplified, s) @@ -1017,7 +1016,7 @@ class TestSymbolic(unittest.TestCase): cond = Variable("s", 0, 3, dtypes.int) < 2 a = Variable("a", 0, 3, dtypes.int) self.assertIs(graph_rewrite(cond.where(a, a+1).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), (a+1).cast(dtypes.half))) - self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half))) + self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), uconst(2.0))) self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid())) def test_where_const_gate_keeps_stated_width(self): diff --git a/test/null/test_uops.py b/test/null/test_uops.py index de18f1a931..0d5cdfe066 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -5,8 +5,8 @@ from tinygrad.tensor import Tensor from tinygrad.helpers import Timing, Context, cdiv from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid # noqa: F401 from tinygrad.device import Device -from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests -from tinygrad.uop.weak import pm_lower_index_dtype +from tinygrad.uop.ops import Ops, AxisType, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests +from tinygrad.uop.weak import pm_lower_weak from tinygrad.uop.spec import spec_program, spec_shared, type_verify from tinygrad.uop.symbolic import sym, pm_remove_invalid from test.helpers import eval_uop, to_uops_list @@ -76,16 +76,18 @@ class TestLowerIndexDtype(unittest.TestCase): buf = UOp.param(0, dtypes.float, (2**31+64,)) i = UOp.variable("i", 0, 2**28) shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4))) - lowered = graph_rewrite(shrink.sink(), pm_lower_index_dtype) - self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint") + lowered = graph_rewrite(shrink.sink(), pm_lower_weak) + self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks), + "lowering must resolve every weak width, except a typed literal's value half") sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK) self.assertEqual(sh.src[1].dtype, dtypes.long) def test_reg_buffer_size_lowers(self): reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG) self.assertEqual(reg.src[0].dtype, dtypes.weakint) - lowered = graph_rewrite(reg.sink(), pm_lower_index_dtype) - self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint") + lowered = graph_rewrite(reg.sink(), pm_lower_weak) + self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks), + "lowering must resolve every weak width, except a typed literal's value half") self.assertEqual(next(u for u in lowered.backward_slice_with_self if u.op is Ops.BUFFER).src[0].dtype, dtypes.int) class TestSafeCast(unittest.TestCase): @@ -457,6 +459,14 @@ class TestUopsObject(unittest.TestCase): self.assertEqual(a.device, Device.DEFAULT) class TestUOpRender(unittest.TestCase): + def test_render_ssimplified_marg_outside_toposort(self): + r = UOp.range(UOp.const(16, dtypes.int), 2, AxisType.WEAK, dtype=dtypes.int) + offset = UOp(Ops.SHL, src=(r, UOp.const(1, dtypes.int))) + shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, (32,)), offset, UOp.const(2, dtypes.int))) + self.assertIsNot(shrink.src[1], shrink.marg[0][0]) + self.assertEqual(shrink.render(simplify=False), "p0.shrink((((r2<<1), 2),))") + self.assertEqual(UOp.range(1, 0, src=(shrink,), dtype=dtypes.int).render(simplify=False), "r0") + def test_render_vectorize_empty(self): u = UOp(Ops.STACK, dtype=dtypes.void, src=()) self.assertEqual(u.render(simplify=False), "{}") diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 59cc8f40f3..fcd8e5fe84 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -4,7 +4,7 @@ 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, GroupOp, dtype_from_uop, graph_rewrite -from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak +from tinygrad.uop.weak import pm_commit_weak from tinygrad.uop.symbolic import symbolic_simple from tinygrad.uop.spec import spec_shared, type_verify from tinygrad.engine.jit import JitError @@ -74,7 +74,7 @@ class TestWeakPromotion(unittest.TestCase): recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL] self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)]) with Context(DEFAULT_FLOAT=dtypes.float16): - committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={}) + committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_commit_weak) self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32]) def test_div_sub_operand_kept_weak(self): @@ -85,7 +85,7 @@ class TestWeakPromotion(unittest.TestCase): def test_cast_weak_expression_commits_at_cast_floor(self): # the floor never narrows: a cast BELOW the default does not pull the compute width down with it with Context(DEFAULT_FLOAT=dtypes.float32): - narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_lower_index_dtype, ctx={}) + narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_commit_weak) self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32)) def test_cast_weak_expression_value_uses_cast_floor(self): @@ -125,16 +125,17 @@ class TestWeakPromotion(unittest.TestCase): with Context(DEFAULT_FLOAT=dtypes.float16): 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={}) + out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_commit_weak) # 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)) def test_weak_srcs_commit_only_at_a_concrete_lub(self): weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0))) - self.assertIs(graph_rewrite(weak_lub, pm_lower_index_dtype, ctx={}), weak_lub) + self.assertIs(graph_rewrite(weak_lub, pm_commit_weak), weak_lub) concrete = UOp.const(2.0).cast(dtypes.float16) - where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_lower_index_dtype, ctx={}) - self.assertEqual(tuple(x.dtype for x in where.src), (dtypes.bool, dtypes.float16, dtypes.float16)) + # the weak arm stays bare: its sibling states the width, so the WHERE already derives float16 for it + where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_commit_weak) + self.assertEqual((where.dtype, tuple(x.dtype for x in where.src)), (dtypes.float16, (dtypes.bool, dtypes.float16, dtypes.weakfloat))) 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`) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 277430f1fe..0a50cbd9a9 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -3,7 +3,7 @@ import itertools, functools 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, NUM_CPU_THREADS, TC_SELECT, TC_OPT, TracingKey, Context, panic from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType -from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak, pm_cast_weak +from tinygrad.uop.weak import pm_lower_weak, pm_commit_weak, pm_cast_const from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program from tinygrad.renderer import Renderer, Estimates @@ -281,10 +281,6 @@ pm_implicit_barriers = PatternMatcher([ (UPat(Ops.END, name="end"), add_war_barrier), ]) -pm_casted_consts = PatternMatcher([ - (UPat(Ops.CONST, dtypes.all, name="c"), lambda c: UOp.cconst(c.val, c.dtype)), -]) - def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST") if DEBUG >= 5: print(pyrender(ast)) @@ -346,11 +342,13 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # extra symbolic before decomp. crashes without this? # NOTE: also run indexing_simplify here, while the index is still weakint and (x+y)*c -> x*c+y*c applies - sink = graph_rewrite(sink, sym+indexing_simplify, name="extra symbolic") + # commit widths minted in this fixpoint before optimization inspects RANGE/INDEX shapes + sink = graph_rewrite(sink, sym+indexing_simplify+pm_commit_weak, name="extra symbolic") - # lower index dtype + # THE BOUNDARY: required compute widths settle here; derivable literal edges may stay bare # NOTE: we need indexing_simplify to remove the cast to long using the Invalid - sink = graph_rewrite(sink, symbolic_simple+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes") + # NOTE: symbolic must NOT be composed here -- its cast collapse eats the weak CAST a lowered node wears, and it cycles + sink = graph_rewrite(sink, pm_lower_weak+indexing_simplify, name="lower all index dtypes") # final symbolic before decomp sink = graph_rewrite(sink, symbolic, name="final symbolic") @@ -374,11 +372,12 @@ 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_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends + pm_final_rewrite = pm_commit_weak+pm_decomp+extra_matcher+pm_split_ends sink = graph_rewrite(sink, pm_final_rewrite+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") + # wrap every const still riding bare so no renderer reads one, and add implicit barriers (stores/loads through LOCAL + # memory ordered by AFTER or across loop iterations need workgroup barriers). neither matcher undoes the other + sink = graph_rewrite(sink, pm_cast_const+pm_implicit_barriers, name="state literal widths, add implicit barriers") # this was the linearizer sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True) @@ -387,10 +386,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: 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) - # spell every literal as a casted const CAST(dt, CONST(value)) - # TODO: remove once consts are always weak - sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", 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 db85526c83..ef4536aab9 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -3,6 +3,7 @@ from tinygrad.dtype import dtypes, DType, truncate from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES, Context, SPEC from tinygrad.uop import GroupOp from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite, ParamArg +from tinygrad.uop.weak import commit_weak_sibling from tinygrad.renderer import Renderer from tinygrad.codegen.decomp.transcendental import exponent_bias, shl, shr @@ -136,6 +137,8 @@ def f2f_store(st, idx, val, fr:DType, to:DType): # tag is the 32-bit word this node becomes - (0 for the low word, 1 for the high, the dtype the consumer wants) pm_long_decomp = PatternMatcher([ + # word splitting needs a bare literal committed at its long sibling's width + (UPat(GroupOp.All, name='x'), lambda x: commit_weak_sibling(x, next((s.dtype for s in x.src if s.dtype in l2i_dt), None))), (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: @@ -147,6 +150,9 @@ pm_long_decomp = PatternMatcher([ split_l2i(ctx, x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))), (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x: split_l2i(ctx, Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]), + # a literal splits by value; the general CAST arm below would drop its high word + (UPat(Ops.CAST, src=(UPat(Ops.CONST, name='c'),), tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), + lambda x,c: UOp.const(truncate[x.tag[1]](c.val >> (32*x.tag[0])), x.tag[1])), (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda ctx,a,x: split_l2i(ctx, x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None), (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x: @@ -161,12 +167,12 @@ pm_long_decomp = PatternMatcher([ 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])) ]) # float decomposition patterns - ctx is (fr, to) tuple pm_float_decomp = PatternMatcher([ + # emulation needs a bare literal committed at its emulated sibling's width + (UPat(GroupOp.All, name='x'), lambda ctx,x: commit_weak_sibling(x, next((s.dtype for s in x.src if s.dtype == ctx[0]), None))), (UPat((*GroupOp.Defines, Ops.INDEX, Ops.SHRINK), name="x"), lambda ctx,x: x.replace(dtype=f2f_dt[ctx[0]], arg=replace(x.arg, dtype=f2f_dt[ctx[0]]) if isinstance(x.arg, ParamArg) else x.arg, tag=ctx[0]) if x.dtype == ctx[0] and (x.op is not Ops.INDEX or x.src[0].op not in {Ops.LOAD, Ops.STACK}) else None), diff --git a/tinygrad/codegen/decomp/op.py b/tinygrad/codegen/decomp/op.py index 6a48cdca53..ddb353c398 100644 --- a/tinygrad/codegen/decomp/op.py +++ b/tinygrad/codegen/decomp/op.py @@ -128,6 +128,7 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa 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))] - pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))] + # mint the 1 bare: the contraction rule below matches it as a bare CONST, and FDIV re-derives the concrete width + pat += [(UPat.var("x").reciprocal(), lambda x: UOp.const(1.0).alu(Ops.FDIV, x))] + pat += [(UPat.var("a") * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))] return PatternMatcher(pat) diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index fb3284b1ee..b90532c100 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -202,7 +202,7 @@ def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]: def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]: bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)] - shape, pos_var = tuple(s.val for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr + shape, pos_var = tuple(s.src[0].val for s in ast.src if s.op is Ops.CAST and s.src[0].op is Ops.CONST), ast.variables()[0].expr bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var]) return [] diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 8cd584f63e..398dad5b1e 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -3,6 +3,7 @@ import math, sys, struct from collections import defaultdict, Counter from tinygrad.codegen.opt import tc from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters +from tinygrad.uop.weak import commit_weak_sibling from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, NUM_CPU_THREADS, IMAGE, FLOAT16, is_image_shape from tinygrad.dtype import dtypes, DType, AddrSpace, truncate, float_to_bf16 from tinygrad.renderer import Renderer @@ -74,6 +75,8 @@ base_rewrite = PatternMatcher([ def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True): patterns = PatternMatcher([ + # a weak CONST states no width and cannot be restated: commit it at the emulated dtype a sibling src states + (UPat(GroupOp.ALU, name="x"), lambda x, dts=dts: commit_weak_sibling(x, next((s.dtype for s in x.src if s.dtype in dts), None))), (UPat(Ops.WHERE, dtype=dts, src=(UPat.var("b"), UPat.var("x"), UPat.var("y")), name="w"), lambda w,b,x,y: b.where(x.cast(dtypes.float), y.cast(dtypes.float)).cast(w.dtype)), (UPat(GroupOp.ALU-{Ops.WHERE}, dtype=dts, name="x"), @@ -520,8 +523,6 @@ class HIPRenderer(CStyleLanguage): (UPat(Ops.WMMA, name="x", dtype=dtypes.float), lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2])) if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None), - # bfloat16 constant casting - (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.val, dtypes.float))), ]) def asm(self, prg:UOp, lin:UOp) -> bytes: diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index e4d2e7471b..fe86869235 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -150,9 +150,10 @@ extra_matcher = PatternMatcher([ # no cmpne for packed ints, y != x => !(y==x) (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"), lambda y,x,cmp: UOp(Ops.CMPEQ, src=(y,x))^True if y.max_numel() > 1 else None), - # float where expects a mask - (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats), UPat.var("b")), - lambda m,a,b: m.cast(a.dtype).ne(0).where(a, b) if m.src[0].dtype not in dtypes.floats else None), + # float WHERE needs a mask unless its comparison already has a float operand + (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats+(dtypes.weakfloat,)), UPat.var("b")).named("w"), + lambda m,a,b,w: m.cast(a.dtype if a.dtype in dtypes.floats else w.dtype).ne(0).where(a, b) + if w.dtype in dtypes.floats and m.src[0].dtype not in dtypes.floats+(dtypes.weakfloat,) else None), # rewrite -x -> 0 - x (UPat(Ops.NEG, name="x"), lambda x: UOp(Ops.SUB, src=(x.const_like(0),) + x.src)), # TODO: add support for mod, requires support for accessing the 2nd+ reg of a multi output instruction diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 3b254ef238..80e34368ec 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -138,7 +138,7 @@ class NIRRenderer(Renderer): # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace( src=(buf,UOp.const(off.val, dtypes.long) if off.op is Ops.CONST else off.cast(dtypes.long))+x.src[2:]) - if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None), + if buf.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) and not is_image_shape(buf._shape) else None), # images need index to be int for nir (coordinates only: the INDEX keeps its access dtype) (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x"), name="x"), lambda x,buf,idx_y,idx_x: x.replace(src=(buf, idx_y.cast(dtypes.int), idx_x.cast(dtypes.int)))), diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index b618db4ddc..625b3ebe1e 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -72,7 +72,8 @@ class WGSLRenderer(CStyleLanguage): (UPat.cvar("c").cast(dtypes.bool), lambda c: "true" if c.val else "false"), (UPat.cvar("c").cast((dtypes.uchar, dtypes.ushort, dtypes.uint32)), lambda c: f"bitcast({c.val})" if c.val < 0 else f"{c.val&0xFFFFFFFF}u"), - (UPat.cvar("c").cast(dtypes.int32, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}"), + # a negative literal must state its type: contextual conversion of a bare abstract int rejects it in a u32 position + (UPat.cvar("c").cast(dtypes.int32, name="x"), lambda ctx,x,c: f"i32({v})" if (v:=truncate[x.dtype](c.val)) < 0 else f"{v}"), (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),),)), diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index ef4ff904b3..88bac3bbb1 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -67,8 +67,8 @@ def make_binary_patch(buf:UOp, blob:bytes) -> UOp: def make_cmdbuf(lin, devs, buf:UOp|None=None): blob, patches = bytearray(), [] for s in (s for ins in lin.src for s in ins.src): - if s.op is not Ops.CONST: patches.append((len(blob), s)) - blob.extend(struct.pack(f'<{s.dtype.fmt}', s.val if s.op is Ops.CONST else 0x0)) + if not (lit:=(s.op is Ops.CAST and s.src[0].op is Ops.CONST)): patches.append((len(blob), s)) + blob.extend(struct.pack(f'<{s.dtype.fmt}', s.src[0].val if lit else 0x0)) cmdbuf = buf if buf is not None else UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf") return cmdbuf.after(make_binary_patch(cmdbuf, bytes(blob)), *make_patches(cmdbuf, patches)) @@ -328,7 +328,8 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp return table, reads, fills, {g:slots[bare[g]] for g in gaddrs} def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patches:list[UOp]) -> dict[UOp, UOp]: - (dst,), words = dedup(p.buf_uop for p in patches), [(off.val, slots[val]) for p in patches for off, val in zip(p.src[0].src[1].src, p.src[1].src)] + (dst,) = dedup(p.buf_uop for p in patches) + words = [(off.src[0].val, slots[val]) for p in patches for off, val in zip(p.src[0].src[1].src, p.src[1].src)] # build a runtime loop that writes every input address pairs = UOp.placeholder((2*len(words),), dtypes.uint32, next(UOp.unique_num), device=dst.device).rtag("systems") @@ -511,7 +512,7 @@ def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp: for off,val in zip(off.src, val.src): for b,v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)): data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype]((v.src[0] if v.op is Ops.CAST else v).val)) - b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(bo:=off.val*buf.dtype.itemsize):bo+len(data)] = data + b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(bo:=off.src[0].val*buf.dtype.itemsize):bo+len(data)] = data return UOp(Ops.NOOP) def resolve_getaddr(buf:UOp, g:UOp) -> UOp: diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 26ed0d880d..0c94435099 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -304,10 +304,12 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([ (UPat(Ops.INDEX, name="idx").f(Ops.STAGE, allow_any_len=True, name="b2"), remove_noop_bufferize), (UPat(Ops.INDEX, src=(UPat(Ops.STAGE),), allow_any_len=True, name="idx").f(Ops.NOOP).f(Ops.STAGE, allow_any_len=True, name="b2"), remove_noop_bufferize), - # no buffers for const (ranges don't matter for const - it's the same value everywhere) + # no buffers for const; STAGE can receive either legal above-boundary spelling (UPat(Ops.CONST, name='c').f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.val)), + (UPat.cvar('c').cast().f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.val)), # indexing a const is a const (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),),), lambda c: c), + (UPat(Ops.INDEX, src=(UPat.cvar().cast(name="c"),),), lambda c: c), # indexing an after with all fully invalid stores is invalid (UPat(Ops.INDEX, src=(UPat(Ops.AFTER, name="after"),), allow_any_len=True, name="idx"), lambda idx,after: idx.const_like(Invalid) if after_all_invalid(after) else None), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index a0c0e72d6e..39d24f3941 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -188,7 +188,8 @@ class UOpMetaClass(type): def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None, metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None): if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void - # CONST derives its dtype by value only when the constructor omits one + # re-mint concrete CAST literals on their grid; derived targets keep the CAST as the conversion + if op is Ops.CAST and src[0].op is Ops.CONST and (c:=UOp(Ops.CONST, arg=dtype.const(src[0].arg))).dtype is not dtype: src = (c,) # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py # an INDEX presents its access dtype, which a still-weak source matches up to weakness if SPEC == 2 and op is not Ops.CONST and \ @@ -614,7 +615,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if isinstance(b, UOp): return b.cast(dtype) # NOTE: it always has to be STACK now, even if they are all the same if isinstance(b, tuple): return UOp.stack(*[UOp.const(c, dtype) for c in b]) - return UOp(Ops.CONST, dtype, arg=dtype.const(b), src=()) + # .cast folds away at exactly the dtypes a CONST derives (bool/weakint/weakfloat): bare there, the pair everywhere else + return UOp(Ops.CONST, arg=dtype.const(b), src=()).cast(dtype) # weak CONST with width on the CAST. TODO: this is the final const @staticmethod def cconst(b:ConstLike, dtype:DType): return UOp(Ops.CAST, dtype, src=(UOp.const(b),), arg=dtype) @@ -990,7 +992,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return unwrap(self.arg.name) def bind(self, val:int|UOp): assert self.is_variable, f"op is {self.op}, need Variable" - uval = self.const_like(val) if isinstance(val, int) else val + # the Variable states the width, so the bound value stays BARE: is_bound_var tests for a CONST there, unbind reads .val + uval = UOp.const(val) if isinstance(val, int) else val assert self.vmin <= uval.vmin and uval.vmax <= self.vmax, f"bind {val} not in range [{self.vmin}, {self.vmax}]" assert uval.divides(self.arg.multiple_of) is not None, f"bind {val} not divisible by {self.arg.multiple_of}" return self.after(self.store(uval)) diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index b005a6eaf2..17e7a87d34 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -1,6 +1,6 @@ from tinygrad.dtype import AddrSpace, dtypes from tinygrad.uop import Ops, GroupOp -from tinygrad.uop.ops import ParamArg, UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort +from tinygrad.uop.ops import ParamArg, UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort, sint from tinygrad.helpers import strip_parens def pretty_print(x:UOp, cache=None, d=0)->str: @@ -69,14 +69,15 @@ renderer_infer = PatternMatcher([ # *** pyrender *** def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})" +# marg is ssimplify'd, so a bound can be a node this graph never contained -- render that one on its own +def marg_str(ctx, a:sint) -> str: return str(a) if not isinstance(a, UOp) else ctx[a] if a in ctx else a.render() + def render_marg(ctx,x:UOp): if x.op is Ops.PERMUTE: return str(x.marg) if x.op is Ops.FLIP: return str(tuple([i for i,x in enumerate(x.marg) if x])) pieces = [] - if x.op in {Ops.RESHAPE, Ops.EXPAND}: - pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg] - if x.op in {Ops.PAD, Ops.SHRINK}: - pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg] + if x.op in {Ops.RESHAPE, Ops.EXPAND}: pieces = [marg_str(ctx, a) for a in x.marg] + if x.op in {Ops.PAD, Ops.SHRINK}: pieces = [f"({marg_str(ctx, a[0])}, {marg_str(ctx, a[1])})" for a in x.marg] return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)" sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY, diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 36d72390d6..5597ab7191 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -203,10 +203,9 @@ spec_tensor = PatternMatcher([ # these ops can exist in programs but not the tensor spec. example: LOAD spec_program = PatternMatcher([ - # a literal is CAST(dt, CONST(value)), so its inner CONST is the one weak node a program may contain - (UPat(Ops.CONST, dtype=dtypes.weaks, name="x"), lambda x: x.dtype is dtypes.from_py(x.val)), - # index and weak dtypes are not allowed in programs - (UPat(GroupOp.All, (dtypes.weakint, dtypes.weakfloat)), lambda: False), + # every width in a program is stated: a CONST appears only under the CAST stating its width, and is the only weak node + (UPat(GroupOp.All, name="x"), lambda x: False if x.op is not Ops.CAST and any(s.op is Ops.CONST for s in x.src) else None), + (UPat(GroupOp.All, dtypes.weaks, name="x"), lambda x: None if x.op is Ops.CONST else False), # allow special SHRINK (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST).or_casted())), lambda: True), @@ -254,8 +253,9 @@ spec_kernel_graph = PatternMatcher([ (UPat(Ops.SINK, dtypes.void), lambda: True), # the store of a bound Variable binds it: AFTER(BUFFER, STORE(BUFFER, CONST)) in call args (UPat(Ops.STORE, dtypes.void, (UPat(Ops.BUFFER, name="b"), UPat(Ops.CONST))), lambda b: b.is_variable), - # const + stack to make vconsts and shape args + # const + stack to make vconsts and shape args. a 0-size/bound reduce keeps its literal in the pair spelling (UPat(Ops.CONST, src=()), lambda: True), + (UPat(Ops.CAST, src=(UPat(Ops.CONST, src=()),)), lambda: True), (UPat(Ops.STACK, name="s"), lambda s: all(x.op in (Ops.CONST, Ops.PARAM) or x.is_variable or x.is_bound_var for x in s.src) or None), # linear for more kernels (TODO: we should enter non sink calls) #(UPat(Ops.LINEAR), lambda: True), diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index d581f86c98..5ada537f35 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -1,12 +1,12 @@ # all of symbolic lives here now import math 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, bitcast, truncate +from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu, promo_dtype +from tinygrad.dtype import PyConst, dtypes, can_lossless_cast, Invalid, bitcast, truncate 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 -from tinygrad.uop.weak import commit_weak +from tinygrad.uop.weak import pm_uncast_const, commit_weak # TODO: symbolic shouldn't be importing from codegen from tinygrad.codegen.decomp.transcendental import xpow @@ -22,20 +22,11 @@ def simplify_pow(x:UOp, c:UOp) -> UOp|None: def fold_bitcast(root:UOp, c:UOp) -> UOp|None: if c.dtype.fmt is None or root.dtype.fmt is None or c.dtype.itemsize != root.dtype.itemsize: return None - return root.const_like(bitcast(truncate[c.dtype](c.val), c.dtype, root.dtype)) + # the value is MATHEMATICAL and may not fit: reading it as bits is the emission that pins it to the stated width + return root.const_like(bitcast(truncate[c.dtype](c.val if c.op is Ops.CONST else c.src[0].val), c.dtype, root.dtype)) -# const folding works for CONST, STACK, and casted CONST -const_folding_pat = UPat.any(UPat((Ops.CONST, Ops.STACK)), UPat(Ops.CAST, src=(UPat(Ops.CONST),))) - -def const_arg(u:UOp) -> ConstType|tuple[ConstType, ...]|None: - if u.op is Ops.CONST: return u.val - if u.op is Ops.CAST and u.src[0].op is Ops.CONST: return u.dtype.const(u.src[0].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 - -def fold_const_alu(a:UOp) -> UOp|None: - vals = [const_arg(s) for s in a.src] - return None if any(v is None for v in vals) else a.const_like(exec_alu(a.op, a.dtype, vals, False)) +# no truncate: ints stay mathematical past the fold (emission truncates); floats re-round in the mint +def fold_const_alu(a:UOp, vals) -> UOp: return a.const_like(exec_alu(a.op, a.dtype, vals, False)) def _quotient_base(q:UOp, base:UOp, div:int) -> UOp|None: # the B with q == B//div and B%div == base%div, or None. only such congruence is needed to recombine, and canonicalization @@ -71,6 +62,12 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: # this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0 invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i") invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat) + +# the two literal spellings and the value reader for each: Invalid carries no width, so it rides bare inside either +bare_lit = UPat.any(UPat(Ops.CONST), UPat(Ops.STACK, src=UPat(Ops.CONST))) +pair_lit = UPat.any(p:=UPat(Ops.CAST, src=(UPat(Ops.CONST),)), UPat(Ops.STACK, src=UPat.any(p, UPat(Ops.CONST, arg=Invalid)))) +def bare_arg(u:UOp): return tuple(s.val for s in u.src) if u.op is Ops.STACK else u.val +def pair_arg(u:UOp): return tuple(s.val if s.is_invalid else s.src[0].val for s in u.src) if u.op is Ops.STACK else u.src[0].val pm_data_invalid = PatternMatcher([ (invalid_pat.broadcast(), lambda i: i), (UPat(GroupOp.Unary|{Ops.CAST, Ops.BITCAST}, src=(invalid_pat,)), lambda i: i), @@ -141,11 +138,14 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ lambda x,mask,c: x // c.val if c.val > 0 and c.val & (c.val-1) == 0 and mask.val | (c.val-1) == -1 else 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=(const_folding_pat,), name="a"), fold_const_alu), + # One rule per spelling combination: bare has no width, pair evaluates at its stated width, and mixed first + # commits the bare operand to the promotion. `symbolic` below removes any newly redundant pair. # NOTE: THREEFRY(const,const) folds via its decomposition - (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=(const_folding_pat,)*2, name="a"), fold_const_alu), - (UPat(GroupOp.Ternary, src=(const_folding_pat,)*3, name="a"), fold_const_alu), + (UPat(GroupOp.ALU-{Ops.THREEFRY}, src=bare_lit, name="a"), lambda a: fold_const_alu(a, [bare_arg(s) for s in a.src])), + (UPat(GroupOp.ALU-{Ops.THREEFRY}, src=pair_lit, name="a"), lambda a: fold_const_alu(a, [pair_arg(s) for s in a.src])), + (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=[pair_lit, bare_lit], name="a"), lambda a: + a.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in a.src)) + if (dt:=promo_dtype(a.src)) not in dtypes.weaks else None), # bool MUL is AND, ADD/MAX is OR. prevents other rules to rewrite bool ADD/MUL incorrectly (UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), lambda x,y: x&y), (UPat.var('x', dtype=dtypes.bool) + UPat.var('y', dtype=dtypes.bool), lambda x,y: x|y), @@ -163,7 +163,9 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ and isinstance(x.val, float) and (math.isnan(x.val) or math.isinf(x.val)) else 0)), # *** cast/bitcast *** (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), + # a BITCAST reads its operand at the width it STATES, so a weak const is nonsense here: the bare arm is bool only + (UPat(Ops.BITCAST, name="root", src=(UPat.any(UPat(Ops.CONST, dtypes.bool, name="c"), UPat(Ops.CAST, src=(UPat(Ops.CONST),), name="c")),)), + fold_bitcast), # b.cast(a).cast(b) -> b if a preserves all values in b (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_lossless_cast(b.dtype, a.dtype) else None), # bitcast twice @@ -289,8 +291,9 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ # cast/long folding # if the intermediate cast doesnt narrow we can do it in one cast (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_lossless_cast(x.dtype, a.dtype) else None), + # commit_weak, not .cast: a weak b.dtype is not a literal spelling, and a CAST(weakfloat, CONST) reaches no commit round (UPat.var('x', dtypes.ints+(dtypes.weakint,)).cast(dtypes.ints+(dtypes.weakint,), name="a").cast(name="b"), - lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None), + lambda x,a,b: commit_weak(x, b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None), # try to do math in int instead of long, keep weak const weak (UPat(GroupOp.Binary, src=(UPat.var("x", (dtypes.long, dtypes.weakint)), UPat.var("y", (dtypes.long, dtypes.weakint))), name="u"), lambda u,x,y: (UOp.const(x.val) if x.op is Ops.CONST else x.cast(dtypes.int)).alu(u.op, @@ -303,7 +306,12 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ else y.src for y in x.src[1:]]))))), # after/end with 1 src is just src[0] (UPat((Ops.AFTER, Ops.END), src=(UPat.var("s"),)), lambda s: s), -])+div_and_mod_symbolic + # a CAST over a typed literal is a value conversion: evaluate at the stated width, keep the outer cast. + # keep out of symbolic_simple: create_non_native_float_pats re-expands the pair and cycles on bf16 + (UPat(Ops.CAST, dtypes.all, name="root", src=(UPat(Ops.CAST, dtypes.all, src=(UPat(Ops.CONST, name="c"),)),)), + lambda root,c: root.const_like(c.val)), + # the rules above key on bare CONSTs, so a const wearing a redundant width statement has to be unwrapped in the same fixpoint +])+div_and_mod_symbolic+pm_uncast_const # ******** we take a small aside to "simplify_valid" to rewrite valids ******** diff --git a/tinygrad/uop/weak.py b/tinygrad/uop/weak.py index ed8eb5d4a2..80c0771f78 100644 --- a/tinygrad/uop/weak.py +++ b/tinygrad/uop/weak.py @@ -1,19 +1,32 @@ from dataclasses import replace from tinygrad.dtype import dtypes, DType, AddrSpace, Invalid, least_upper_dtype, strong_dtype, weak_dtype from tinygrad.helpers import unwrap -from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, GroupOp, graph_rewrite, dtype_from_uop +from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, GroupOp, dtype_from_uop, promo_dtype def default_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 commit_weak(s:UOp, dt:DType) -> UOp: - # a CONST commits directly at dt (the value stays mathematical, emission truncates), a non-const src takes the cast + # a CONST re-mints, never takes a cast: at bool/weakint/weakfloat a CAST would be a second spelling of one literal return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt) -def commit_srcs_at(u:UOp, dt:DType) -> UOp: +# the decomps and non-native-float emulation ask this with the width a sibling src states: a bare literal commits there +def commit_weak_sibling(u:UOp, dt:DType|None) -> UOp|None: + return None if dt is None else u.replace(src=tuple(commit_weak(s, dt) if s.op is Ops.CONST and s.dtype in dtypes.weaks else s for s in u.src)) + +# the concrete widths src state: where the operands meet and what u derives (different for comparisons and shifts) +def derived_widths(u:UOp, src:tuple[UOp, ...]) -> tuple[DType, DType]|None: + if u.op not in GroupOp.Broadcastable or (meet:=promo_dtype(src)) in dtypes.weaks \ + or (result:=unwrap(dtype_from_uop(u.op, src, u.arg))) in dtypes.weaks: return None + return meet, result + +def commit_srcs_at(u:UOp, dt:DType) -> UOp|None: # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too - return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)) + widths = derived_widths(u, u.src) + ret = u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks and + not (s.op is Ops.CONST and widths is not None) else s for s in u.src)) + return None if ret is u else ret def commit_weak_srcs(u:UOp) -> UOp|None: if not any(s.dtype in dtypes.weaks for s in u.src) or (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None @@ -21,65 +34,77 @@ def commit_weak_srcs(u:UOp) -> UOp|None: # a concrete CAST over a weak node states the width the value will live at. that width is a floor, never a narrowing def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None: + # only within the kind: an int cast of a weakfloat node is a value conversion, not a statement about the node's width if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None - return commit_srcs_at(u, least_upper_dtype(c.dtype, default_dtype(u))).cast(c.dtype) + return None if (ret:=commit_srcs_at(u, least_upper_dtype(c.dtype, default_dtype(u)))) is None else ret.cast(c.dtype) -pm_cast_weak = PatternMatcher([ - (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs), - (UPat(Ops.CAST, name="c", src=(UPat(Ops.CONST, dtype=dtypes.weaks, name="u"),)), lambda c,u: commit_weak(u, c.dtype)), -]) - -# 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 +# rides every round that can mint a weak const, and must reach fixpoint before pm_lower_weak below hands one its own default 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:]))), + # NOTE: no CONST arm. a concrete CAST over a weak CONST is already the committed pair, minted that way by UOp.const + (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs), ]) +# Every consumer edge absorbs the weak CAST off its srcs (the consumer states that width) and defaults literals whose +# consumer does not derive a width; width-producing ops also settle their compute floor, leaving a trailing weak CAST. # A weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition. +# This round cannot compose symbolic: its cast collapse eats that CAST and cycles. _lower_weak_ops = GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL} def lower_weak_node(u:UOp) -> UOp|None: + if u.op is Ops.CAST and u.src[0].op is Ops.CONST: return None # a literal, not a consumer: its CONST is the value half src = tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src) + widths = derived_widths(u, src) + src = tuple(commit_weak(s, default_dtype(s)) if s.op is Ops.CONST and s.dtype in dtypes.weaks and widths is None else s + for s in src) + if u.op not in _lower_weak_ops: return None if src == u.src else u.replace(dtype=None, src=src) start = 1 if u.op is Ops.WHERE else 0 # WHERE's cond is bool, never part of the width unification - if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None - # Binary can widen from the bounds, all other nodes derive from the lowered sources. + # derivable literals inherit this node's width; wait only on unresolved weak expressions + if src == u.src or any(s.dtype in dtypes.weaks and s.op is not Ops.CONST for s in src[start:]): + return None if src == u.src else u.replace(dtype=None, src=src) + # a Binary widens from its own bounds as well as the lowered srcs, every other op derives from the lowered srcs alone dt = strong_dtype(least_upper_dtype(default_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary else unwrap(dtype_from_uop(u.op, src, u.arg))) - return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else commit_weak(s, dt) for s in src[start:])).cast(u.dtype) + return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid or s.dtype in dtypes.weaks else commit_weak(s, dt) + for s in src[start:])).cast(u.dtype) pm_lower_weak = PatternMatcher([ - (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, default_dtype(u)).cast(u.dtype)), + # a gated long index into a small buffer narrows; its out-of-gate value is discarded + (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))), + allow_any_len=True, name="u"), + lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None), # two stacked weak casts are two kind conversions: each resolves at its own 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(default_dtype(u.src[0])).cast(default_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None), (UPat((Ops.PARAM, Ops.BUFFER), dtype=dtypes.weakint, name="u"), lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=default_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None), - (UPat(_lower_weak_ops, name="u"), lower_weak_node), + (UPat(GroupOp.All, name="u"), lower_weak_node), ]) -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 +# drop a CAST over a weak const where the consumer restores the width anyway, so bare-CONST rules keep matching. two +# statements must survive the drop: the width the operands MEET at, and the node's own DERIVED dtype +def uncast_const(u:UOp) -> UOp|None: + # only a strong CAST is a width statement; a weak CAST over a const is a kind conversion, still resolving + src = tuple(s.src[0] if s.op is Ops.CAST and s.dtype not in dtypes.weaks and s.src[0].op is Ops.CONST + and s.src[0].dtype in dtypes.weaks else s for s in u.src) + if src == u.src or (widths:=derived_widths(u, src)) is None or widths[0] != promo_dtype(u.src) or widths[1] is not u.dtype: return None + return u.replace(src=src) -pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([ - # a CAST between two concrete dtypes over a CONST is a value conversion: evaluate it once, at the width the CAST states - # TODO: delete this once CONST has no dtype - (UPat(Ops.CAST, dtypes.all, name="root", src=(UPat.cvar("c", dtypes.all),)), lambda root, c: root.const_like(c.val)), - (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), - # 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))), - allow_any_len=True, name="u"), - lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None), -]) +# the inverse of pm_cast_const: drop a width statement the consumer re-derives, so bare-keyed rules keep matching. +# composes into symbolic and never symbolic_simple: the widths it drops are first stated by the round above it +pm_uncast_const = PatternMatcher([(UPat(GroupOp.Broadcastable, name="u"), uncast_const)]) + +def cast_const(u:UOp, s:UOp) -> UOp: + if s.op is not Ops.CONST or s.is_invalid: return s # Invalid carries no width: the door never wraps it + # bool is the one strong bare dtype, and its CAST is built raw: .cast(bool) would fold at construction + if s.dtype is dtypes.bool: return UOp(Ops.CAST, src=(s,), arg=s.dtype) + # the width its consumer derives; where nothing does, commit_weak is the identity and spec_program rejects it LOUDLY + return commit_weak(s, widths[0]) if (widths:=derived_widths(u, u.src)) is not None else s + +# THE DOOR: wrap every remaining bare const in its CAST. keyed on the CONSUMER -- "bare" is a property of the edge +def cast_consts(u:UOp) -> UOp|None: + if u.op is Ops.CAST and u.src[0].op is Ops.CONST: return None # a pair's CONST is the value half, not an edge + return None if (src:=tuple(cast_const(u, s) for s in u.src)) == u.src else u.replace(src=src) + +pm_cast_const = PatternMatcher([(UPat(GroupOp.All, name="u"), cast_consts)])