tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 14:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fold_bitcast should truncate its input [pr] (#17601)

Browse Source
This commit is contained in:
chenyu
2026-08-19 14:39:36 -04:00
committed by GitHub
parent bd6e70ac15
commit c218b4842d
2 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_const_folding.py
+6
View File
@@ -23,6 +23,12 @@ class TestWeakConstFolding(unittest.TestCase):
self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
class TestBitcastConstFolding(unittest.TestCase):
def test_out_of_range_source_value(self):
for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000),
(70000, dtypes.int16, dtypes.uint16, 4464),
(-5, dtypes.uint32, dtypes.int32, -5)):
self.assertEqual(UOp.const(val, src_dt).bitcast(dst_dt).simplify().val, bits)
def test_scalar_bitcast(self):
def t(cases: dict[DType, ConstType]):
for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
tinygrad/uop/symbolic.py
+2 -2
View File
@@ -2,7 +2,7 @@
import math
from collections import defaultdict
from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid, bitcast
from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid, bitcast, truncate
from tinygrad.helpers import partition, all_same, prod, flatten, unwrap, IMAGE, dedup
from tinygrad.uop.divandmod import div_and_mod_symbolic
from tinygrad.uop.movement import mop_cleanup
@@ -22,7 +22,7 @@ def simplify_pow(x:UOp, c:UOp) -> UOp|None:
def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
if c.dtype.fmt is None or root.dtype.fmt is None or c.dtype.itemsize != root.dtype.itemsize: return None
return root.const_like(bitcast(c.val, c.dtype, root.dtype))
return root.const_like(bitcast(truncate[c.dtype](c.val), c.dtype, root.dtype))
# const folding works for CONST, STACK, and casted CONST
const_folding_pat = UPat.any(UPat((Ops.CONST, Ops.STACK)), UPat(Ops.CAST, src=(UPat(Ops.CONST),)))