From c8fde8c5d39f9dc338ffa9a866f9b6a49f9d4409 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Tue, 2 Dec 2025 15:12:59 -0500 Subject: [PATCH 01/21] initial implementation and tests --- .../Transforms/MergeRotationsPatterns.cpp | 61 +++++++++++++++++++ .../lib/Quantum/Transforms/merge_rotation.cpp | 2 + mlir/test/Quantum/MergeRotationsTest.mlir | 48 ++++----------- 3 files changed, 75 insertions(+), 36 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index b66bb62fcf..b98b5d8a8f 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -396,6 +396,66 @@ struct MergePPRRewritePattern : public OpRewritePattern { } }; +struct MergePPRArbitraryRewritePattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(PPRotationArbitraryOp op, + PatternRewriter &rewriter) const override + { + ValueRange inQubits = op.getInQubits(); + auto definingOp = inQubits[0].getDefiningOp(); + + if (!definingOp) { + return failure(); + } + + auto prevOp = dyn_cast(definingOp); + + if (!prevOp) { + return failure(); + } + + // verify that prevOp agrees on all qubits, not just the first + for (auto qubit : inQubits) { + if (qubit.getDefiningOp() != prevOp) { + return failure(); + } + } + + // check same pauli strings + if (op.getPauliProduct() != prevOp.getPauliProduct()) { + return failure(); + } + + // check same conditionals + if (op.getCondition() != prevOp.getCondition()) { + return failure(); + } + + auto opRotation = op.getArbitraryAngle(); + auto prevOpRotation = prevOp.getArbitraryAngle(); + + if (!opRotation || !prevOpRotation) { + return failure(); + } + + auto loc = op.getLoc(); + // replace references to current op with prevOp + // create merged op + mlir::Value newAngleOp = + rewriter.create(loc, opRotation, prevOpRotation).getResult(); + + auto mergeOp = rewriter.create( + loc, op.getOutQubits().getTypes(), op.getPauliProduct(), newAngleOp, + prevOp.getInQubits(), op.getCondition()); + + rewriter.replaceOp(op, mergeOp); + rewriter.eraseOp(prevOp); + + return success(); + } +}; + struct MergeMultiRZRewritePattern : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -442,6 +502,7 @@ void populateMergeRotationsPatterns(RewritePatternSet &patterns) patterns.add>(patterns.getContext(), 1); patterns.add(patterns.getContext(), 1); patterns.add(patterns.getContext(), 1); + patterns.add(patterns.getContext(), 1); } } // namespace quantum diff --git a/mlir/lib/Quantum/Transforms/merge_rotation.cpp b/mlir/lib/Quantum/Transforms/merge_rotation.cpp index 3bd79a3b87..c9db18fb14 100644 --- a/mlir/lib/Quantum/Transforms/merge_rotation.cpp +++ b/mlir/lib/Quantum/Transforms/merge_rotation.cpp @@ -55,6 +55,8 @@ struct MergeRotationsPass : impl::MergeRotationsPassBase { &getContext()); catalyst::qec::PPRotationOp::getCanonicalizationPatterns(patternsCanonicalization, &getContext()); + catalyst::qec::PPRotationArbitraryOp::getCanonicalizationPatterns(patternsCanonicalization, + &getContext()); if (failed(applyPatternsGreedily(module, std::move(patternsCanonicalization)))) { return signalPassFailure(); } diff --git a/mlir/test/Quantum/MergeRotationsTest.mlir b/mlir/test/Quantum/MergeRotationsTest.mlir index 9fced31b6a..a1343795da 100644 --- a/mlir/test/Quantum/MergeRotationsTest.mlir +++ b/mlir/test/Quantum/MergeRotationsTest.mlir @@ -902,42 +902,18 @@ func.func public @half_compatible_qubits(%q1: !quantum.bit, %q2: !quantum.bit, % // ----- -// merge correctly when conditionals are present - -// CHECK-LABEL: merge_conditionals -func.func public @merge_conditionals(%q1: !quantum.bit, %q2: !quantum.bit, %arg0: i1) { - // CHECK: qec.ppr ["X", "Z"](4) {{%.+}}, {{%.+}} cond({{%.+}}) - // CHECK-NOT: qec.ppr ["X", "Z"](8) - %0, %1 = qec.ppr ["X", "Z"](8) %q1, %q2 cond(%arg0): !quantum.bit, !quantum.bit - %2, %3 = qec.ppr ["X", "Z"](8) %0, %1 cond(%arg0): !quantum.bit, !quantum.bit - func.return -} - -// ----- - -// don't merge compatible Pauli words when conditionals differ - -// CHECK-LABEL: dont_merge_incompatible_conditionals -func.func public @dont_merge_incompatible_conditionals(%q1: !quantum.bit, %q2: !quantum.bit, %arg0: i1, %arg1: i1) { - // CHECK-NOT: qec.ppr ["X", "Z"](4) - // CHECK: [[in:%.+]]:2 = qec.ppr ["X", "Z"](8) - // CHECK: qec.ppr ["X", "Z"](8) [[in]]#0, [[in]]#1 - %0:2 = qec.ppr ["X", "Z"](8) %q1, %q2 cond(%arg0): !quantum.bit, !quantum.bit - %1:2 = qec.ppr ["X", "Z"](8) %0#0, %0#1 cond(%arg1): !quantum.bit, !quantum.bit - func.return -} - -// ----- - -// don't merge PPRs with conditional PPRs - -// CHECK-LABEL: dont_merge_conditionals -func.func public @dont_merge_conditionals(%q1: !quantum.bit, %q2: !quantum.bit, %arg0: i1) { - // CHECK-NOT: qec.ppr ["X", "Z"](4) - // CHECK: [[in:%.+]]:2 = qec.ppr ["X", "Z"](8) - // CHECK: qec.ppr ["X", "Z"](8) [[in]]#0, [[in]]#1 - %0:2 = qec.ppr ["X", "Z"](8) %q1, %q2: !quantum.bit, !quantum.bit - %1:2 = qec.ppr ["X", "Z"](8) %0#0, %0#1 cond(%arg0): !quantum.bit, !quantum.bit +// Arbitrary Angle PPR (AAPPR) Tests + +// simple AAPPR merge + +// CHECK-LABEL: merge_Y +func.func public @merge_Y(%q1: !quantum.bit) { + // CHECK: [[angle:%.+]] = arith.addf [[cst1]], [[cst2]] + // CHECK: qec.ppr.arbitrary ["Y"]([[angle]]) + %0 = arith.constant 0.9 : f64 + %1 = arith.constant 0.4 : f64 + %2 = qec.ppr.arbitrary ["Y"](%0) %q1: !quantum.bit + %3 = qec.ppr.arbitrary ["Y"](%1) %2: !quantum.bit func.return } From 9761ddd9565f2e4aac58f7cca5d4ff87f9202e51 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 3 Dec 2025 11:07:34 -0500 Subject: [PATCH 02/21] added check for differently ordered but equivalent pauli string+qubit combinations --- .../Transforms/MergeRotationsPatterns.cpp | 22 ++- mlir/test/Quantum/MergeRotationsTest.mlir | 144 +++++++++++++++++- 2 files changed, 160 insertions(+), 6 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index b98b5d8a8f..b03fa8be57 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -322,6 +322,24 @@ struct MergeRotationsRewritePattern : public OpRewritePattern { } }; +bool matchingQubitsAndPaulis(PPRotationArbitraryOp op, PPRotationArbitraryOp parentOp) +{ + auto pauliProduct = op.getPauliProduct(); + auto parentPauliProduct = parentOp.getPauliProduct(); + + for (auto [i, qubit] : llvm::enumerate(op.getInQubits())) { + for (auto [j, parentQubit] : llvm::enumerate(parentOp.getOutQubits())) { + if (qubit == parentQubit) { + if (pauliProduct[i] != parentPauliProduct[j]) { + return false; + } + } + } + } + + return true; +} + struct MergePPRRewritePattern : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -422,8 +440,8 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern !quantum.bit { + // CHECK-DAG: [[angle:%.+]] = arith.addf + // CHECK-DAG: quantum.custom + // CHECK-DAG: qec.ppr.arbitrary ["X"]([[angle]]) + %2 = qec.ppr.arbitrary ["X"](%0) %q1: !quantum.bit + %3 = quantum.custom "Hadamard"() %q2: !quantum.bit + %4 = qec.ppr.arbitrary ["X"](%1) %2: !quantum.bit + func.return %3: !quantum.bit +} + +// ----- + +// don't merge through other operations + +// CHECK-LABEL: mixed_operations +func.func public @mixed_operations(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64) { + // CHECK-NOT: arith.addf + // CHECK: qec.ppr.arbitrary ["Z", "X"] + // CHECK: quantum.custom + // CHECK: qec.ppr.arbitrary ["Z", "X"] + %2:2 = qec.ppr.arbitrary ["Z", "X"](%0) %q1, %q2: !quantum.bit, !quantum.bit + %3 = quantum.custom "Hadamard"() %2#1: !quantum.bit + %5:2 = qec.ppr.arbitrary ["Z", "X"](%1) %2#0, %3: !quantum.bit, !quantum.bit + func.return +} + +// ----- + +// don't merge if only one qubit matches + +// CHECK-LABEL: half_compatible_qubits +func.func public @half_compatible_qubits(%q1: !quantum.bit, %q2: !quantum.bit, %q3: !quantum.bit, %0: f64, %1: f64) { + // CHECK: qec.ppr.arbitrary ["X", "Z"] + %2:2 = qec.ppr.arbitrary ["X", "Z"](%0) %q1, %q2: !quantum.bit, !quantum.bit + %3:2 = qec.ppr.arbitrary ["X", "Z"](%1) %q3, %2#1 : !quantum.bit, !quantum.bit + func.return +} + +// ----- + +// re-arranging qubits is ok as long as the pauli words are re-arranged too + +// CHECK-LABEL: mix_and_match +func.func public @mix_and_match(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64) { + // CHECK: [[angle:%.+]] = arith.addf + // CHECK: qec.ppr.arbitrary [{{.+}}]([[angle]]) + %2:2 = qec.ppr.arbitrary ["Z", "Y"](%0) %q1, %q2: !quantum.bit, !quantum.bit + %3:2 = qec.ppr.arbitrary ["Y", "Z"](%1) %2#1, %2#0: !quantum.bit, !quantum.bit + func.return +} + +// ----- + +// re-arranging qubits without re-arranging the pauli word is NOT okay + +// CHECK-LABEL: mix_dont_match +func.func public @mix_dont_match(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64) { + // CHECK: qec.ppr.arbitrary ["Y", "X"] + // CHECK: qec.ppr.arbitrary ["Y", "X"] + %2:2 = qec.ppr.arbitrary ["Y", "X"](%0) %q1, %q2: !quantum.bit, !quantum.bit + %3:2 = qec.ppr.arbitrary ["Y", "X"](%1) %2#1, %2#0: !quantum.bit, !quantum.bit + func.return +} From c25cff6d7fc9b89b44a112cd937133ea062be9f1 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 3 Dec 2025 11:36:49 -0500 Subject: [PATCH 03/21] added integration test --- .../pytest/test_peephole_optimizations.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/frontend/test/pytest/test_peephole_optimizations.py b/frontend/test/pytest/test_peephole_optimizations.py index 0a17a0e06a..1d297ff614 100644 --- a/frontend/test/pytest/test_peephole_optimizations.py +++ b/frontend/test/pytest/test_peephole_optimizations.py @@ -518,6 +518,33 @@ def circuit(): assert 'qec.ppr ["X", "Y", "Z"](2)' in ir_opt +@pytest.mark.usefixtures("use_capture") +def test_merge_rotation_arbitrary_angle_ppr(): + """Test that the merge_rotation pass correctly merges arbtirary angle PPRs.""" + + my_pipeline = [("pipe", ["quantum-compilation-stage"])] + + @qml.qjit(pipelines=my_pipeline, target="mlir") + def test_merge_rotation_ppr_workflow(): + @qml.transforms.merge_rotations # have to use qml to be capture-compatible + @qml.qnode(qml.device("lightning.qubit", wires=2)) + def circuit(x, y): + qml.PauliRot(x, pauli_word="ZY", wires=[0, 1]) + qml.PauliRot(y, pauli_word="ZY", wires=[0, 1]) + + return circuit(2.6, 0.3) + + ir = test_merge_rotation_ppr_workflow.mlir + ir_opt = test_merge_rotation_ppr_workflow.mlir_opt + + assert 'transform.apply_registered_pass "merge-rotations"' in ir + assert "qec.ppr.arbitrary" in ir + assert "arith.addf" not in ir + + assert "arith.addf" in ir_opt + assert 'qec.ppr.arbitrary ["Z", "Y"]' in ir_opt + + def test_clifford_to_ppm(): pipe = [("pipe", ["quantum-compilation-stage"])] From 0585074a3ca309bde0efa6ca7b9e5d9871506b16 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 3 Dec 2025 11:37:00 -0500 Subject: [PATCH 04/21] changelog --- doc/releases/changelog-dev.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index f4def776b0..cfa4ed59e1 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -306,13 +306,14 @@ gradient dialect and the `lower-gradients` compilation stage. [(#2241)](https://github.com/PennyLaneAI/catalyst/pull/2241) - * Added support for PPRs to the :func:`~.passes.merge_rotations` pass to merge PPRs with - equivalent angles, and cancelling of PPRs with opposite angles, or angles - that sum to identity. Also supports conditions on PPRs, merging when conditions are - identical and not merging otherwise. + * Added support for PPRs and arbitrary angle PPRs to the :func:`~.passes.merge_rotations` pass. + This pass now merges PPRs with equivalent angles, and cancels PPRs with opposite angles, or + angles that sum to identity when the angles are known. The pass also supports conditions on PPRs, + merging when conditions are identical and not merging otherwise. [(#2224)](https://github.com/PennyLaneAI/catalyst/pull/2224) [(#2245)](https://github.com/PennyLaneAI/catalyst/pull/2245) [(#2254)](https://github.com/PennyLaneAI/catalyst/pull/2254) + [(#2258)](https://github.com/PennyLaneAI/catalyst/pull/2258)

Documentation 📝

From 9bbf993602751e8c73538c2fc6f47b2d49c8c4dc Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 3 Dec 2025 14:04:27 -0500 Subject: [PATCH 05/21] improved complexity of compatibility helper function --- .../Transforms/MergeRotationsPatterns.cpp | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index b03fa8be57..3aac39c8e1 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -324,16 +324,16 @@ struct MergeRotationsRewritePattern : public OpRewritePattern { bool matchingQubitsAndPaulis(PPRotationArbitraryOp op, PPRotationArbitraryOp parentOp) { - auto pauliProduct = op.getPauliProduct(); - auto parentPauliProduct = parentOp.getPauliProduct(); - - for (auto [i, qubit] : llvm::enumerate(op.getInQubits())) { - for (auto [j, parentQubit] : llvm::enumerate(parentOp.getOutQubits())) { - if (qubit == parentQubit) { - if (pauliProduct[i] != parentPauliProduct[j]) { - return false; - } - } + // construct map + std::unordered_map qubit_to_pauli; + for (auto [qubit, pauli] : llvm::zip(parentOp.getOutQubits(), parentOp.getPauliProduct())) { + qubit_to_pauli[hash_value(qubit)] = pauli; + } + + // check pairings + for (auto [qubit, pauli] : llvm::zip(op.getInQubits(), op.getPauliProduct())) { + if (qubit_to_pauli[hash_value(qubit)] != pauli) { + return false; } } From 5ec8f17494f73b604c9893827f588f162b2c8fc8 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 3 Dec 2025 14:38:53 -0500 Subject: [PATCH 06/21] added tests for conditions --- mlir/test/Quantum/MergeRotationsTest.mlir | 87 ++++++++++++++++------- 1 file changed, 63 insertions(+), 24 deletions(-) diff --git a/mlir/test/Quantum/MergeRotationsTest.mlir b/mlir/test/Quantum/MergeRotationsTest.mlir index 5bb1ebdeab..53dbbd7807 100644 --- a/mlir/test/Quantum/MergeRotationsTest.mlir +++ b/mlir/test/Quantum/MergeRotationsTest.mlir @@ -902,15 +902,15 @@ func.func public @half_compatible_qubits(%q1: !quantum.bit, %q2: !quantum.bit, % // ----- -// Arbitrary Angle PPR (AAPPR) Tests +// Arbitrary Angle PPR Tests -// simple AAPPR merge +// simple merge // CHECK-LABEL: merge_Y -func.func public @merge_Y(%q1: !quantum.bit, %0: f64, %1: f64) { +func.func public @merge_Y(%q0: !quantum.bit, %0: f64, %1: f64) { // CHECK: [[angle:%.+]] = arith.addf // CHECK: qec.ppr.arbitrary ["Y"]([[angle]]) - %2 = qec.ppr.arbitrary ["Y"](%0) %q1: !quantum.bit + %2 = qec.ppr.arbitrary ["Y"](%0) %q0: !quantum.bit %3 = qec.ppr.arbitrary ["Y"](%1) %2: !quantum.bit func.return } @@ -920,12 +920,12 @@ func.func public @merge_Y(%q1: !quantum.bit, %0: f64, %1: f64) { // multiple merges // CHECK-LABEL: merge_multi_Z -func.func public @merge_multi_Z(%q1: !quantum.bit, %0: f64, %1: f64, %2: f64) { +func.func public @merge_multi_Z(%q0: !quantum.bit, %0: f64, %1: f64, %2: f64) { // CHECK: [[angle:%.+]] = arith.addf // CHECK: [[angle2:%.+]] = arith.addf // CHECK: qec.ppr.arbitrary ["Z"]([[angle2]]) // CHECK-NOT: qec.ppr.arbitrary - %3 = qec.ppr.arbitrary ["Z"](%0) %q1: !quantum.bit + %3 = qec.ppr.arbitrary ["Z"](%0) %q0: !quantum.bit %4 = qec.ppr.arbitrary ["Z"](%1) %3: !quantum.bit %5 = qec.ppr.arbitrary ["Z"](%2) %4: !quantum.bit func.return @@ -936,7 +936,7 @@ func.func public @merge_multi_Z(%q1: !quantum.bit, %0: f64, %1: f64, %2: f64) { // not merging when incompatible // CHECK-LABEL: dont_merge -func.func public @dont_merge(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64, %2: f64, %3: f64, %4: f64, %5: f64) { +func.func public @dont_merge(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, %1: f64, %2: f64, %3: f64, %4: f64, %5: f64) { // CHECK-NOT: arith.addf // CHECK: qec.ppr.arbitrary ["Z", "X"] // CHECK: qec.ppr.arbitrary ["Y", "X"] @@ -944,7 +944,7 @@ func.func public @dont_merge(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: // CHECK: qec.ppr.arbitrary ["X", "Z"] // CHECK: qec.ppr.arbitrary ["X", "Y"] // CHECK: qec.ppr.arbitrary ["Z", "Y"] - %6:2 = qec.ppr.arbitrary ["Z", "X"](%0) %q1, %q2: !quantum.bit, !quantum.bit + %6:2 = qec.ppr.arbitrary ["Z", "X"](%0) %q0, %q1: !quantum.bit, !quantum.bit %7:2 = qec.ppr.arbitrary ["Y", "X"](%1) %6#0, %6#1: !quantum.bit, !quantum.bit %8:2 = qec.ppr.arbitrary ["Y", "Z"](%2) %7#0, %7#1: !quantum.bit, !quantum.bit %9:2 = qec.ppr.arbitrary ["X", "Z"](%3) %8#0, %8#1: !quantum.bit, !quantum.bit @@ -958,12 +958,12 @@ func.func public @dont_merge(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: // updating references // CHECK-LABEL: merge_correct_references -func.func public @merge_correct_references(%q1: !quantum.bit, %0: f64, %1: f64, %2: f64, %3: f64) { +func.func public @merge_correct_references(%q0: !quantum.bit, %0: f64, %1: f64, %2: f64, %3: f64) { // CHECK-DAG: [[angle:%.+]] = arith.addf // CHECK-DAG: [[in:%.+]] = qec.ppr.arbitrary ["X"] // CHECK: [[out:%.+]] = qec.ppr.arbitrary ["Z"]([[angle]]) [[in]] // CHECK: qec.ppr.arbitrary ["Y"]({{%.+}}) [[out]] - %4 = qec.ppr.arbitrary ["X"](%0) %q1: !quantum.bit + %4 = qec.ppr.arbitrary ["X"](%0) %q0: !quantum.bit %5 = qec.ppr.arbitrary ["Z"](%1) %4: !quantum.bit %6 = qec.ppr.arbitrary ["Z"](%2) %5: !quantum.bit %7 = qec.ppr.arbitrary ["Y"](%3) %6: !quantum.bit @@ -975,11 +975,11 @@ func.func public @merge_correct_references(%q1: !quantum.bit, %0: f64, %1: f64, // multi-qubit merge // CHECK-LABEL: merge_multi_XZY -func.func public @merge_multi_XZY(%q1: !quantum.bit, %q2: !quantum.bit, %q3: !quantum.bit, %0: f64, %1: f64, %2: f64) { +func.func public @merge_multi_XZY(%q0: !quantum.bit, %q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64, %2: f64) { // CHECK: [[angle1:%.+]] = arith.addf // CHECK: [[angle2:%.+]] = arith.addf // CHECK: qec.ppr.arbitrary ["X", "Z", "Y"]([[angle2]]) - %3:3 = qec.ppr.arbitrary ["X", "Z", "Y"](%0) %q1, %q2, %q3: !quantum.bit, !quantum.bit, !quantum.bit + %3:3 = qec.ppr.arbitrary ["X", "Z", "Y"](%0) %q0, %q1, %q2: !quantum.bit, !quantum.bit, !quantum.bit %4:3 = qec.ppr.arbitrary ["X", "Z", "Y"](%1) %3#0, %3#1, %3#2: !quantum.bit, !quantum.bit, !quantum.bit %5:3 = qec.ppr.arbitrary ["X", "Z", "Y"](%2) %4#0, %4#1, %4#2: !quantum.bit, !quantum.bit, !quantum.bit func.return @@ -990,12 +990,12 @@ func.func public @merge_multi_XZY(%q1: !quantum.bit, %q2: !quantum.bit, %q3: !qu // merge through other ops // CHECK-LABEL: merge_through -func.func public @merge_through(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64) -> !quantum.bit { +func.func public @merge_through(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, %1: f64) -> !quantum.bit { // CHECK-DAG: [[angle:%.+]] = arith.addf // CHECK-DAG: quantum.custom // CHECK-DAG: qec.ppr.arbitrary ["X"]([[angle]]) - %2 = qec.ppr.arbitrary ["X"](%0) %q1: !quantum.bit - %3 = quantum.custom "Hadamard"() %q2: !quantum.bit + %2 = qec.ppr.arbitrary ["X"](%0) %q0: !quantum.bit + %3 = quantum.custom "Hadamard"() %q1: !quantum.bit %4 = qec.ppr.arbitrary ["X"](%1) %2: !quantum.bit func.return %3: !quantum.bit } @@ -1005,12 +1005,12 @@ func.func public @merge_through(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, % // don't merge through other operations // CHECK-LABEL: mixed_operations -func.func public @mixed_operations(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64) { +func.func public @mixed_operations(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, %1: f64) { // CHECK-NOT: arith.addf // CHECK: qec.ppr.arbitrary ["Z", "X"] // CHECK: quantum.custom // CHECK: qec.ppr.arbitrary ["Z", "X"] - %2:2 = qec.ppr.arbitrary ["Z", "X"](%0) %q1, %q2: !quantum.bit, !quantum.bit + %2:2 = qec.ppr.arbitrary ["Z", "X"](%0) %q0, %q1: !quantum.bit, !quantum.bit %3 = quantum.custom "Hadamard"() %2#1: !quantum.bit %5:2 = qec.ppr.arbitrary ["Z", "X"](%1) %2#0, %3: !quantum.bit, !quantum.bit func.return @@ -1021,10 +1021,10 @@ func.func public @mixed_operations(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64 // don't merge if only one qubit matches // CHECK-LABEL: half_compatible_qubits -func.func public @half_compatible_qubits(%q1: !quantum.bit, %q2: !quantum.bit, %q3: !quantum.bit, %0: f64, %1: f64) { +func.func public @half_compatible_qubits(%q0: !quantum.bit, %q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64) { // CHECK: qec.ppr.arbitrary ["X", "Z"] - %2:2 = qec.ppr.arbitrary ["X", "Z"](%0) %q1, %q2: !quantum.bit, !quantum.bit - %3:2 = qec.ppr.arbitrary ["X", "Z"](%1) %q3, %2#1 : !quantum.bit, !quantum.bit + %2:2 = qec.ppr.arbitrary ["X", "Z"](%0) %q0, %q1: !quantum.bit, !quantum.bit + %3:2 = qec.ppr.arbitrary ["X", "Z"](%1) %q2, %2#1 : !quantum.bit, !quantum.bit func.return } @@ -1033,10 +1033,10 @@ func.func public @half_compatible_qubits(%q1: !quantum.bit, %q2: !quantum.bit, % // re-arranging qubits is ok as long as the pauli words are re-arranged too // CHECK-LABEL: mix_and_match -func.func public @mix_and_match(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64) { +func.func public @mix_and_match(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, %1: f64) { // CHECK: [[angle:%.+]] = arith.addf // CHECK: qec.ppr.arbitrary [{{.+}}]([[angle]]) - %2:2 = qec.ppr.arbitrary ["Z", "Y"](%0) %q1, %q2: !quantum.bit, !quantum.bit + %2:2 = qec.ppr.arbitrary ["Z", "Y"](%0) %q0, %q1: !quantum.bit, !quantum.bit %3:2 = qec.ppr.arbitrary ["Y", "Z"](%1) %2#1, %2#0: !quantum.bit, !quantum.bit func.return } @@ -1046,10 +1046,49 @@ func.func public @mix_and_match(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, % // re-arranging qubits without re-arranging the pauli word is NOT okay // CHECK-LABEL: mix_dont_match -func.func public @mix_dont_match(%q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64) { +func.func public @mix_dont_match(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, %1: f64) { // CHECK: qec.ppr.arbitrary ["Y", "X"] // CHECK: qec.ppr.arbitrary ["Y", "X"] - %2:2 = qec.ppr.arbitrary ["Y", "X"](%0) %q1, %q2: !quantum.bit, !quantum.bit + %2:2 = qec.ppr.arbitrary ["Y", "X"](%0) %q0, %q1: !quantum.bit, !quantum.bit %3:2 = qec.ppr.arbitrary ["Y", "X"](%1) %2#1, %2#0: !quantum.bit, !quantum.bit func.return } + +// ----- + +// check equivalent conditions are merged + +// CHECK-LABEL: merge_condition +func.func public @merge_condition(%q0: !quantum.bit, %0: f64, %1: f64, %b0: i1) { + // CHECK: [[angle:%.+]] = arith.addf + // CHECK: qec.ppr.arbitrary ["X"]([[angle]]) {{%.+}} cond({{%.+}}) + %2 = qec.ppr.arbitrary ["X"](%0) %q0 cond(%b0): !quantum.bit + %3 = qec.ppr.arbitrary ["X"](%1) %2 cond(%b0): !quantum.bit + func.return +} + +// ----- + +// dont merge different conditions + +// CHECK-LABEL: dont_merge_condition +func.func public @dont_merge_condition(%q0: !quantum.bit, %0: f64, %1: f64, %b0: i1, %b1: i1) { + // CHECK: [[in:%.+]] = qec.ppr.arbitrary ["X"]({{%.+}}) {{%.+}} cond({{%.+}}) + // CHECK: qec.ppr.arbitrary ["X"]({{%.+}}) [[in]] cond({{%.+}}) + %2 = qec.ppr.arbitrary ["X"](%0) %q0 cond(%b0): !quantum.bit + %3 = qec.ppr.arbitrary ["X"](%1) %2 cond(%b1): !quantum.bit + func.return +} + +// ----- + +// don't merge conditions and non-conditions + +// CHECK-LABEL: dont_merge_mixed +func.func public @dont_merge_mixed(%q0: !quantum.bit, %0: f64, %1: f64, %b0: i1) { + // CHECK: [[in:%.+]] = qec.ppr.arbitrary ["X"]({{%.+}}) {{%.+}} cond({{%.+}}) + // CHECK: qec.ppr.arbitrary ["X"]({{%.+}}) [[in]] + %2 = qec.ppr.arbitrary ["X"](%0) %q0 cond(%b0): !quantum.bit + %3 = qec.ppr.arbitrary ["X"](%1) %2: !quantum.bit + func.return +} From e57b49088eb23f29ff07aa20ef55b0ce5a06e72d Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 3 Dec 2025 14:43:43 -0500 Subject: [PATCH 07/21] improved comments --- mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index 3aac39c8e1..4bcbae7673 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -457,9 +457,8 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern(loc, opRotation, prevOpRotation).getResult(); @@ -467,6 +466,7 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern Date: Wed, 3 Dec 2025 14:44:02 -0500 Subject: [PATCH 08/21] re-added accidentally deleted tests for PPR conditions --- mlir/test/Quantum/MergeRotationsTest.mlir | 41 +++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/mlir/test/Quantum/MergeRotationsTest.mlir b/mlir/test/Quantum/MergeRotationsTest.mlir index 53dbbd7807..1497168a6a 100644 --- a/mlir/test/Quantum/MergeRotationsTest.mlir +++ b/mlir/test/Quantum/MergeRotationsTest.mlir @@ -902,6 +902,47 @@ func.func public @half_compatible_qubits(%q1: !quantum.bit, %q2: !quantum.bit, % // ----- +// merge correctly when conditionals are present + +// CHECK-LABEL: merge_conditionals +func.func public @merge_conditionals(%q1: !quantum.bit, %q2: !quantum.bit, %arg0: i1) { + // CHECK: qec.ppr ["X", "Z"](4) {{%.+}}, {{%.+}} cond({{%.+}}) + // CHECK-NOT: qec.ppr ["X", "Z"](8) + %0, %1 = qec.ppr ["X", "Z"](8) %q1, %q2 cond(%arg0): !quantum.bit, !quantum.bit + %2, %3 = qec.ppr ["X", "Z"](8) %0, %1 cond(%arg0): !quantum.bit, !quantum.bit + func.return +} + +// ----- + +// don't merge compatible Pauli words when conditionals differ + +// CHECK-LABEL: dont_merge_incompatible_conditionals +func.func public @dont_merge_incompatible_conditionals(%q1: !quantum.bit, %q2: !quantum.bit, %arg0: i1, %arg1: i1) { + // CHECK-NOT: qec.ppr ["X", "Z"](4) + // CHECK: [[in:%.+]]:2 = qec.ppr ["X", "Z"](8) + // CHECK: qec.ppr ["X", "Z"](8) [[in]]#0, [[in]]#1 + %0:2 = qec.ppr ["X", "Z"](8) %q1, %q2 cond(%arg0): !quantum.bit, !quantum.bit + %1:2 = qec.ppr ["X", "Z"](8) %0#0, %0#1 cond(%arg1): !quantum.bit, !quantum.bit + func.return +} + +// ----- + +// don't merge PPRs with conditional PPRs + +// CHECK-LABEL: dont_merge_conditionals +func.func public @dont_merge_conditionals(%q1: !quantum.bit, %q2: !quantum.bit, %arg0: i1) { + // CHECK-NOT: qec.ppr ["X", "Z"](4) + // CHECK: [[in:%.+]]:2 = qec.ppr ["X", "Z"](8) + // CHECK: qec.ppr ["X", "Z"](8) [[in]]#0, [[in]]#1 + %0:2 = qec.ppr ["X", "Z"](8) %q1, %q2: !quantum.bit, !quantum.bit + %1:2 = qec.ppr ["X", "Z"](8) %0#0, %0#1 cond(%arg0): !quantum.bit, !quantum.bit + func.return +} + +// ----- + // Arbitrary Angle PPR Tests // simple merge From dbf09e68ae94c53e430bb94e0c6cf91fb05c5fbd Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 3 Dec 2025 15:13:04 -0500 Subject: [PATCH 09/21] simplified mapping in helper --- mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index 4bcbae7673..dc16d97857 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -325,14 +325,14 @@ struct MergeRotationsRewritePattern : public OpRewritePattern { bool matchingQubitsAndPaulis(PPRotationArbitraryOp op, PPRotationArbitraryOp parentOp) { // construct map - std::unordered_map qubit_to_pauli; + llvm::DenseMap qubit_to_pauli; for (auto [qubit, pauli] : llvm::zip(parentOp.getOutQubits(), parentOp.getPauliProduct())) { - qubit_to_pauli[hash_value(qubit)] = pauli; + qubit_to_pauli[qubit] = pauli; } // check pairings for (auto [qubit, pauli] : llvm::zip(op.getInQubits(), op.getPauliProduct())) { - if (qubit_to_pauli[hash_value(qubit)] != pauli) { + if (qubit_to_pauli[qubit] != pauli) { return false; } } From 4ba82c01e28d26ec5124c5eab4ed0d9fdcb21d9e Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 3 Dec 2025 15:52:28 -0500 Subject: [PATCH 10/21] Update mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp Co-authored-by: Paul <79805239+paul0403@users.noreply.github.com> --- mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index dc16d97857..208d6eb9a5 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -421,7 +421,7 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern Date: Wed, 3 Dec 2025 15:57:07 -0500 Subject: [PATCH 11/21] removed unnecessary check --- mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp | 4 ---- 1 file changed, 4 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index 208d6eb9a5..c752dcd328 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -453,10 +453,6 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern Date: Wed, 3 Dec 2025 15:58:33 -0500 Subject: [PATCH 12/21] made typing more explicit --- mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index c752dcd328..53f9bcb983 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -434,7 +434,7 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern(loc, opRotation, prevOpRotation).getResult(); From 2fc4a226ad6119ea41d9d309e487ebc28a663b91 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 3 Dec 2025 16:20:51 -0500 Subject: [PATCH 13/21] corrected and improved test for permuted qubits/pauli --- mlir/test/Quantum/MergeRotationsTest.mlir | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/mlir/test/Quantum/MergeRotationsTest.mlir b/mlir/test/Quantum/MergeRotationsTest.mlir index 1497168a6a..2d9e3261db 100644 --- a/mlir/test/Quantum/MergeRotationsTest.mlir +++ b/mlir/test/Quantum/MergeRotationsTest.mlir @@ -1074,11 +1074,13 @@ func.func public @half_compatible_qubits(%q0: !quantum.bit, %q1: !quantum.bit, % // re-arranging qubits is ok as long as the pauli words are re-arranged too // CHECK-LABEL: mix_and_match -func.func public @mix_and_match(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, %1: f64) { +func.func public @mix_and_match(%z1: !quantum.bit, %y1: !quantum.bit, %0: f64, %1: f64, %2: f64) { // CHECK: [[angle:%.+]] = arith.addf - // CHECK: qec.ppr.arbitrary [{{.+}}]([[angle]]) - %2:2 = qec.ppr.arbitrary ["Z", "Y"](%0) %q0, %q1: !quantum.bit, !quantum.bit - %3:2 = qec.ppr.arbitrary ["Y", "Z"](%1) %2#1, %2#0: !quantum.bit, !quantum.bit + // CHECK: [[zOut:%.+]], [[yOut:%.+]] = qec.ppr.arbitrary ["Z", "Y"]([[angle]]) %arg0, %arg1 + // CHECK: qec.ppr.arbitrary ["Y", "X"]({{%.+}}) [[yOut]], [[zOut]] + %z2, %y2 = qec.ppr.arbitrary ["Z", "Y"](%0) %zIn, %yIn: !quantum.bit, !quantum.bit + %y3, %z3 = qec.ppr.arbitrary ["Y", "Z"](%1) %y2, %z2: !quantum.bit, !quantum.bit + %6:2 = qec.ppr.arbitrary ["Y", "X"](%2) %y3, %z3: !quantum.bit, !quantum.bit func.return } From 98206113329e44cd8ae5c985a2b9ffc1ea6e0df5 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 4 Dec 2025 11:18:13 -0500 Subject: [PATCH 14/21] updated test and implemented permutation handling --- .../Transforms/MergeRotationsPatterns.cpp | 90 +++++++++++-------- mlir/test/Quantum/MergeRotationsTest.mlir | 18 ++-- 2 files changed, 62 insertions(+), 46 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index 53f9bcb983..1a0403fba5 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -322,24 +322,6 @@ struct MergeRotationsRewritePattern : public OpRewritePattern { } }; -bool matchingQubitsAndPaulis(PPRotationArbitraryOp op, PPRotationArbitraryOp parentOp) -{ - // construct map - llvm::DenseMap qubit_to_pauli; - for (auto [qubit, pauli] : llvm::zip(parentOp.getOutQubits(), parentOp.getPauliProduct())) { - qubit_to_pauli[qubit] = pauli; - } - - // check pairings - for (auto [qubit, pauli] : llvm::zip(op.getInQubits(), op.getPauliProduct())) { - if (qubit_to_pauli[qubit] != pauli) { - return false; - } - } - - return true; -} - struct MergePPRRewritePattern : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -420,51 +402,81 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern(definingOp); - - if (!prevOp) { + auto parentOp = dyn_cast(definingOp); + if (!parentOp) { return failure(); } - // verify that prevOp agrees on all qubits, not just the first - for (mlir::Value qubit : inQubits) { - if (qubit.getDefiningOp() != prevOp) { + // verify that parentOp is parent of all qubits + for (mlir::Value qubit : opInQubits) { + if (qubit.getDefiningOp() != parentOp) { return failure(); } } - - // check that the same pauli operators are applied to the same qubits - if (!matchingQubitsAndPaulis(op, prevOp)) { + ValueRange parentOpOutQubits = parentOp.getOutQubits(); + if (parentOpOutQubits.size() != opInQubits.size()) { return failure(); } + // When two rotations have permuted Pauli strings, we can still merge them, we just need to + // correctly re-map the inputs. This map stores the index of a qubit in parentOp's out + // qubits at the index it appears in op's in qubits. + uint16_t *inverse_permutation = (uint16_t *)malloc(opInQubits.size() * sizeof(uint16_t)); + for (auto [i, parentQubit] : llvm::enumerate(parentOpOutQubits)) { + for (auto [j, opQubit] : llvm::enumerate(opInQubits)) { + if (parentQubit == opQubit) { + inverse_permutation[j] = i; + } + } + } + + // check Pauli + qubit pairings + mlir::ArrayAttr opPauliProduct = op.getPauliProduct(); + mlir::ArrayAttr parentOpPauliProduct = parentOp.getPauliProduct(); + for (size_t i = 0; i < opInQubits.size(); i++) { + if (opPauliProduct[i] != parentOpPauliProduct[inverse_permutation[i]]) { + return failure(); + } + } + // check same conditionals - if (op.getCondition() != prevOp.getCondition()) { + mlir::Value opCondition = op.getCondition(); + if (opCondition != parentOp.getCondition()) { return failure(); } - mlir::Value opRotation = op.getArbitraryAngle(); - mlir::Value prevOpRotation = prevOp.getArbitraryAngle(); - - // create merged op mlir::Location loc = op.getLoc(); + + mlir::Value opRotation = op.getArbitraryAngle(); + mlir::Value parentOpRotation = parentOp.getArbitraryAngle(); mlir::Value newAngleOp = - rewriter.create(loc, opRotation, prevOpRotation).getResult(); + rewriter.create(loc, opRotation, parentOpRotation).getResult(); + + // We need to construct the Pauli string + inQubits for new op. The simplest way to ensure + // that permuted PPRs can merge correctly is to maintain output qubits order and permute + // input qubits + mlir::ValueRange parentOpInQubits = parentOp.getInQubits(); + SmallVector newInQubits; + for (size_t i = 0; i < parentOpInQubits.size(); i++) { + newInQubits.push_back(parentOpInQubits[inverse_permutation[i]]); + } - auto mergeOp = rewriter.create( - loc, op.getOutQubits().getTypes(), op.getPauliProduct(), newAngleOp, - prevOp.getInQubits(), op.getCondition()); + auto mergeOp = rewriter.create(loc, parentOpOutQubits.getTypes(), + opPauliProduct, newAngleOp, + newInQubits, opCondition); // replace and erase old ops rewriter.replaceOp(op, mergeOp); - rewriter.eraseOp(prevOp); + rewriter.eraseOp(parentOp); + + free(inverse_permutation); return success(); } diff --git a/mlir/test/Quantum/MergeRotationsTest.mlir b/mlir/test/Quantum/MergeRotationsTest.mlir index 2d9e3261db..2bf5287a93 100644 --- a/mlir/test/Quantum/MergeRotationsTest.mlir +++ b/mlir/test/Quantum/MergeRotationsTest.mlir @@ -1074,13 +1074,17 @@ func.func public @half_compatible_qubits(%q0: !quantum.bit, %q1: !quantum.bit, % // re-arranging qubits is ok as long as the pauli words are re-arranged too // CHECK-LABEL: mix_and_match -func.func public @mix_and_match(%z1: !quantum.bit, %y1: !quantum.bit, %0: f64, %1: f64, %2: f64) { - // CHECK: [[angle:%.+]] = arith.addf - // CHECK: [[zOut:%.+]], [[yOut:%.+]] = qec.ppr.arbitrary ["Z", "Y"]([[angle]]) %arg0, %arg1 - // CHECK: qec.ppr.arbitrary ["Y", "X"]({{%.+}}) [[yOut]], [[zOut]] - %z2, %y2 = qec.ppr.arbitrary ["Z", "Y"](%0) %zIn, %yIn: !quantum.bit, !quantum.bit - %y3, %z3 = qec.ppr.arbitrary ["Y", "Z"](%1) %y2, %z2: !quantum.bit, !quantum.bit - %6:2 = qec.ppr.arbitrary ["Y", "X"](%2) %y3, %z3: !quantum.bit, !quantum.bit +func.func public @mix_and_match(%z0: !quantum.bit, %y0: !quantum.bit, %0: f64, %1: f64, %2: f64, %3: f64) { + // think of X as placeholder here, they're necessary to prevent merging, but the other variable + // intended to match throughout the test, i.e. y0 -> X -> Y -> Y -> Y, never goes through a Z + // CHECK-DAG: [[res0:%.+]]:2 = qec.ppr.arbitrary ["Z", "X"]({{%.+}}) [[z0:%.+]], [[y0:%.+]] + // CHECK-DAG: [[angle:%.+]] = arith.addf + // CHECK: [[res1:%.+]]:2 = qec.ppr.arbitrary ["Y", "Z"]([[angle]]) [[res0]]#1, [[res0]]#0 + // CHECK: [[res2:%.+]]:2 = qec.ppr.arbitrary ["Y", "X"]({{%.+}}) [[res1]]#0, [[res1]]#1 + %z1, %y1 = qec.ppr.arbitrary ["Z", "X"](%0) %z0, %y0: !quantum.bit, !quantum.bit + %z2, %y2 = qec.ppr.arbitrary ["Z", "Y"](%1) %z1, %y1: !quantum.bit, !quantum.bit + %y3, %z3 = qec.ppr.arbitrary ["Y", "Z"](%2) %y2, %z2: !quantum.bit, !quantum.bit + %6:2 = qec.ppr.arbitrary ["Y", "X"](%3) %y3, %z3: !quantum.bit, !quantum.bit func.return } From 831620e313638f2184d39bc9d9a3f58d99063f0a Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 4 Dec 2025 14:30:04 -0500 Subject: [PATCH 15/21] improved testing for permuted qubits/paulis --- mlir/test/Quantum/MergeRotationsTest.mlir | 48 ++++++++++++++++------- 1 file changed, 34 insertions(+), 14 deletions(-) diff --git a/mlir/test/Quantum/MergeRotationsTest.mlir b/mlir/test/Quantum/MergeRotationsTest.mlir index 2bf5287a93..c3f1e5179d 100644 --- a/mlir/test/Quantum/MergeRotationsTest.mlir +++ b/mlir/test/Quantum/MergeRotationsTest.mlir @@ -1073,27 +1073,34 @@ func.func public @half_compatible_qubits(%q0: !quantum.bit, %q1: !quantum.bit, % // re-arranging qubits is ok as long as the pauli words are re-arranged too -// CHECK-LABEL: mix_and_match -func.func public @mix_and_match(%z0: !quantum.bit, %y0: !quantum.bit, %0: f64, %1: f64, %2: f64, %3: f64) { - // think of X as placeholder here, they're necessary to prevent merging, but the other variable - // intended to match throughout the test, i.e. y0 -> X -> Y -> Y -> Y, never goes through a Z - // CHECK-DAG: [[res0:%.+]]:2 = qec.ppr.arbitrary ["Z", "X"]({{%.+}}) [[z0:%.+]], [[y0:%.+]] +// CHECK-LABEL: merge_permutations +func.func public @merge_permutations(%z0: !quantum.bit, %y0: !quantum.bit, %0: f64, %1: f64, %2: f64, %3: f64) { // CHECK-DAG: [[angle:%.+]] = arith.addf - // CHECK: [[res1:%.+]]:2 = qec.ppr.arbitrary ["Y", "Z"]([[angle]]) [[res0]]#1, [[res0]]#0 - // CHECK: [[res2:%.+]]:2 = qec.ppr.arbitrary ["Y", "X"]({{%.+}}) [[res1]]#0, [[res1]]#1 - %z1, %y1 = qec.ppr.arbitrary ["Z", "X"](%0) %z0, %y0: !quantum.bit, !quantum.bit - %z2, %y2 = qec.ppr.arbitrary ["Z", "Y"](%1) %z1, %y1: !quantum.bit, !quantum.bit - %y3, %z3 = qec.ppr.arbitrary ["Y", "Z"](%2) %y2, %z2: !quantum.bit, !quantum.bit - %6:2 = qec.ppr.arbitrary ["Y", "X"](%3) %y3, %z3: !quantum.bit, !quantum.bit + // CHECK: qec.ppr.arbitrary ["Y", "Z"]([[angle]]) %arg1, %arg0 + %z1, %y1 = qec.ppr.arbitrary ["Z", "Y"](%1) %z0, %y0: !quantum.bit, !quantum.bit + %y2, %z2 = qec.ppr.arbitrary ["Y", "Z"](%2) %y1, %z1: !quantum.bit, !quantum.bit func.return } // ----- -// re-arranging qubits without re-arranging the pauli word is NOT okay +// check permutations with duplicate Pauli symbols -// CHECK-LABEL: mix_dont_match -func.func public @mix_dont_match(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, %1: f64) { +// CHECK-LABEL: merge_permutations_with_duplicates +func.func public @merge_permutations_with_duplicates(%q0: !quantum.bit, %q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64, %2:f64) { + // CHECK: [[angle:%.+]] = arith.addf + // CHECK: qec.ppr.arbitrary ["Y", "X", "X"]([[angle]]) %arg1, %arg2, %arg0 + %3:3 = qec.ppr.arbitrary ["X", "Y", "X"](%0) %q0, %q1, %q2: !quantum.bit, !quantum.bit, !quantum.bit + %4:3 = qec.ppr.arbitrary ["Y", "X", "X"](%1) %3#1, %3#2, %3#0: !quantum.bit, !quantum.bit, !quantum.bit + func.return +} + +// ----- + +// re-arranging qubits without re-arranging the Pauli word is NOT okay + +// CHECK-LABEL: dont_merge_permutations_qubits +func.func public @dont_merge_permutations_qubits(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, %1: f64) { // CHECK: qec.ppr.arbitrary ["Y", "X"] // CHECK: qec.ppr.arbitrary ["Y", "X"] %2:2 = qec.ppr.arbitrary ["Y", "X"](%0) %q0, %q1: !quantum.bit, !quantum.bit @@ -1103,6 +1110,19 @@ func.func public @mix_dont_match(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, // ----- +// re-arranging Pauli word without re-arranging qubits is not okay + +// CHECK-LABEL: dont_merge_permutations_pauli +func.func public @dont_merge_permutations_pauli(%q0: !quantum.bit, %q1: !quantum.bit, %0: f64, %1: f64) { + // CHECK: qec.ppr.arbitrary ["Z", "Y"] + // CHECK: qec.ppr.arbitrary ["Y", "Z"] + %2:2 = qec.ppr.arbitrary ["Z", "Y"](%0) %q0, %q1: !quantum.bit, !quantum.bit + %3:2 = qec.ppr.arbitrary ["Y", "Z"](%1) %2#0, %2#1: !quantum.bit, !quantum.bit + return +} + +// ----- + // check equivalent conditions are merged // CHECK-LABEL: merge_condition From 0afae006e7f0ef1db0630509efd0dd0f9e00ef35 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 4 Dec 2025 15:01:30 -0500 Subject: [PATCH 16/21] simplified map creation --- mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index 1aa598e1ec..7e9a0e7248 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -429,12 +429,8 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern(qubit).getResultNumber(); } // check Pauli + qubit pairings From 362c001a6c9fb2eb25cae4dd2eb4acc2f10fbf84 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 4 Dec 2025 15:11:35 -0500 Subject: [PATCH 17/21] further improved map implementation --- mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index 7e9a0e7248..af415611a6 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -428,9 +428,9 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern(qubit).getResultNumber(); + SmallVector inverse_permutation; + for (auto qubit : opInQubits) { + inverse_permutation.push_back(cast(qubit).getResultNumber()); } // check Pauli + qubit pairings @@ -472,8 +472,6 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern Date: Fri, 5 Dec 2025 10:08:52 -0500 Subject: [PATCH 18/21] Update mlir/test/Quantum/MergeRotationsTest.mlir Co-authored-by: Paul <79805239+paul0403@users.noreply.github.com> --- mlir/test/Quantum/MergeRotationsTest.mlir | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mlir/test/Quantum/MergeRotationsTest.mlir b/mlir/test/Quantum/MergeRotationsTest.mlir index c3f1e5179d..f43a56bc95 100644 --- a/mlir/test/Quantum/MergeRotationsTest.mlir +++ b/mlir/test/Quantum/MergeRotationsTest.mlir @@ -1074,11 +1074,11 @@ func.func public @half_compatible_qubits(%q0: !quantum.bit, %q1: !quantum.bit, % // re-arranging qubits is ok as long as the pauli words are re-arranged too // CHECK-LABEL: merge_permutations -func.func public @merge_permutations(%z0: !quantum.bit, %y0: !quantum.bit, %0: f64, %1: f64, %2: f64, %3: f64) { +func.func public @merge_permutations(%z0: !quantum.bit, %y0: !quantum.bit, %0: f64, %1: f64) { // CHECK-DAG: [[angle:%.+]] = arith.addf // CHECK: qec.ppr.arbitrary ["Y", "Z"]([[angle]]) %arg1, %arg0 - %z1, %y1 = qec.ppr.arbitrary ["Z", "Y"](%1) %z0, %y0: !quantum.bit, !quantum.bit - %y2, %z2 = qec.ppr.arbitrary ["Y", "Z"](%2) %y1, %z1: !quantum.bit, !quantum.bit + %z1, %y1 = qec.ppr.arbitrary ["Z", "Y"](%0) %z0, %y0: !quantum.bit, !quantum.bit + %y2, %z2 = qec.ppr.arbitrary ["Y", "Z"](%1) %y1, %z1: !quantum.bit, !quantum.bit func.return } From 820532e22cb5f638374d977deb0628e2274be053 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Fri, 5 Dec 2025 13:09:28 -0500 Subject: [PATCH 19/21] removed unnecessary argument --- mlir/test/Quantum/MergeRotationsTest.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/test/Quantum/MergeRotationsTest.mlir b/mlir/test/Quantum/MergeRotationsTest.mlir index f43a56bc95..4dfa396862 100644 --- a/mlir/test/Quantum/MergeRotationsTest.mlir +++ b/mlir/test/Quantum/MergeRotationsTest.mlir @@ -1087,7 +1087,7 @@ func.func public @merge_permutations(%z0: !quantum.bit, %y0: !quantum.bit, %0: f // check permutations with duplicate Pauli symbols // CHECK-LABEL: merge_permutations_with_duplicates -func.func public @merge_permutations_with_duplicates(%q0: !quantum.bit, %q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64, %2:f64) { +func.func public @merge_permutations_with_duplicates(%q0: !quantum.bit, %q1: !quantum.bit, %q2: !quantum.bit, %0: f64, %1: f64) { // CHECK: [[angle:%.+]] = arith.addf // CHECK: qec.ppr.arbitrary ["Y", "X", "X"]([[angle]]) %arg1, %arg2, %arg0 %3:3 = qec.ppr.arbitrary ["X", "Y", "X"](%0) %q0, %q1, %q2: !quantum.bit, !quantum.bit, !quantum.bit From ebec260eb192d669c9bea817a1fdcb29b5870241 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Fri, 5 Dec 2025 15:07:22 -0500 Subject: [PATCH 20/21] cleanup --- .../test/pytest/test_peephole_optimizations.py | 2 +- .../Quantum/Transforms/MergeRotationsPatterns.cpp | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/frontend/test/pytest/test_peephole_optimizations.py b/frontend/test/pytest/test_peephole_optimizations.py index d3893cdb5f..db938b81a6 100644 --- a/frontend/test/pytest/test_peephole_optimizations.py +++ b/frontend/test/pytest/test_peephole_optimizations.py @@ -533,7 +533,7 @@ def test_merge_rotation_arbitrary_angle_ppr(): @qml.qjit(pipelines=my_pipeline, target="mlir") def test_merge_rotation_ppr_workflow(): - @qml.transforms.merge_rotations # have to use qml to be capture-compatible + @qml.transforms.merge_rotations @qml.qnode(qml.device("lightning.qubit", wires=2)) def circuit(x, y): qml.PauliRot(x, pauli_word="ZY", wires=[0, 1]) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index af415611a6..02869dc1fc 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -430,12 +430,12 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern inverse_permutation; for (auto qubit : opInQubits) { - inverse_permutation.push_back(cast(qubit).getResultNumber()); + inverse_permutation.push_back(cast(qubit).getResultNumber()); } // check Pauli + qubit pairings - mlir::ArrayAttr opPauliProduct = op.getPauliProduct(); - mlir::ArrayAttr parentOpPauliProduct = parentOp.getPauliProduct(); + ArrayAttr opPauliProduct = op.getPauliProduct(); + ArrayAttr parentOpPauliProduct = parentOp.getPauliProduct(); for (size_t i = 0; i < opInQubits.size(); i++) { if (opPauliProduct[i] != parentOpPauliProduct[inverse_permutation[i]]) { return failure(); @@ -448,17 +448,17 @@ struct MergePPRArbitraryRewritePattern : public OpRewritePattern(loc, opRotation, parentOpRotation).getResult(); // We need to construct the Pauli string + inQubits for new op. The simplest way to ensure // that permuted PPRs can merge correctly is to maintain output qubits order and permute // input qubits - mlir::ValueRange parentOpInQubits = parentOp.getInQubits(); + ValueRange parentOpInQubits = parentOp.getInQubits(); SmallVector newInQubits; for (size_t i = 0; i < parentOpInQubits.size(); i++) { newInQubits.push_back(parentOpInQubits[inverse_permutation[i]]); @@ -501,7 +501,7 @@ struct MergeMultiRZRewritePattern : public OpRewritePattern { auto parentTheta = parentOp.getTheta(); auto theta = op.getTheta(); - mlir::Value sumParam = rewriter.create(loc, parentTheta, theta).getResult(); + Value sumParam = rewriter.create(loc, parentTheta, theta).getResult(); auto mergeOp = rewriter.create(loc, outQubitsTypes, outQubitsCtrlTypes, sumParam, parentInQubits, nullptr, parentInCtrlQubits, From ee9da71cf9959a6bb7f4613167f853bdcc95eafe Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Fri, 5 Dec 2025 15:13:48 -0500 Subject: [PATCH 21/21] missed one --- mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp index 02869dc1fc..960f812636 100644 --- a/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp +++ b/mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp @@ -501,7 +501,7 @@ struct MergeMultiRZRewritePattern : public OpRewritePattern { auto parentTheta = parentOp.getTheta(); auto theta = op.getTheta(); - Value sumParam = rewriter.create(loc, parentTheta, theta).getResult(); + mlir::Value sumParam = rewriter.create(loc, parentTheta, theta).getResult(); auto mergeOp = rewriter.create(loc, outQubitsTypes, outQubitsCtrlTypes, sumParam, parentInQubits, nullptr, parentInCtrlQubits,