diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index ee1b3b0caf..2dd337be02 100644
--- a/doc/releases/changelog-dev.md
+++ b/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)
+
Breaking changes 💔
* Catalyst's xDSL dependencies have been updated to `xdsl` 0.63.0 and `xdsl-jax` 0.5.2.
diff --git a/frontend/test/pytest/test_reroll_loops.py b/frontend/test/pytest/test_reroll_loops.py
new file mode 100644
index 0000000000..852190572b
--- /dev/null
+++ b/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__])
diff --git a/mlir/include/Catalyst/Transforms/Passes.td b/mlir/include/Catalyst/Transforms/Passes.td
index c5f87f0d7c..6196e4eff7 100644
--- a/mlir/include/Catalyst/Transforms/Passes.td
+++ b/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.";
diff --git a/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h b/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h
index 9c652e7732..01bbf4349e 100644
--- a/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h
+++ b/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",
diff --git a/mlir/lib/Catalyst/Transforms/CMakeLists.txt b/mlir/lib/Catalyst/Transforms/CMakeLists.txt
index bb794d3b4c..15cefe6437 100644
--- a/mlir/lib/Catalyst/Transforms/CMakeLists.txt
+++ b/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
diff --git a/mlir/lib/Catalyst/Transforms/RerollLoopsPass.cpp b/mlir/lib/Catalyst/Transforms/RerollLoopsPass.cpp
new file mode 100644
index 0000000000..b60e5265bf
--- /dev/null
+++ b/mlir/lib/Catalyst/Transforms/RerollLoopsPass.cpp
@@ -0,0 +1,545 @@
+// 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.
+
+#define DEBUG_TYPE "reroll-loops"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace llvm;
+using namespace mlir;
+
+// Tracing unrolls Python loops: N identical circuit segments (Trotter steps,
+// layers, folds) arrive as N copies of the same op sequence, and every later
+// stage amplifies each copy. This pass reconstructs the loops in four steps:
+// structural hashing of each op, tandem-repeat detection on the hash sequence
+// (maximal runs with H[i] == H[i+p]), semantic verification that consecutive
+// windows are isomorphic with cross-window dataflow limited to threaded values
+// (e.g. qubits) plus loop invariants, and replacement of the repeat with an
+// scf.for whose iter_args are the threaded values.
+
+namespace {
+
+/// Verifying a candidate period is linear in the block, so bound how many
+/// distinct periods (most-promising first) are tried per block.
+constexpr size_t kMaxPeriodsPerBlock = 16;
+
+/// Minimum number of windows for a repeat to become a loop. Two windows are
+/// cheaper as straight-line code and give extension no signal to distinguish
+/// coincidental hash matches from real repetition.
+constexpr size_t kMinWindows = 3;
+
+/// Rerolling creates new blocks (loop bodies) that may themselves contain
+/// repeats (nested loops), so the driver iterates; each round strictly
+/// shrinks the IR, and in practice one nesting level per round suffices.
+constexpr unsigned kMaxRounds = 4;
+
+/// Number of rounds of operand-hash mixing (see refineHashes). Deep enough to
+/// tell apart structurally similar ops at different positions of a window
+/// (e.g. the same gate applied to different qubits of a chain), shallow enough
+/// that an op's hash rarely depends on ops more than a window away.
+constexpr unsigned kHashDepth = 3;
+
+/// Ops that may participate in a rerolled window: single-block-region-free,
+/// no successors, and not a terminator.
+bool isRerollableOp(Operation *op)
+{
+ return op->getNumRegions() == 0 && op->getNumSuccessors() == 0 &&
+ !op->hasTrait();
+}
+
+/// Compute structural hashes for all ops of a block (excluding the
+/// terminator). Non-rerollable ops receive unique sentinel hashes. The hash
+/// covers the op name, attributes, and result types; in-block operand defs
+/// contribute only their result number (hashing them by identity would make
+/// every window distinct), out-of-block values their identity. The weak
+/// discrimination this leaves is compensated by semantic verification.
+void computeHashes(ArrayRef ops, const llvm::DenseMap &indexOf,
+ SmallVectorImpl &hashes)
+{
+ uint64_t sentinel = 0;
+ for (Operation *op : ops) {
+ if (!isRerollableOp(op)) {
+ hashes.push_back(hash_combine(0xdeadbeefULL, ++sentinel));
+ continue;
+ }
+ hash_code h =
+ hash_combine(op->getName().getTypeID(), op->getAttrDictionary().getAsOpaquePointer());
+ for (Type t : op->getResultTypes()) {
+ h = hash_combine(h, t.getAsOpaquePointer());
+ }
+ for (Value v : op->getOperands()) {
+ Operation *def = v.getDefiningOp();
+ auto it = def ? indexOf.find(def) : indexOf.end();
+ if (it != indexOf.end()) {
+ h = hash_combine(h, 1, cast(v).getResultNumber());
+ }
+ else {
+ // Block argument or out-of-block def: hash the value identity.
+ h = hash_combine(h, 2, v.getAsOpaquePointer());
+ }
+ }
+ hashes.push_back(h);
+ }
+}
+
+/// One round of mixing each op's hash with the hashes of its in-block operand
+/// defs. The base hashes alone can be too coarse for period detection: after
+/// CSE a long window may consist of a handful of op kinds (e.g. identical
+/// gates on different qubits of a chain), and the gap statistics in
+/// findCandidates only see the period if some op hash recurs exactly once per
+/// window. Mixing distinguishes positions within a window by their local
+/// dataflow ancestry while staying equal across windows of a repeat. It also
+/// pollutes hashes near the start of the run (their ancestry reaches the
+/// prologue), so candidates are harvested both before and after refinement.
+void refineHashes(ArrayRef ops, const llvm::DenseMap &indexOf,
+ SmallVectorImpl &hashes)
+{
+ SmallVector prev(hashes.begin(), hashes.end());
+ for (const auto &[i, op] : llvm::enumerate(ops)) {
+ if (!isRerollableOp(op)) {
+ continue;
+ }
+ hash_code h = prev[i];
+ for (Value v : op->getOperands()) {
+ Operation *def = v.getDefiningOp();
+ auto it = def ? indexOf.find(def) : indexOf.end();
+ if (it != indexOf.end()) {
+ h = hash_combine(h, prev[it->second]);
+ }
+ }
+ hashes[i] = h;
+ }
+}
+
+struct Candidate {
+ size_t start; // index of the first op of the first window
+ size_t period; // ops per window
+ size_t count; // number of windows
+};
+
+/// Harvest candidate periods from the gap distribution between successive
+/// occurrences of equal hashes, then find maximal H[i] == H[i+p] runs.
+void findCandidates(ArrayRef hashes, unsigned minPeriod, unsigned maxPeriods,
+ SmallVectorImpl &out)
+{
+ size_t n = hashes.size();
+ llvm::DenseMap lastSeen;
+ llvm::DenseMap gapWeight;
+ for (size_t i = 0; i < n; ++i) {
+ auto [it, inserted] = lastSeen.try_emplace(hashes[i], i);
+ if (!inserted) {
+ size_t gap = i - it->second;
+ if (gap >= minPeriod) {
+ gapWeight[gap]++;
+ }
+ it->second = i;
+ }
+ }
+
+ SmallVector> periods(gapWeight.begin(), gapWeight.end());
+ // Favor periods that explain the most ops.
+ llvm::sort(periods, [](auto &a, auto &b) { return a.first * a.second > b.first * b.second; });
+ if (periods.size() > maxPeriods) {
+ periods.resize(maxPeriods);
+ }
+
+ for (auto &[p, weight] : periods) {
+ // A repeat of period p with >= kMinWindows windows produces at
+ // least (kMinWindows - 1) * p same-hash pairs at gap p; anything below
+ // that cannot yield an acceptable candidate.
+ if (weight < (kMinWindows - 1) * p) {
+ continue;
+ }
+ size_t runStart = SIZE_MAX;
+ for (size_t i = 0; i + p <= n; ++i) {
+ bool match = (i + p < n) && hashes[i] == hashes[i + p];
+ if (match && runStart == SIZE_MAX) {
+ runStart = i;
+ }
+ if (!match && runStart != SIZE_MAX) {
+ size_t runLen = i - runStart;
+ size_t count = runLen / p + 1;
+ if (count >= kMinWindows) {
+ out.push_back({runStart, p, count});
+ }
+ runStart = SIZE_MAX;
+ }
+ }
+ }
+}
+
+struct RerollPlan {
+ Candidate cand;
+ // Threaded slots: values flowing window -> next window, identified by
+ // (defining op position within the window, result number). Order defines
+ // the iter_args order.
+ SmallVector> slots;
+ // Initial value of each slot (operand of the first window).
+ SmallVector inits;
+ // Classification of each cross-window operand use:
+ // (window-relative op position, operand number) -> slot index.
+ llvm::DenseMap, unsigned> threadedUse;
+};
+
+/// Verify that the candidate's windows are isomorphic with dataflow limited to
+/// threaded slots + loop-invariant values, and build the reroll plan.
+std::optional verifyCandidate(ArrayRef ops,
+ const llvm::DenseMap &indexOf,
+ Candidate cand)
+{
+ size_t start = cand.start, p = cand.period, count = cand.count;
+ size_t end = start + count * p;
+ if (end > ops.size()) {
+ return std::nullopt;
+ }
+
+ RerollPlan plan;
+ plan.cand = cand;
+ llvm::DenseMap, unsigned> slotIndex; // (defPos,resNo) -> idx
+
+ auto getSlot = [&](unsigned defPos, unsigned resNo) -> unsigned {
+ auto [it, inserted] = slotIndex.try_emplace({defPos, resNo}, plan.slots.size());
+ if (inserted) {
+ plan.slots.push_back({defPos, resNo});
+ plan.inits.push_back(Value());
+ }
+ return it->second;
+ };
+
+ for (size_t w = 1; w < count; ++w) {
+ for (size_t j = 0; j < p; ++j) {
+ Operation *a = ops[start + (w - 1) * p + j];
+ Operation *b = ops[start + w * p + j];
+ if (!isRerollableOp(a) || !isRerollableOp(b)) {
+ return std::nullopt;
+ }
+ if (a->getName() != b->getName() || a->getAttrDictionary() != b->getAttrDictionary() ||
+ a->getResultTypes() != b->getResultTypes() ||
+ a->getNumOperands() != b->getNumOperands()) {
+ return std::nullopt;
+ }
+ for (unsigned t = 0; t < b->getNumOperands(); ++t) {
+ Value vb = b->getOperand(t);
+ Value va = a->getOperand(t);
+ Operation *defB = vb.getDefiningOp();
+ auto itB = defB ? indexOf.find(defB) : indexOf.end();
+ size_t ib = (itB != indexOf.end()) ? itB->second : SIZE_MAX;
+
+ if (ib != SIZE_MAX && ib >= start + w * p) {
+ // Within current window: the counterpart must reference the
+ // same relative position.
+ Operation *defA = va.getDefiningOp();
+ auto itA = defA ? indexOf.find(defA) : indexOf.end();
+ if (itA == indexOf.end() || itA->second + p != ib ||
+ cast(va).getResultNumber() !=
+ cast(vb).getResultNumber()) {
+ return std::nullopt;
+ }
+ }
+ else if (ib != SIZE_MAX && ib >= start + (w - 1) * p) {
+ // Threaded from the previous window.
+ unsigned defPos = ib - (start + (w - 1) * p);
+ unsigned resNo = cast(vb).getResultNumber();
+ unsigned slot = getSlot(defPos, resNo);
+ auto [uit, uinserted] = plan.threadedUse.try_emplace({(unsigned)j, t}, slot);
+ if (!uinserted && uit->second != slot) {
+ return std::nullopt;
+ }
+ // The counterpart operand must thread identically.
+ if (w == 1) {
+ // va is the init value; it must be loop-invariant w.r.t.
+ // the region (defined before it).
+ Operation *defA = va.getDefiningOp();
+ auto itA = defA ? indexOf.find(defA) : indexOf.end();
+ if (itA != indexOf.end() && itA->second >= start) {
+ return std::nullopt;
+ }
+ if (plan.inits[slot] && plan.inits[slot] != va) {
+ return std::nullopt;
+ }
+ plan.inits[slot] = va;
+ }
+ else {
+ Operation *defA = va.getDefiningOp();
+ auto itA = defA ? indexOf.find(defA) : indexOf.end();
+ if (itA == indexOf.end() || itA->second + p != ib ||
+ cast(va).getResultNumber() != resNo) {
+ return std::nullopt;
+ }
+ }
+ }
+ else {
+ // Loop-invariant: must be the exact same value, defined
+ // before the region.
+ if (va != vb) {
+ return std::nullopt;
+ }
+ if (ib != SIZE_MAX && ib >= start) {
+ return std::nullopt;
+ }
+ }
+ }
+ }
+ }
+
+ // Every slot must have an init.
+ for (Value init : plan.inits) {
+ if (!init) {
+ return std::nullopt;
+ }
+ }
+
+ // Uses of window results outside the allowed range:
+ // - windows 0..count-2: results may only be used inside their own window or
+ // the next one (threaded uses were verified above; any other use pattern
+ // is unsupported).
+ // - last window: external uses allowed only for threaded slots (they become
+ // loop results).
+ for (size_t w = 0; w < count; ++w) {
+ bool isLast = (w == count - 1);
+ for (size_t j = 0; j < p; ++j) {
+ Operation *op = ops[start + w * p + j];
+ for (OpResult res : op->getResults()) {
+ for (OpOperand &use : res.getUses()) {
+ Operation *user = use.getOwner();
+ auto uit = indexOf.find(user);
+ size_t ui = (uit != indexOf.end()) ? uit->second : SIZE_MAX;
+ bool inOwnWindow =
+ ui != SIZE_MAX && ui >= start + w * p && ui < start + (w + 1) * p;
+ bool inNextWindow = !isLast && ui != SIZE_MAX && ui >= start + (w + 1) * p &&
+ ui < start + (w + 2) * p;
+ if (inOwnWindow || inNextWindow) {
+ continue;
+ }
+ // External use.
+ if (!isLast) {
+ return std::nullopt;
+ }
+ if (!slotIndex.count({(unsigned)j, res.getResultNumber()})) {
+ return std::nullopt;
+ }
+ }
+ }
+ }
+ }
+
+ return plan;
+}
+
+/// Try to extend a verified candidate by whole windows to the left/right; the
+/// hash sequence misses the first window (its cross-window references point at
+/// the prologue, at different distances), so this recovers it. Each
+/// verification is linear in the candidate size, so the number of windows
+/// added per attempt grows geometrically (and resets on failure) to keep the
+/// total cost O(size * log(windows)) rather than quadratic.
+RerollPlan extendCandidate(ArrayRef ops,
+ const llvm::DenseMap &indexOf, RerollPlan plan)
+{
+ size_t step = 1;
+ while (plan.cand.start >= plan.cand.period * step) {
+ Candidate c = plan.cand;
+ c.start -= c.period * step;
+ c.count += step;
+ if (auto extended = verifyCandidate(ops, indexOf, c)) {
+ plan = *extended;
+ step *= 2;
+ }
+ else if (step == 1) {
+ break;
+ }
+ else {
+ step = 1;
+ }
+ }
+ step = 1;
+ while (true) {
+ Candidate c = plan.cand;
+ c.count += step;
+ if (c.start + c.count * c.period <= ops.size()) {
+ if (auto extended = verifyCandidate(ops, indexOf, c)) {
+ plan = *extended;
+ step *= 2;
+ continue;
+ }
+ }
+ if (step == 1) {
+ break;
+ }
+ step = 1;
+ }
+ return plan;
+}
+
+/// Replace the repeat with an scf.for and erase the original ops.
+void materialize(ArrayRef ops, const RerollPlan &plan)
+{
+ size_t start = plan.cand.start, p = plan.cand.period, count = plan.cand.count;
+ Operation *first = ops[start];
+ Location loc = first->getLoc();
+ OpBuilder builder(first);
+
+ Value lb = arith::ConstantIndexOp::create(builder, loc, 0);
+ Value ub = arith::ConstantIndexOp::create(builder, loc, count);
+ Value step = arith::ConstantIndexOp::create(builder, loc, 1);
+
+ auto forOp = scf::ForOp::create(builder, loc, lb, ub, step, plan.inits);
+ Block *body = forOp.getBody();
+ builder.setInsertionPointToStart(body);
+
+ // Clone window 0 as the loop body, remapping operands per classification.
+ SmallVector cloned(p);
+ IRMapping mapping; // within-window result mapping
+ for (size_t j = 0; j < p; ++j) {
+ Operation *proto = ops[start + j];
+ Operation *clone = proto->cloneWithoutRegions(mapping);
+ // Fix up operands that are threaded from the previous iteration: the
+ // prototype (window 0) uses the init values there.
+ for (unsigned t = 0; t < clone->getNumOperands(); ++t) {
+ auto it = plan.threadedUse.find({(unsigned)j, t});
+ if (it != plan.threadedUse.end()) {
+ clone->setOperand(t, forOp.getRegionIterArg(it->second));
+ }
+ }
+ builder.insert(clone);
+ cloned[j] = clone;
+ }
+ SmallVector yields;
+ for (auto [defPos, resNo] : plan.slots) {
+ yields.push_back(cloned[defPos]->getResult(resNo));
+ }
+ scf::YieldOp::create(builder, loc, yields);
+
+ // Rewire external uses of the last window's results to the loop results.
+ for (const auto &[slotIdx, slot] : llvm::enumerate(plan.slots)) {
+ auto [defPos, resNo] = slot;
+ Operation *lastOp = ops[start + (count - 1) * p + defPos];
+ lastOp->getResult(resNo).replaceAllUsesWith(forOp.getResult(slotIdx));
+ }
+
+ // Erase the original ops, last first (uses before defs). Verification
+ // guarantees no surviving uses; erase() asserts use_empty, so a
+ // verification bug fails loudly here instead of producing invalid IR.
+ for (size_t i = start + count * p; i-- > start;) {
+ assert(ops[i]->use_empty() && "rerolled op still has uses; verification is unsound");
+ ops[i]->erase();
+ }
+}
+
+bool processBlock(Block &block, unsigned minPeriod, unsigned minSavings)
+{
+ SmallVector ops;
+ for (Operation &op : block.without_terminator()) {
+ ops.push_back(&op);
+ }
+ if (ops.size() < 2 * minPeriod) {
+ return false;
+ }
+
+ llvm::DenseMap indexOf;
+ for (const auto &[i, op] : llvm::enumerate(ops)) {
+ indexOf[op] = i;
+ }
+
+ // Harvest candidates at every hash refinement depth: shallow hashes see
+ // repeats whose windows contain few distinct op kinds only as noise, deep
+ // hashes lose the first windows of a run to prologue ancestry. Duplicated
+ // candidates are cheap (verification dedups via the overlap check).
+ SmallVector hashes;
+ computeHashes(ops, indexOf, hashes);
+ SmallVector candidates;
+ findCandidates(hashes, minPeriod, kMaxPeriodsPerBlock, candidates);
+ for (unsigned depth = 0; depth < kHashDepth; ++depth) {
+ refineHashes(ops, indexOf, hashes);
+ findCandidates(hashes, minPeriod, kMaxPeriodsPerBlock, candidates);
+ }
+
+ // Verify, extend, and pick non-overlapping plans greedily by savings.
+ SmallVector plans;
+ for (Candidate cand : candidates) {
+ auto plan = verifyCandidate(ops, indexOf, cand);
+ if (!plan) {
+ continue;
+ }
+ *plan = extendCandidate(ops, indexOf, *plan);
+ if ((plan->cand.count - 1) * plan->cand.period >= minSavings) {
+ plans.push_back(std::move(*plan));
+ }
+ }
+ llvm::sort(plans, [](const RerollPlan &a, const RerollPlan &b) {
+ return (a.cand.count - 1) * a.cand.period > (b.cand.count - 1) * b.cand.period;
+ });
+
+ SmallVector> used;
+ SmallVector accepted;
+ for (const RerollPlan &plan : plans) {
+ size_t s = plan.cand.start, e = s + plan.cand.count * plan.cand.period;
+ bool overlaps =
+ llvm::any_of(used, [&](auto range) { return s < range.second && range.first < e; });
+ if (!overlaps) {
+ used.push_back({s, e});
+ accepted.push_back(&plan);
+ }
+ }
+
+ // Materialize from the highest start index down so op indices of pending
+ // plans remain valid.
+ llvm::sort(accepted, [](const RerollPlan *a, const RerollPlan *b) {
+ return a->cand.start > b->cand.start;
+ });
+ for (const RerollPlan *plan : accepted) {
+ LLVM_DEBUG(dbgs() << "rerolling: start=" << plan->cand.start
+ << " period=" << plan->cand.period << " count=" << plan->cand.count
+ << " slots=" << plan->slots.size() << "\n");
+ materialize(ops, *plan);
+ }
+ return !accepted.empty();
+}
+
+} // namespace
+
+namespace catalyst {
+
+// GEN_PASS_DECL is needed in addition to GEN_PASS_DEF to declare the
+// RerollLoopsPassOptions struct that the generated base class references
+// (only passes with options need this).
+#define GEN_PASS_DECL_REROLLLOOPSPASS
+#define GEN_PASS_DEF_REROLLLOOPSPASS
+#include "Catalyst/Transforms/Passes.h.inc"
+
+struct RerollLoopsPass : public impl::RerollLoopsPassBase {
+ using impl::RerollLoopsPassBase::RerollLoopsPassBase;
+
+ void runOnOperation() override
+ {
+ bool changed = true;
+ unsigned rounds = 0;
+ while (changed && rounds++ < kMaxRounds) {
+ changed = false;
+ SmallVector blocks;
+ getOperation()->walk([&](Block *block) { blocks.push_back(block); });
+ for (Block *block : blocks) {
+ changed |= processBlock(*block, minPeriod, minSavings);
+ }
+ }
+ }
+};
+
+} // namespace catalyst
diff --git a/mlir/test/Catalyst/RerollLoopsTest.mlir b/mlir/test/Catalyst/RerollLoopsTest.mlir
new file mode 100644
index 0000000000..5fe1ee9ad0
--- /dev/null
+++ b/mlir/test/Catalyst/RerollLoopsTest.mlir
@@ -0,0 +1,127 @@
+// 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.
+
+// RUN: quantum-opt --reroll-loops="min-period=2 min-savings=4" --split-input-file %s | FileCheck %s
+
+// A scalar chain of alternating ops rerolls into an scf.for threading one
+// value through all eight iterations.
+
+// CHECK-LABEL: @scalar_chain
+// CHECK: %[[FOR:.+]] = scf.for {{.*}} iter_args(%[[IT:.+]] = %arg0) -> (f64)
+// CHECK: %[[A:.+]] = arith.addf %[[IT]],
+// CHECK: %[[M:.+]] = arith.mulf %[[A]],
+// CHECK: scf.yield %[[M]] : f64
+// CHECK: return %[[FOR]]
+func.func @scalar_chain(%arg0: f64, %c: f64) -> f64 {
+ %0 = arith.addf %arg0, %c : f64
+ %1 = arith.mulf %0, %c : f64
+ %2 = arith.addf %1, %c : f64
+ %3 = arith.mulf %2, %c : f64
+ %4 = arith.addf %3, %c : f64
+ %5 = arith.mulf %4, %c : f64
+ %6 = arith.addf %5, %c : f64
+ %7 = arith.mulf %6, %c : f64
+ %8 = arith.addf %7, %c : f64
+ %9 = arith.mulf %8, %c : f64
+ %10 = arith.addf %9, %c : f64
+ %11 = arith.mulf %10, %c : f64
+ %12 = arith.addf %11, %c : f64
+ %13 = arith.mulf %12, %c : f64
+ %14 = arith.addf %13, %c : f64
+ %15 = arith.mulf %14, %c : f64
+ return %15 : f64
+}
+
+// -----
+
+// A repeated gate sequence threading two qubits rerolls with both qubit
+// values as iter_args; the rotation angle is loop-invariant.
+
+// CHECK-LABEL: @gate_sequence
+// CHECK: quantum.alloc
+// CHECK: %[[FOR:.+]]:2 = scf.for {{.*}} iter_args(%[[Q0:.+]] = %{{.+}}, %[[Q1:.+]] = %{{.+}}) -> (!quantum.bit, !quantum.bit)
+// CHECK: %[[H:.+]] = quantum.custom "Hadamard"() %[[Q0]]
+// CHECK: %[[RZ:.+]] = quantum.custom "RZ"(%{{.+}}) %[[Q1]]
+// CHECK: %[[CNOT:.+]]:2 = quantum.custom "CNOT"() %[[H]], %[[RZ]]
+// CHECK: scf.yield %[[CNOT]]#0, %[[CNOT]]#1
+// CHECK: quantum.insert %{{.+}}[ 0], %[[FOR]]#0
+func.func @gate_sequence(%theta: f64) -> !quantum.reg {
+ %r0 = quantum.alloc( 2) : !quantum.reg
+ %q0 = quantum.extract %r0[ 0] : !quantum.reg -> !quantum.bit
+ %q1 = quantum.extract %r0[ 1] : !quantum.reg -> !quantum.bit
+
+ %h0 = quantum.custom "Hadamard"() %q0 : !quantum.bit
+ %z0 = quantum.custom "RZ"(%theta) %q1 : !quantum.bit
+ %c0:2 = quantum.custom "CNOT"() %h0, %z0 : !quantum.bit, !quantum.bit
+
+ %h1 = quantum.custom "Hadamard"() %c0#0 : !quantum.bit
+ %z1 = quantum.custom "RZ"(%theta) %c0#1 : !quantum.bit
+ %c1:2 = quantum.custom "CNOT"() %h1, %z1 : !quantum.bit, !quantum.bit
+
+ %h2 = quantum.custom "Hadamard"() %c1#0 : !quantum.bit
+ %z2 = quantum.custom "RZ"(%theta) %c1#1 : !quantum.bit
+ %c2:2 = quantum.custom "CNOT"() %h2, %z2 : !quantum.bit, !quantum.bit
+
+ %h3 = quantum.custom "Hadamard"() %c2#0 : !quantum.bit
+ %z3 = quantum.custom "RZ"(%theta) %c2#1 : !quantum.bit
+ %c3:2 = quantum.custom "CNOT"() %h3, %z3 : !quantum.bit, !quantum.bit
+
+ %r1 = quantum.insert %r0[ 0], %c3#0 : !quantum.reg, !quantum.bit
+ %r2 = quantum.insert %r1[ 1], %c3#1 : !quantum.reg, !quantum.bit
+ return %r2 : !quantum.reg
+}
+
+// -----
+
+// Iterations with *different* angle values must not be rerolled: the varying
+// operand fails loop-invariance verification.
+
+// CHECK-LABEL: @varying_angles
+// CHECK-NOT: scf.for
+func.func @varying_angles(%t0: f64, %t1: f64, %t2: f64, %t3: f64) -> !quantum.reg {
+ %r0 = quantum.alloc( 1) : !quantum.reg
+ %q0 = quantum.extract %r0[ 0] : !quantum.reg -> !quantum.bit
+ %a = quantum.custom "RZ"(%t0) %q0 : !quantum.bit
+ %b = quantum.custom "RX"(%t0) %a : !quantum.bit
+ %c = quantum.custom "RZ"(%t1) %b : !quantum.bit
+ %d = quantum.custom "RX"(%t1) %c : !quantum.bit
+ %e = quantum.custom "RZ"(%t2) %d : !quantum.bit
+ %f = quantum.custom "RX"(%t2) %e : !quantum.bit
+ %g = quantum.custom "RZ"(%t3) %f : !quantum.bit
+ %h = quantum.custom "RX"(%t3) %g : !quantum.bit
+ %r1 = quantum.insert %r0[ 0], %h : !quantum.reg, !quantum.bit
+ return %r1 : !quantum.reg
+}
+
+// -----
+
+// Results used outside the repeat (other than the final threaded values) block
+// rerolling: here every intermediate feeds the final sum.
+
+// CHECK-LABEL: @external_uses
+// CHECK-NOT: scf.for
+func.func @external_uses(%arg0: f64, %c: f64) -> f64 {
+ %0 = arith.addf %arg0, %c : f64
+ %1 = arith.mulf %0, %c : f64
+ %2 = arith.addf %1, %c : f64
+ %3 = arith.mulf %2, %c : f64
+ %4 = arith.addf %3, %c : f64
+ %5 = arith.mulf %4, %c : f64
+ %6 = arith.addf %5, %c : f64
+ %7 = arith.mulf %6, %c : f64
+ %s0 = arith.addf %1, %3 : f64
+ %s1 = arith.addf %s0, %5 : f64
+ %s2 = arith.addf %s1, %7 : f64
+ return %s2 : f64
+}