Fix BufferizationStage IR blowup for runtime-parameterized Hamiltonians (Reroll)

doc/releases/changelog-dev.md

@@ -166,6 +166,15 @@
 * Added ``CZ`` support to ``to-ppr`` pass.
   [(#3009)](https://github.com/PennyLaneAI/catalyst/pull/3009)
 
+* IR amplification for circuits whose gate parameters are computed from runtime values
+  (e.g. `qml.TrotterProduct` with runtime Hamiltonian coefficients) is drastically reduced.
+  The new `scalarize-tensor-extracts` pass turns gate-angle dataflow into scalar arithmetic
+  instead of thousands of small tensors that each survive bufferization as an allocation,
+  and the new `reroll-loops` pass reconstructs the loops that tracing unrolled by rewriting
+  repeated op sequences as `scf.for` loops. On a Trotterized QPE workload with runtime
+  coefficients, compile time, peak memory, and final IR size all drop by large factors.
+  [(#3036)](https://github.com/PennyLaneAI/catalyst/pull/3036)
+
 <h3>Breaking changes 💔</h3>
 
 * Catalyst's xDSL dependencies have been updated to `xdsl` 0.63.0 and `xdsl-jax` 0.5.2.

frontend/test/pytest/test_reroll_loops.py

@@ -0,0 +1,143 @@
+# Copyright 2026 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Integration tests for the reroll-loops pass in the default pipeline:
+rerolling the unrolled Trotter steps of a QPE workload must preserve numerics
+and must actually reconstruct scf.for loops."""
+
+import re
+
+import numpy as np
+import pennylane as qml
+import pytest
+from jax import numpy as jnp
+
+from catalyst import qjit
+from catalyst.debug import get_compilation_stage
+
+# H2/STO-3G-like coefficients; the structure (15 terms, runtime values) is what
+# exercises the pipeline, the values just need to be a valid Hamiltonian.
+COEFFS = [
+    -0.0996,
+    0.1711,
+    0.1711,
+    -0.2225,
+    -0.2225,
+    0.1686,
+    0.0453,
+    -0.0453,
+    -0.0453,
+    0.0453,
+    0.1205,
+    0.1658,
+    0.1658,
+    0.1205,
+    0.1743,
+]
+
+
+def make_ops():
+    """Build a fresh list of Hamiltonian terms (operators can't be reused)."""
+    return [
+        qml.Identity(0),
+        qml.PauliZ(0),
+        qml.PauliZ(1),
+        qml.PauliZ(2),
+        qml.PauliZ(3),
+        qml.PauliZ(0) @ qml.PauliZ(1),
+        qml.PauliY(0) @ qml.PauliX(1) @ qml.PauliX(2) @ qml.PauliY(3),
+        qml.PauliY(0) @ qml.PauliY(1) @ qml.PauliX(2) @ qml.PauliX(3),
+        qml.PauliX(0) @ qml.PauliX(1) @ qml.PauliY(2) @ qml.PauliY(3),
+        qml.PauliX(0) @ qml.PauliY(1) @ qml.PauliY(2) @ qml.PauliX(3),
+        qml.PauliZ(0) @ qml.PauliZ(2),
+        qml.PauliZ(0) @ qml.PauliZ(3),
+        qml.PauliZ(1) @ qml.PauliZ(2),
+        qml.PauliZ(1) @ qml.PauliZ(3),
+        qml.PauliZ(2) @ qml.PauliZ(3),
+    ]
+
+
+N_QUBITS = 4
+N_EST = 2
+N_TROTTER = 6  # enough repetitions for reroll-loops to fire
+
+
+def make_qpe(dev, runtime: bool):
+    """Controlled-power Trotterized QPE, with runtime or compile-time coeffs."""
+
+    @qml.qnode(dev)
+    def qpe_circuit(coeffs):
+        qml.PauliX(0)
+        qml.PauliX(1)
+        for k in range(N_EST):
+            qml.Hadamard(wires=N_QUBITS + k)
+        H = qml.dot(coeffs if runtime else COEFFS, make_ops())
+        for k in range(N_EST):
+            t = 2 ** (N_EST - 1 - k)
+            qml.ctrl(
+                qml.adjoint(
+                    qml.TrotterProduct(H, time=t, n=N_TROTTER, order=2, check_hermitian=False)
+                ),
+                control=N_QUBITS + k,
+            )
+        qml.adjoint(qml.QFT)(wires=range(N_QUBITS, N_QUBITS + N_EST))
+        return qml.probs(wires=range(N_QUBITS, N_QUBITS + N_EST))
+
+    return qpe_circuit
+
+
+class TestRerollLoops:
+    """Numerical equivalence and loop reconstruction for Trotterized workloads
+    through the default pipeline (which rerolls unrolled repeats)."""
+
+    def test_runtime_matches_fixed(self):
+        """qml.dot with traced coefficients must produce the same distribution
+        as the same Hamiltonian with Python-float coefficients."""
+        dev = qml.device("lightning.qubit", wires=N_QUBITS + N_EST)
+        coeffs = jnp.array(COEFFS)
+
+        dyn = qjit(make_qpe(dev, runtime=True))(coeffs)
+        fixed = qjit(make_qpe(dev, runtime=False))(coeffs)
+
+        assert np.allclose(np.asarray(dyn), np.asarray(fixed), atol=1e-9)
+
+    def test_reroll_recovers_loops(self):
+        """The IR after HLO lowering must contain scf.for loops recovered from
+        the unrolled Trotter steps, and fewer gate ops than the unrolled
+        circuit (guards against silent regression of reroll-loops in the
+        default pipeline)."""
+        dev = qml.device("lightning.qubit", wires=N_QUBITS + N_EST)
+        coeffs = jnp.array(COEFFS)
+
+        compiled = qjit(make_qpe(dev, runtime=True), keep_intermediate=True)
+        compiled(coeffs)
+        try:
+            traced = get_compilation_stage(compiled, "QuantumCompilationStage")
+            lowered = get_compilation_stage(compiled, "HLOLoweringStage")
+            # Rerolled Trotter steps are scf.for loops threading qubit values
+            # through iter_args.
+            qubit_loops = re.findall(r"scf\.for .*iter_args\([^)]*\).*->.*!quantum\.bit", lowered)
+            assert qubit_loops, "reroll-loops did not produce qubit-threading loops"
+            unrolled_gates = traced.count("quantum.custom")
+            rerolled_gates = lowered.count("quantum.custom")
+            assert rerolled_gates < unrolled_gates / 2, (
+                f"expected reroll to shrink gate volume by >2x, got "
+                f"{unrolled_gates} -> {rerolled_gates}"
+            )
+        finally:
+            compiled.workspace.cleanup()
+
+
+if __name__ == "__main__":
+    pytest.main(["-x", __file__])

mlir/include/Catalyst/Transforms/Passes.td

@@ -333,6 +333,47 @@ def RegisterDecompRuleResourcePass : Pass<"register-decomp-rule-resource"> {
     }];
 }
 
+def RerollLoopsPass : Pass<"reroll-loops"> {
+    let summary = "Reconstruct loops from unrolled repetitive op sequences.";
+    let description = [{
+        Tracing a Python program unrolls its loops: N structurally identical
+        circuit segments (Trotter steps, layers, folds) arrive as N copies of the
+        same op sequence, and every downstream stage amplifies each copy.
+
+        This pass detects maximal tandem repeats of structurally isomorphic op
+        windows via shift-invariant structural hashing, verifies that
+        cross-window dataflow is limited to values threaded from the directly
+        preceding window (e.g. qubit SSA values) plus loop-invariant values, and
+        replaces the repeat with an `scf.for` whose iter_args are the threaded
+        values. A repeat of multiplicity k shrinks that region k-fold.
+    }];
+
+    let dependentDialects = [
+        "mlir::arith::ArithDialect",
+        "mlir::scf::SCFDialect"
+    ];
+
+    let options = [
+        Option<
+            /*C++ var name=*/"minPeriod",
+            /*CLI arg name=*/"min-period",
+            /*type=*/"unsigned",
+            /*default=*/"2",
+            /*description=*/"Minimum number of ops per repeated window."
+        >,
+        Option<
+            /*C++ var name=*/"minSavings",
+            /*CLI arg name=*/"min-savings",
+            /*type=*/"unsigned",
+            /*default=*/"32",
+            /*description=*/
+            "Minimum number of ops a reroll must eliminate to be applied. The "
+            "default is deliberately conservative so that only substantial "
+            "repeats (e.g. unrolled Trotter steps) are converted to loops."
+        >
+    ];
+}
+
 def EmptyPass : Pass<"empty"> {
     let summary = "Empty pass that does nothing.";
 

mlir/include/Driver/DefaultPipelines/DefaultPipelines.h

@@ -71,6 +71,13 @@ const PipelineList pipelineList{
       "scatter-lowering",
       "hlo-custom-call-lowering",
       "cse",
+      // Reconstruct the loops that tracing unrolled (e.g. Trotter steps); a
+      // repeat of multiplicity k shrinks its region k-fold before the
+      // bufferization and LLVM stages amplify it. Computations duplicated per
+      // iteration sit inside each repeated window and reroll as part of it,
+      // and constants (the only cross-window operands that must be identical
+      // SSA values) are already uniqued by the preceding cse.
+      "reroll-loops",
       "func.func(linalg-detensorize{aggressive-mode})",
       "detensorize-scf",
       "detensorize-function-boundary",

mlir/lib/Catalyst/Transforms/CMakeLists.txt

@@ -22,6 +22,7 @@ file(GLOB SRC
     qnode_to_async_lowering.cpp
     QnodeToAsyncPatterns.cpp
     RegisterInactiveCallbackPass.cpp
+    RerollLoopsPass.cpp
     ResourceAnalysisPass.cpp
     RegisterDecompRuleResourcePass.cpp
     SplitMultipleTapes.cpp