From 5fb3cfb9bcdd8e40e0187ace2d9cd5324e9fe1bf Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 3 Jul 2026 15:00:35 -0400 Subject: [PATCH] remove early THREEFRY const folding [PR] (#16850) it folds automatically once decomposed --- test/backend/test_const_folding.py | 4 +++- tinygrad/uop/symbolic.py | 4 +--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/test/backend/test_const_folding.py b/test/backend/test_const_folding.py index 3731195bf2..4a8d4a350f 100644 --- a/test/backend/test_const_folding.py +++ b/test/backend/test_const_folding.py @@ -2,6 +2,7 @@ import unittest, math from tinygrad import Tensor, Device, dtypes from tinygrad.dtype import DTYPES_DICT from tinygrad.uop.ops import Ops, UOp +from tinygrad.codegen.decomp.op import threefry2x32 import numpy as np from test.helpers import not_support_multi_device @@ -167,7 +168,8 @@ class TestMultiConstFolding(unittest.TestCase): class TestThreefryConstFolding(unittest.TestCase): def test_threefry(self): - x = UOp.const(dtypes.uint64, 5).threefry(UOp.const(dtypes.uint64, 10)) + # THREEFRY(const,const) folds to a const once decomposed + x = threefry2x32(UOp.const(dtypes.uint64, 5), UOp.const(dtypes.uint64, 10)) self.assertIs(x.simplify().op, Ops.CONST) class TestTautologicalCompare(unittest.TestCase): diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 5df25543f9..54a4b43cf4 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -7,7 +7,6 @@ from tinygrad.helpers import partition, all_same, prod, flatten, get_single_elem from tinygrad.uop.divandmod import div_and_mod_symbolic # TODO: symbolic shouldn't be importing from codegen -from tinygrad.codegen.decomp.op import threefry2x32 from tinygrad.codegen.decomp.transcendental import xpow # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ******** @@ -130,9 +129,8 @@ symbolic_simple = propagate_invalid + PatternMatcher([ lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints) # ** constant folding ** (UPat(GroupOp.Unary, src=(UPat((Ops.CONST, Ops.STACK)),), 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(Ops.THREEFRY, src=(UPat.cvar("x"), UPat.cvar("key")), name="a"), - lambda a, x, key: a.const_like(threefry2x32(x, key).simplify().arg)), (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),