From c9849b015fa60335bba7a486803202424b1dfe95 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Wed, 12 Aug 2026 16:03:07 +0200 Subject: [PATCH] Inline SymbolicOps at the start of the fusion pass so their bodies can fuse --- pytensor/compile/mode.py | 4 +- pytensor/tensor/rewriting/ofg.py | 51 ++++++++++++++++++++++- tests/tensor/rewriting/test_elemwise.py | 54 +++++++++++++++++++++++++ tests/tensor/rewriting/test_special.py | 5 ++- tests/tensor/test_special.py | 6 ++- 5 files changed, 116 insertions(+), 4 deletions(-) diff --git a/pytensor/compile/mode.py b/pytensor/compile/mode.py index 241c65b3ca..fe159d6be8 100644 --- a/pytensor/compile/mode.py +++ b/pytensor/compile/mode.py @@ -524,7 +524,9 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs): ) PYTORCH = Mode( PytorchLinker(), - RewriteDatabaseQuery(include=["fast_run"]), + # `Composite` has no PyTorch dispatch (it has no `nfunc_spec` to look up in torch), + # and torch fuses its own kernels anyway. + RewriteDatabaseQuery(include=["fast_run"], exclude=["fusion"]), ) MLX = Mode( diff --git a/pytensor/tensor/rewriting/ofg.py b/pytensor/tensor/rewriting/ofg.py index 2f98d1654e..e9cacde348 100644 --- a/pytensor/tensor/rewriting/ofg.py +++ b/pytensor/tensor/rewriting/ofg.py @@ -1,8 +1,17 @@ from pytensor.compile.rewriting import inline_ofg_node from pytensor.graph import node_rewriter +from pytensor.graph.rewriting.basic import MergeOptimizer, dfs_rewriter from pytensor.tensor.basic import AllocDiag from pytensor.tensor.rewriting.basic import register_specialize -from pytensor.tensor.special import LogAddExp, LogSumExp, XLog1PY, XLogY +from pytensor.tensor.rewriting.elemwise import fuse_seqopt +from pytensor.tensor.special import ( + LogAddExp, + LogSoftmax, + LogSumExp, + Softmax, + XLog1PY, + XLogY, +) @register_specialize("inline_ofg") @@ -32,3 +41,43 @@ def late_inline_OpFromGraph(fgraph, node): """ return inline_ofg_node(node) + + +@node_rewriter([Softmax, LogSoftmax, LogSumExp, LogAddExp, XLogY, XLog1PY]) +def inline_symbolic_for_fusion(fgraph, node): + """Inline `SymbolicOp`s so their bodies can fuse. + + The inner graph is otherwise compiled on its own, making the op a fusion barrier + both inside and across its boundary. Ops are listed explicitly, as some are worth + keeping opaque anyway (`KroneckerProduct` is: its body is one broadcasting `mul` + behind a `reshape`). Backends that dispatch these ops directly (JAX, PyTorch and MLX + have their own `Softmax`) exclude the whole fusion pass, so they keep them whole. + + Ops that `late_inline_OpFromGraph` already inlines at specialize are listed here + too, since fusion also runs in modes that skip specialize. + """ + return inline_ofg_node(node) + + +# Fusion is the only thing that needs the body, so inline as its first step: every +# Op-level rewrite still matches on the op. +fuse_seqopt.register( + "inline_symbolic_for_fusion", + dfs_rewriter(inline_symbolic_for_fusion), + "fast_run", + "fusion", + position=0, +) +# Each op is inlined from its own inner graph, so two ops over the same inputs (a +# `Softmax` and a `LogSoftmax` over the same logits) yield structurally identical but +# distinct copies of everything they share. The `merge2` pass cannot clean this up: it +# sits at the same optdb position as the fusion pass, where ties are broken by name, so +# it only runs once fusion has pulled both copies into the `Composite`. +fuse_seqopt.register( + "merge_inlined_symbolic", + MergeOptimizer(), + "fast_run", + "fusion", + "merge", + position=0.5, +) diff --git a/tests/tensor/rewriting/test_elemwise.py b/tests/tensor/rewriting/test_elemwise.py index 61ead5e87f..9ad6e065f5 100644 --- a/tests/tensor/rewriting/test_elemwise.py +++ b/tests/tensor/rewriting/test_elemwise.py @@ -1458,6 +1458,60 @@ def test_many_fused_subgraphs(self): fgraph = FunctionGraph([x_row, x_mat], [m]) FusionOptimizer().apply(fgraph) + def test_fuse_across_symbolic_op(self): + # A SymbolicOp with no backend dispatch is a fusion barrier: its body is + # compiled on its own, so the ops around it cannot join the ops inside it. + # Inlining it at the start of the fusion pass lets the whole chain fuse -- + # the subtraction from the log_softmax body ends up in the same Composite + # as the surrounding exp and add. + rng = np.random.default_rng(75) + x = pt.tensor("x", shape=(5, 4)) + out = exp(pt.special.log_softmax(x * 2.0, axis=-1)) + 1.0 + + def output_scalar_ops(fn): + scalar_op = fn.maker.fgraph.outputs[0].owner.op.scalar_op + assert isinstance(scalar_op, Composite) + return {node.op for node in scalar_op.fgraph.apply_nodes} + + fn = function([x], out, mode=self.mode) + assert output_scalar_ops(fn) == {ps.add, ps.exp, ps.sub} + + # Without the inlining the body stays sealed off, so only the outer ops fuse + barrier_fn = function( + [x], out, mode=self.mode.excluding("inline_symbolic_for_fusion") + ) + assert output_scalar_ops(barrier_fn) == {ps.add, ps.exp} + + x_val = rng.normal(size=(5, 4)).astype(config.floatX) + np.testing.assert_allclose(fn(x_val), barrier_fn(x_val), rtol=1e-6) + + def test_merge_inlined_symbolic_ops(self): + # Each SymbolicOp is inlined from its own inner graph, so a Softmax and a + # LogSoftmax over the same logits come out with distinct copies of everything + # their bodies share. Merging them before fusion is what keeps the Composite + # from evaluating the exponential over each copy. + x = pt.tensor("x", shape=(5, 4)) + outs = [pt.special.softmax(x, axis=-1), pt.special.log_softmax(x, axis=-1)] + + def n_exp(fn): + scalar_ops = [] + for node in fn.maker.fgraph.toposort(): + scalar_op = getattr(node.op, "scalar_op", None) + if isinstance(scalar_op, Composite): + scalar_ops.extend(n.op for n in scalar_op.fgraph.apply_nodes) + elif scalar_op is not None: + scalar_ops.append(scalar_op) + return scalar_ops.count(ps.exp) + + fn = function([x], outs, mode=self.mode) + assert n_exp(fn) == 1 + + # Without the merge the two copies survive into the Composite + unmerged_fn = function( + [x], outs, mode=self.mode.excluding("merge_inlined_symbolic") + ) + assert n_exp(unmerged_fn) == 2 + class TimesN(ps.basic.UnaryScalarOp): """ diff --git a/tests/tensor/rewriting/test_special.py b/tests/tensor/rewriting/test_special.py index f2293c3e54..9f118f3e21 100644 --- a/tests/tensor/rewriting/test_special.py +++ b/tests/tensor/rewriting/test_special.py @@ -26,7 +26,10 @@ from tests.unittest_tools import RewriteTester -_fast_run_rewrites = RewriteDatabaseQuery(include=["fast_run"]) +# Fusion inlines `Softmax` and `LogSoftmax`, which would hide the ops these tests check for +_fast_run_rewrites = RewriteDatabaseQuery( + include=["fast_run"], exclude=["inline_symbolic_for_fusion"] +) _fast_run_rewrites = optdb.query(_fast_run_rewrites) diff --git a/tests/tensor/test_special.py b/tests/tensor/test_special.py index cc50f1175e..49ddbb15f4 100644 --- a/tests/tensor/test_special.py +++ b/tests/tensor/test_special.py @@ -69,7 +69,11 @@ def test_infer_shape(self): rng = np.random.default_rng(utt.fetch_seed()) admat_val = rng.random((3, 4)).astype(config.floatX) self._compile_and_check( - [admat], [Softmax(axis=-1)(admat)], [admat_val], Softmax + [admat], + [Softmax(axis=-1)(admat)], + [admat_val], + Softmax, + excluding=["inline_symbolic_for_fusion"], ) def test_vector_perform(self):