nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

no pm_fold_cast_const in UOp.simplify and hcq2 [pr] (#17540)

Browse Source
This commit is contained in:
chenyu
2026-08-14 21:39:47 -04:00
committed by GitHub
parent 64ccbde3bb
commit 4b0525e594
4 changed files with 18 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_const_folding.py
+1 -12
View File
@@ -1,6 +1,6 @@
import unittest, itertools, math
from tinygrad import Tensor, dtypes, Context
from tinygrad.dtype import DType, ConstType, truncate
from tinygrad.dtype import DType, ConstType
from tinygrad.uop.ops import Ops, UOp
from test.helpers import full_rewrite
import numpy as np
@@ -59,17 +59,6 @@ class TestWeakConstFolding(unittest.TestCase):
# correctly rounded: within half a float32 ulp of the exact value (folding at float32 lands 0.66 ulp off)
self.assertLess(abs(const - 1/math.log(10)), 2**-26)
def test_cast_commits_to_dtype_grid(self):
# committing a weak const to a stated width puts the value on that width's grid, same as storage packing and native compilers
v = 1/123008 # not representable in float16
out = UOp.const(v).cast(dtypes.half).simplify()
self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.half, truncate[dtypes.half](v)))
self.assertNotEqual(out.val, v)
# the grid commit preserves the sign of zero
self.assertEqual(math.copysign(1, UOp.const(-0.0).cast(dtypes.half).simplify().val), -1)
# observable at tensor level: the const-folded comparison agrees with the committed value
self.assertTrue((Tensor(-3.2).cast(dtypes.float32) <= truncate[dtypes.float32](-3.2)).item())
class TestBinaryOpsConstFolding(unittest.TestCase):
def test_add_literal_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
tinygrad/runtime/support/hcq2.py
+5 -5
View File
@@ -7,7 +7,7 @@ from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, Co
from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker
from tinygrad.device import ProfileDeviceEvent, ProfileGraphEntry, ProfileGraphEvent
from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, rewrite_group, GroupOp
from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
from tinygrad.uop.symbolic import symbolic
from tinygrad.dtype import dtypes, truncate
from tinygrad.runtime.support.hcq import MMIOInterface, HCQBuffer
from tinygrad.runtime.support.memory import BumpAllocator
@@ -445,7 +445,7 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
linear = graph_rewrite(linear, pm_encode_cmdbufs+pm_pack_placeholders, walk=True, name="encode and pack", enter_calls=True)
# patches and runtime uops
linear = graph_rewrite(linear, pm_early_simplify+symbolic+pm_fold_cast_const, bottom_up=False, name="simplify patches", enter_calls=True)
linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True)
linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
# and compile it
@@ -477,7 +477,7 @@ def fold_binary(buf:UOp, blob:UOp) -> UOp:
def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
for off,val in zip(off.src, val.src):
for b,v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.val))
data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype]((v.src[0] if v.op is Ops.CAST else v).val))
b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(bo:=off.val*buf.dtype.itemsize):bo+len(data)] = data
return UOp(Ops.NOOP)
@@ -492,7 +492,7 @@ def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
pm_resolve_patches = PatternMatcher([
# multi
(UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat(Ops.CONST)], name="op"), push_stack),
(UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat.any(UPat(Ops.CONST), UPat(Ops.CAST, src=(UPat(Ops.CONST),)))], name="op"), push_stack),
(UPat(Ops.CAST, src=(UPat(Ops.STACK, name="s"),), name="op"), push_stack),
# getaddr
@@ -518,7 +518,7 @@ def hcq_link(linear:UOp, cache=True) -> UOp:
bufs = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1)
if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS}
linear = linear.substitute({x:link_buf_cache[k] for a in bufs.values() if (k:=link_buf_key(a)) in link_buf_cache for x in (a, a.src[0])}, walk=True)
linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_fold_cast_const+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False,
linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False,
name="resolve patches")
for (j,i),a in bufs.items(): link_buf_cache.setdefault(link_buf_key(a), linear.src[j].src[i])
if cache: link_linear_cache[linear_key] = linear
tinygrad/uop/ops.py
+5 -3
View File
@@ -517,10 +517,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
if self.op is Ops.CONST: return self
if self.op is Ops.SINK and all(s.op is Ops.CONST or (s.op is Ops.STACK and len(s.src) == 0) for s in self.src): return self
# late import!
from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
from tinygrad.uop.symbolic import symbolic
with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value):
return graph_rewrite(self, symbolic+pm_fold_cast_const, name="simplify")
def ssimplify(self) -> UOp|ConstType: return ret.val if (ret:=self.simplify()).op is Ops.CONST else ret
return graph_rewrite(self, symbolic, name="simplify")
def ssimplify(self) -> UOp|ConstType:
if (ret := self.simplify()).op is Ops.CAST and ret.src[0].op is Ops.CONST: return ret.dtype.const(ret.src[0].val)
return ret.val if ret.op is Ops.CONST else ret
def _eval(self, dtype, expected_type:Type[T]) -> T:
assert self.dtype in dtype, f"eval with wrong dtype {self}"
vmin, vmax = (simple_self:=self.simplify())._min_max
tinygrad/uop/symbolic.py
+7 -3
View File
@@ -23,8 +23,12 @@ 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))
# 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),)))
def const_arg(u:UOp) -> ConstType|tuple[ConstType, ...]|None:
if u.op is Ops.CONST: return u.val
if u.op is Ops.CAST and u.src[0].op is Ops.CONST: return u.dtype.const(u.src[0].val)
if u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src): return tuple(s.val for s in u.src)
return None
@@ -136,10 +140,10 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"),
lambda x: x.const_like(False, dtypes.bool)), # x != x -> False (only ints)
# ** constant folding **
(UPat(GroupOp.Unary, src=(UPat((Ops.CONST, Ops.STACK)),), name="a"), fold_const_alu),
(UPat(GroupOp.Unary, src=(const_folding_pat,), 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(GroupOp.Ternary, src=(UPat((Ops.CONST, Ops.STACK)),)*3, name="a"), fold_const_alu),
(UPat(GroupOp.Binary-{Ops.THREEFRY}, src=(const_folding_pat,)*2, name="a"), fold_const_alu),
(UPat(GroupOp.Ternary, src=(const_folding_pat,)*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),
(UPat.var('x', dtype=dtypes.bool) + UPat.var('y', dtype=dtypes.bool), lambda x,y: x|y),