From 86e908db5769715780d7246cc0815f87cd9902da Mon Sep 17 00:00:00 2001 From: Sieds Lykles <93992551+S-Lykles@users.noreply.github.com> Date: Wed, 3 Sep 2025 11:05:04 +0200 Subject: [PATCH] cast parents of int64 alu to int32 if possible (#11977) * add overflows helper * add rules * x -> y * check overflow of u too * cleaner * use alu instead of replace to preserve vectorization * just one rule * add test --- test/unit/test_uop_symbolic.py | 6 ++++++ tinygrad/uop/ops.py | 1 + tinygrad/uop/symbolic.py | 2 ++ 3 files changed, 9 insertions(+) diff --git a/test/unit/test_uop_symbolic.py b/test/unit/test_uop_symbolic.py index fa2e90df5a..c852bde02e 100644 --- a/test/unit/test_uop_symbolic.py +++ b/test/unit/test_uop_symbolic.py @@ -732,6 +732,12 @@ class TestSymbolic(unittest.TestCase): a = Variable("a", 1, 10, dtypes.int) self.helper_test_variable(a.trunc(), 1, 10, "a", test_z3=False) + def test_do_math_in_int32(self): + a = Variable("a", 1, 10) + b = Variable("b", 1, 10) + self.helper_test_variable(a.cast(dtypes.long)+b.cast(dtypes.long), 2, 20, "(long)((a+b))") + self.helper_test_variable(a.cast(dtypes.long)*b.cast(dtypes.long), 1, 100, "(long)((a*b))") + class TestSymbolicNumeric(unittest.TestCase): def helper_test_numeric(self, f): MIN, MAX = 0, 10 diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 09f76e97b7..f8e4aa71ee 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -327,6 +327,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): def allreduce(self, op, device:str|tuple[str, ...]|UOp): assert isinstance(self.device, tuple), f"allreduce must be on tuple {self.device} isn't" return UOp(Ops.ALLREDUCE, self.dtype, (self, UOp(Ops.DEVICE, arg=device) if not isinstance(device, UOp) else device), op) + def overflows(self, dtype:DType) -> bool: return self.vmin < dtype.min or dtype.max < self.vmax # *** from MultiLazyBuffer *** diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index fe5a29fd78..97b4e201d2 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -289,6 +289,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ (-1 * (UPat.var("x") + UPat.cvar("c")), lambda x,c: (-x)+(-c)), # -(x+c) -> -x + -c (UPat.var('x', dtypes.ints).cast(dtypes.ints, 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), + (UPat(GroupOp.Binary, src=(UPat.var("x",dtypes.long), UPat.var("y", dtypes.long)), name="u"), lambda u,x,y: + x.cast(dtypes.int).alu(u.op, y.cast(dtypes.int)).cast(u.dtype) if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None), # a conditional with the same results either way is a noop, also fold const conditionals (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val), (UPat.cvar("gate", vec=False).where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),