From 539a03343a45a94893f07ececeea6e80b651e8dc Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 15 Aug 2026 07:47:40 -0400 Subject: [PATCH] no casted const from sub and div [pr] (#17543) --- test/unit/test_dtype_weak.py | 5 +++++ tinygrad/mixin/elementwise.py | 8 +++++--- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index a4e9740eb9..ea6a612a3c 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -76,6 +76,11 @@ class TestWeakPromotion(unittest.TestCase): 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_div_sub_operand_kept_weak(self): + a = Tensor.empty(4, dtype=dtypes.float32) + for t in (a / 1, a - 0): + self.assertEqual(t.uop.src[1].dtype, dtypes.weakfloat) + 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): diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index f06fd5aa52..3818ef0b5e 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -115,7 +115,8 @@ class ElementwiseMixin(CreationMixin): ``` """ a, b = self._broadcasted(x, reverse) - return a + (-b) + # alu, not +: _broadcasted already promoted these, and a second promote would cast -b (only a bare weak CONST is kept weak) + return a.alu(Ops.ADD, -b) def mul(self, x: Self | ConstType, reverse: bool = False) -> Self: """ @@ -245,8 +246,9 @@ class ElementwiseMixin(CreationMixin): if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype): if rounding_mode == "trunc": return a.alu(Ops.CDIV, b) if rounding_mode == "floor": return a.alu(Ops.FLOORDIV, b) - a = a.cast(dtypes.default_float) - d = a * b.reciprocal() + if dtypes.is_int(a.dtype) or a.dtype == dtypes.bool: a = a.cast(dtypes.default_float) + # alu, not *: _broadcasted already promoted these, and a second promote would cast 1/b (only a bare weak CONST is kept weak) + d = a.alu(Ops.MUL, b.reciprocal()) if rounding_mode is None: return d if rounding_mode == "trunc": return d.trunc() if rounding_mode == "floor": return d.floor()