sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

weak 1 for FDIV in get_late_rewrite_patterns [PR] (#17655)

Browse Source
This commit is contained in:
chenyu
2026-08-21 09:33:43 -04:00
committed by GitHub
parent 1c3c9e96f6
commit 176377ff6e
2 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/__init__.py
+4 -4
View File
@@ -377,6 +377,10 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
pm_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends
sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")
# spell every literal as a casted const CAST(dt, CONST(value))
# TODO: remove once consts are always weak
sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
# add implicit barriers (stores/loads through LOCAL memory ordered by AFTER or across loop iterations need workgroup barriers)
sink = graph_rewrite(sink, pm_implicit_barriers, name="add implicit barriers")
@@ -387,10 +391,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1])
sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
# spell every literal as a casted const CAST(dt, CONST(value))
# TODO: remove once consts are always weak
sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
if SPEC: type_verify(sink, spec_program)
tinygrad/codegen/decomp/op.py
+2 -2
View File
@@ -128,6 +128,6 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1<<n.val), c))]
# some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
if Ops.FDIV in ops:
pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
pat += [(UPat.var("x").reciprocal(), lambda x: UOp.const(1.0).alu(Ops.FDIV, x))]
pat += [(UPat.var("a") * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
return PatternMatcher(pat)