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

doc/releases/changelog-dev.md

@@ -189,6 +189,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.
+  [(#3013)](https://github.com/PennyLaneAI/catalyst/pull/3013)
+
 <h3>Breaking changes 💔</h3>
 
 * Python 3.11 is no longer supported. Catalyst now requires Python 3.12 or newer.

frontend/test/lit/test_detensorize.py

@@ -45,15 +45,15 @@ def f_with_cond(a, b):
     # CHECK-NOT: linalg.generic
     a2 = a + a
     if a2 > b:
-        # CHECK: arith.subf
+        # CHECK-DAG: arith.subf
         a = a - 2.0
-        # CHECK: arith.mulf
+        # CHECK-DAG: arith.mulf
         b = b * 2.0
         c = a + b
     else:
-        # CHECK: arith.mulf
+        # CHECK-DAG: arith.mulf
         a = a * 2.0
-        # CHECK: arith.subf
+        # CHECK-DAG: arith.subf
         b = b - 2.0
         c = a + b
     return c * 2.0

frontend/test/pytest/test_trotter_runtime_coeffs.py

@@ -0,0 +1,117 @@
+# 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 IR-amplification fixes for runtime-coefficient
+Hamiltonians: scalarize-tensor-extracts and elementwise fusion in the default
+pipeline must preserve numerics for Trotterized workloads with runtime
+coefficients."""
+
+import numpy as np
+import pennylane as qml
+import pytest
+from jax import numpy as jnp
+
+from catalyst import qjit
+
+# 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 the Hamiltonian's Pauli-word operators."""
+    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
+
+
+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 TestRuntimeCoefficientTrotter:
+    """Numerical equivalence of runtime- and fixed-coefficient Trotterization
+    through the default pipeline (which scalarizes and fuses)."""
+
+    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)
+
+
+if __name__ == "__main__":
+    pytest.main(["-x", __file__])

mlir/include/Catalyst/Transforms/Passes.td

@@ -333,6 +333,30 @@ def RegisterDecompRuleResourcePass : Pass<"register-decomp-rule-resource"> {
     }];
 }
 
+def ScalarizeTensorExtractsPass : Pass<"scalarize-tensor-extracts"> {
+    let summary = "Sink scalar tensor.extract ops through small-tensor producers.";
+    let description = [{
+        Programs that compute quantum gate parameters from runtime inputs (e.g.
+        Trotterization with runtime Hamiltonian coefficients) produce long chains of
+        small tensor operations whose only consumers are scalar `tensor.extract`
+        operations. Each intermediate tensor survives bufferization as an allocation
+        with copies and each `linalg.generic` is later unrolled into loops, causing
+        severe IR amplification.
+
+        This pass folds `tensor.extract` through elementwise `linalg.generic` (by
+        inlining the scalar payload), `tensor.collapse_shape`, and
+        `tensor.extract_slice`, so the extracted element is computed directly in
+        scalar arithmetic and the intermediate tensors become dead. Payload inlining
+        is limited to small statically-shaped tensors to bound code growth.
+    }];
+
+    let dependentDialects = [
+        "mlir::arith::ArithDialect",
+        "mlir::linalg::LinalgDialect",
+        "mlir::tensor::TensorDialect"
+    ];
+}
+
 def EmptyPass : Pass<"empty"> {
     let summary = "Empty pass that does nothing.";
 

mlir/include/Driver/DefaultPipelines/DefaultPipelines.h

@@ -71,9 +71,20 @@ const PipelineList pipelineList{
       "scatter-lowering",
       "hlo-custom-call-lowering",
       "cse",
+      // Sink scalar extractions through small-tensor producers and fuse the
+      // remaining elementwise ops. Traced gate-parameter dataflow (e.g. runtime
+      // Hamiltonian coefficients) otherwise reaches bufferization as thousands
+      // of tiny tensor ops that each become an alloc + copy.
+      "scalarize-tensor-extracts",
+      "func.func(linalg-fuse-elementwise-ops)",
+      "canonicalize",
       "func.func(linalg-detensorize{aggressive-mode})",
       "detensorize-scf",
       "detensorize-function-boundary",
+      // Detensorization is what materializes tensor.extract on the gate-angle
+      // dataflow, so scalarization must run again here to fold the
+      // extract_slice/collapse_shape chains it exposes.
+      "scalarize-tensor-extracts",
       "canonicalize",
       "symbol-dce"}},
     {"gradient-lowering-stage",

mlir/lib/Catalyst/Transforms/CMakeLists.txt

@@ -23,6 +23,7 @@ file(GLOB SRC
     QnodeToAsyncPatterns.cpp
     RegisterInactiveCallbackPass.cpp
     ResourceAnalysisPass.cpp
+    ScalarizeTensorExtractsPass.cpp
     RegisterDecompRuleResourcePass.cpp
     SplitMultipleTapes.cpp
     TBAAPatterns.cpp