From 18cdbec447b9504c0af5a3e6b6712c433c1b3f75 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 12 Aug 2025 12:56:05 -0700 Subject: [PATCH] split decompositions pass (#11638) * split decompositions pass * fix ptx * pack load store early * restore that --- tinygrad/codegen/__init__.py | 6 +++++- tinygrad/renderer/ptx.py | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 2f11807df0..c876b48309 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -82,11 +82,15 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC supported_ops = tuple(opts.code_for_op.keys()) extra_matcher = opts.extra_matcher if opts.extra_matcher is not None else PatternMatcher([]) + # decompositions + pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2) + ret.append(RewriteStep(pm_decomp, name="decompositions")) + # optional pre matcher if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher")) # final rules for the renderer (without sym) - pm_final_rewrite = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2)+pm_render+extra_matcher + pm_final_rewrite = pm_decomp+pm_render+extra_matcher ret.append(RewriteStep(pm_final_rewrite, lambda _: opts.device, name="final rewrite")) # return the list (with optional linearizer) diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index 5f957fbdcc..1fd63ee8d9 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -38,6 +38,7 @@ doesnt_support_half: tuple[Ops, ...] = tuple(op for op in asm_for_op.keys() if o ptx_matcher = PatternMatcher([ # bool CMPNE is XOR, bool CMPLT is XOR+AND (universal makes this slow, this is for renderer only) (UPat.var('x', dtype=dtypes.bool).ne(UPat.var('y')), lambda x,y: x^y), + (UPat.var('x', dtype=dtypes.bool).alu(Ops.CMPEQ, UPat.var('y')), lambda x,y: (x^y)^True), (UPat.var('x', dtype=dtypes.bool)