forked from tinygrad/tinygrad
remove early THREEFRY const folding [PR] (#16850)
it folds automatically once decomposed
This commit is contained in:
@@ -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):
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user