From c2f1e5ae2a1aa03206734f46ec2c2bcf7300a9e4 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 5 Aug 2026 15:32:23 -0400 Subject: [PATCH] fix weak cast to strong dtype [pr] (#17418) weak can mean higher than that strong dtype, so always use that strong dtype is wrong --- test/unit/test_dtype_weak.py | 18 ++++++++++++++++-- tinygrad/uop/ops.py | 13 ++++++++----- 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index ae3672d064..9d50529273 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -64,14 +64,28 @@ class TestWeakPromotion(unittest.TestCase): def test_weak_expression_anchors_at_strong_lub(self): # regression test for the HALF bert nan (#17408, reverted in #17409): lub(int32, weakfloat)==weakfloat makes - # `loss_mask.sum() + 1e-5` a weakfloat EXPRESSION. Meeting a strong float in a binop must pin it at the lub, otherwise - # nothing owns a width until the bufferize/codegen commits it at default_float: a HALF reciprocal of (sum+1e-5) is inf + # `loss_mask.sum() + 1e-5` a weakfloat EXPRESSION. Meeting a strong float in a binop must pin it at the lub denom = (Tensor.zeros(912, dtype=dtypes.int32) != Tensor.zeros(912, dtype=dtypes.float32)).sum() + 1e-5 self.assertIs(denom.dtype, dtypes.weakfloat) # the setup: the denominator expression itself is weak x, y = Tensor([2048.0], dtype=dtypes.float32)._broadcasted(denom) self.assertIs(y.dtype, dtypes.float32) 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={}) + self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32]) + + 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={}) + self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32)) + + def test_cast_weak_expression_value_uses_cast_floor(self): + with Context(DEFAULT_FLOAT=dtypes.float16): + denom = Tensor.ones(1, dtype=dtypes.int32, device="CPU").sum() * 70000 + 1e-5 + out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom + self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10) def test_uop_scalar_const_lifts_kind(self): for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index ccf39be4dc..2e0573edb4 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1805,14 +1805,17 @@ pm_commit_weak = PatternMatcher([ lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))), ]) -# push cast to weak src +# 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: + if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None + dt = least_upper_dtype(c.dtype, select_dtype(u)) + return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype) + 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), + (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs), ]) -pm_lower_index_dtype = pm_commit_weak+PatternMatcher([ +pm_lower_index_dtype = pm_commit_weak+pm_cast_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), # 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)