tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-09-06 01:06:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

Double cast is one cast if intermediate cast is safe (#11939)

Browse Source 
* add rule

* add some tests

* prevent infinite loop with bfloat16

* prevent more infinite rewrite
This commit is contained in:
Sieds Lykles
2025-09-01 00:36:29 +02:00
committed by GitHub
parent 7e68045fb2
commit 1c6e43c203
2 changed files with 24 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_uops.py
+20
View File
@@ -177,6 +177,26 @@ class TestBoolUOps(TestUOps):
def test_cmplt_bool(self): self._test_bop_bool_fxn(Ops.CMPLT, lambda a,b: a < b)
def test_where_bool(self): self._test_top_bool_fxn(Ops.WHERE, lambda a,b,c: b if a else c)
class TestSafeCast(TestUOps):
def test_cast_folds(self):
a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a)
# TODO: double preserves certain int dtypes
# self.assertEqual(a.cast(dtypes.double).cast(dtypes.int32).simplify(), a)
a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
self.assertEqual(a.cast(dtypes.int64).cast(dtypes.uint8).simplify(), a)
self.assertEqual(a.cast(dtypes.uint32).cast(dtypes.uint8).simplify(), a)
def test_remove_intermediate_cast(self):
a = UOp.variable("a", 0., 100., dtype=dtypes.half)
self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
# TODO: double preserves certain int dtypes
# self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int16).simplify(), a.cast(dtypes.int16))
a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a.cast(dtypes.int32))
class TestExecALU(TestUOps):
def test_sqrt(self):
self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
tinygrad/uop/symbolic.py
+4 -1
View File
@@ -68,7 +68,10 @@ symbolic_simple = PatternMatcher([
(UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None),
(UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
# b.cast(a).cast(b) -> b if a preserves all values in b
(UPat.var('x').cast().named('a').cast().named('b'), lambda x,a,b: x if x.dtype == b.dtype and can_safe_cast(b.dtype, a.dtype) else None),
(UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_safe_cast(b.dtype, a.dtype) else None),
# if the intermediate cast doesnt narrow we can do it in one cast, we have to be carefull with bfloat16
(UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_safe_cast(x.dtype, a.dtype) and
not (a.dtype==dtypes.float and (b.dtype==dtypes.bfloat16 or x.dtype==dtypes.bfloat16)) else None),
# ** pow **
(UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow),
# positive const ** x