diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index d88e12dfa7..14ea2e85f9 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -138,6 +138,21 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { */ Value floatConstant(double value); + /** + * @brief Create a constant boolean value + * @param value The value to store in the constant + * @return The value produced by the constant operation + * + * @par Example: + * ```c++ + * auto c = builder.boolConstant(true); + * ``` + * ```mlir + * %c = arith.constant 1 : i1 + * ``` + */ + Value boolConstant(bool value); + //===--------------------------------------------------------------------===// // Memory Management //===--------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td index 971dc18f71..9a95a0642e 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td @@ -201,4 +201,38 @@ def HadamardLifting : Pass<"hadamard-lifting", "mlir::ModuleOp"> { }]; } +def MeasurementLifting : Pass<"measurement-lifting", "mlir::ModuleOp"> { + let dependentDialects = ["mlir::qco::QCODialect", + "::mlir::arith::ArithDialect", + ]; + let summary = "This pass attempts to move measurements as far up as " + "possible, shifting them above gates that commute with them. " + "This is done to enable qubit reuse and other optimizations."; + let description = [{ + This pass applies measurement lifting, moving measurements up the code as far as possible. + Measurement lifting is a subroutine of the qubit reuse routine. The goal is to measure qubits earlier in the + circuit to reuse them and to potentially remove some quantum gates. + + Measurement lifting uses the following commutation rules: + ┌──────┐ ┌──────┐ + ──■──┤ Meas │────── ─┤ Meas ├──■─── + │ └──────┘ └──────┘ │ + ┌─┴─┐ => ┌─┴─┐ + ┤ U ├──────────────── ─────────┤ U ├─ + └───┘ └───┘ + (Where U is any (controlled) unitary gate) + + ┌───┐┌──────┐ ┌──────┐┌───┐ + ┤ P ├┤ Meas ├ => ┤ Meas ├┤ P ├ + └───┘└──────┘ └──────┘└───┘ + (Where P is any diagonal gate, e.g., `z`, `s`, ...) + + ┌───┐┌──────┐ ┌───────┐┌───┐ + ┤ X ├┤ Meas ├ => ┤ Meas* ├┤ X ├ + └───┘└──────┘ └───────┘└───┘ + (Where Meas* is a measurement after which the outcome is classically negated) + + }]; +} + #endif // MLIR_DIALECT_QCO_TRANSFORMS_PASSES_TD diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index f766602a24..5d800fb99e 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -92,6 +92,11 @@ Value QCOProgramBuilder::floatConstant(const double value) { return arith::ConstantOp::create(*this, getF64FloatAttr(value)).getResult(); } +Value QCOProgramBuilder::boolConstant(const bool value) { + checkFinalized(); + return arith::ConstantOp::create(*this, getBoolAttr(value)).getResult(); +} + Value& QCOProgramBuilder::QubitRegister::operator[](const size_t index) { if (index >= qubits.size()) { llvm::reportFatalUsageError("Qubit index out of bounds"); diff --git a/mlir/lib/Dialect/QCO/Transforms/Optimizations/MeasurementLifting.cpp b/mlir/lib/Dialect/QCO/Transforms/Optimizations/MeasurementLifting.cpp new file mode 100644 index 0000000000..8e752eeca7 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/MeasurementLifting.cpp @@ -0,0 +1,243 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +// +// Created by damian on 5/13/26. +// + +#include "mlir/Dialect/QCO/IR/QCOInterfaces.h" +#include "mlir/Dialect/QCO/IR/QCOOps.h" +#include "mlir/Dialect/QCO/Transforms/Passes.h" + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace mlir::qco { + +#define GEN_PASS_DEF_MEASUREMENTLIFTING +#include "mlir/Dialect/QCO/Transforms/Passes.h.inc" + +/** + * @brief Checks if the given operation is an inverting gate. + * @param op The operation to check. + * @return True if the operation is an inverting gate, false otherwise. + */ +static bool isInverting(Operation* op) { return isa(op); } + +/** + * @brief Checks if the given operation is a diagonal gate. + * @param op The operation to check. + * @return True if the operation is a diagonal gate, false otherwise. + */ +static bool isDiagonal(Operation* op) { + if (auto c = dyn_cast(op)) { + if (c.getNumBodyUnitaries() != 1) { + return false; + } + return isDiagonal(c.getBodyUnitary(0)); + } + if (auto i = dyn_cast(op)) { + if (i.getNumBodyUnitaries() != 1) { + return false; + } + return isDiagonal(i.getBodyUnitary(0)); + } + return isa(op); +} + +/** + * @brief This method swaps a gate with a measurement. + * @param gate The gate to swap. + * @param measurement The measurement to swap. + * @param rewriter The used rewriter. + */ +static void swapGateWithMeasurement(UnitaryOpInterface gate, + MeasureOp measurement, + mlir::PatternRewriter& rewriter) { + auto measurementInput = measurement.getQubitIn(); + auto gateInput = gate.getInputForOutput(measurementInput); + rewriter.replaceUsesWithIf(measurementInput, gateInput, + [&](mlir::OpOperand& operand) { + // We only replace the single use by the + // measure op + return operand.getOwner() == measurement; + }); + rewriter.replaceUsesWithIf(gateInput, measurement.getQubitOut(), + [&](mlir::OpOperand& operand) { + // We only replace the single use by the + // predecessor + return operand.getOwner() == gate; + }); + rewriter.replaceUsesWithIf(measurement.getQubitOut(), measurementInput, + [&](mlir::OpOperand& operand) { + // All further uses of the measurement output now + // use the gate output + return operand.getOwner() != gate; + }); + rewriter.moveOpBefore(measurement, gate); +} + +namespace { +/** + * @brief This pattern is responsible for lifting measurements above any phase + * gates. + */ +struct LiftMeasurementsAbovePhaseGatesPattern final + : mlir::OpRewritePattern { + + explicit LiftMeasurementsAbovePhaseGatesPattern(mlir::MLIRContext* context) + : OpRewritePattern(context) {} + + mlir::LogicalResult + matchAndRewrite(MeasureOp op, + mlir::PatternRewriter& rewriter) const override { + const auto qubitVariable = op.getQubitIn(); + auto* predecessor = qubitVariable.getDefiningOp(); + + auto predecessorUnitary = mlir::dyn_cast(predecessor); + + if (!predecessorUnitary) { + return mlir::failure(); + } + + if (isDiagonal(predecessor)) { + swapGateWithMeasurement(predecessorUnitary, op, rewriter); + return mlir::success(); + } + + return mlir::failure(); + } +}; + +/** + * @brief This pattern is responsible for lifting measurements above any + * non-phase gates. + */ +struct LiftMeasurementsAboveInvertingGatesPattern final + : mlir::OpRewritePattern { + + explicit LiftMeasurementsAboveInvertingGatesPattern( + mlir::MLIRContext* context) + : OpRewritePattern(context) {} + + /** + * @brief Checks if the given qubit is not used anymore. + * @param outQubit The output qubit to check. + * @return True if all users are resets/deallocs, false otherwise. + */ + static bool outputQubitRemainsUnused(mlir::Value outQubit) { + return llvm::all_of(outQubit.getUsers(), [](mlir::Operation* user) { + return mlir::isa(user) || mlir::isa(user); + }); + } + + mlir::LogicalResult + matchAndRewrite(MeasureOp op, + mlir::PatternRewriter& rewriter) const override { + const auto qubitVariable = op.getQubitIn(); + auto* predecessor = qubitVariable.getDefiningOp(); + + auto predecessorUnitary = mlir::dyn_cast(predecessor); + + if (!predecessorUnitary) { + return mlir::failure(); + } + + if (isInverting(predecessor) && + predecessorUnitary.getInputQubits().size() == 1) { + swapGateWithMeasurement(predecessorUnitary, op, rewriter); + rewriter.setInsertionPointAfter(op); + const mlir::Value trueConstant = rewriter.create( + op.getLoc(), rewriter.getBoolAttr(true)); + auto inversion = rewriter.create( + op.getLoc(), op.getResult(), trueConstant); + // We need `replaceUsesWithIf` so that we can replace all uses except for + // the one use that defines the inverted bit. + rewriter.replaceUsesWithIf(op.getResult(), inversion.getResult(), + [&](mlir::OpOperand& operand) { + return operand.getOwner() != inversion; + }); + return mlir::success(); + } + + return mlir::failure(); + } +}; + +/** + * @brief This pattern is responsible for applying the "deferred measurement + * principle", lifting measurements above controls. + */ +struct LiftMeasurementsAboveControlsPattern final + : mlir::OpRewritePattern { + + explicit LiftMeasurementsAboveControlsPattern(mlir::MLIRContext* context) + : OpRewritePattern(context) {} + + mlir::LogicalResult + matchAndRewrite(MeasureOp op, + mlir::PatternRewriter& rewriter) const override { + const auto qubitVariable = op.getQubitIn(); + auto* predecessor = qubitVariable.getDefiningOp(); + auto predecessorCtrl = mlir::dyn_cast(predecessor); + + if (!predecessorCtrl) { + return mlir::failure(); + } + + if (llvm::find(predecessorCtrl.getControlsOut(), qubitVariable) == + predecessorCtrl.getControlsOut().end()) { + // The measured qubit is a target, not a control of the gate. + return mlir::failure(); + } + + swapGateWithMeasurement(predecessorCtrl, op, rewriter); + + return mlir::success(); + } +}; + +/** + * @brief Pass raises Measurements above controlled and uncontrolled gates + * gates. + */ +struct MeasurementLifting final + : impl::MeasurementLiftingBase { + using MeasurementLiftingBase::MeasurementLiftingBase; + +protected: + void runOnOperation() override { + const auto op = getOperation(); + auto* ctx = &getContext(); + + // Define the set of patterns to use. + RewritePatternSet patterns(ctx); + patterns.add(patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add(patterns.getContext()); + + // Apply patterns in an iterative and greedy manner. + if (failed(applyPatternsGreedily(op, std::move(patterns)))) { + signalPassFailure(); + } + } +}; + +} // namespace + +} // namespace mlir::qco diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt index e80e957680..bdc78401c8 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt @@ -7,7 +7,7 @@ # Licensed under the MIT License set(target_name mqt-core-mlir-unittest-optimizations) -add_executable(${target_name} test_qco_hadamard_lifting.cpp +add_executable(${target_name} test_qco_hadamard_lifting.cpp test_qco_measurement_lifting.cpp test_qco_merge_single_qubit_rotation.cpp test_quantum_loop_unroll.cpp) target_link_libraries( diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_measurement_lifting.cpp b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_measurement_lifting.cpp new file mode 100644 index 0000000000..f197a467f7 --- /dev/null +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_measurement_lifting.cpp @@ -0,0 +1,448 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +// +// Created by damian on 5/21/26. +// + +#include "mlir/Dialect/QCO/Builder/QCOProgramBuilder.h" +#include "mlir/Dialect/QCO/IR/QCODialect.h" +#include "mlir/Dialect/QCO/Transforms/Passes.h" +#include "mlir/Support/IRVerification.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace { + +using namespace mlir; +using namespace mlir::qco; + +class QCOMeasurementLiftingTest : public testing::Test { + +protected: + MLIRContext context; + QCOProgramBuilder programBuilder; + QCOProgramBuilder referenceBuilder; + OwningOpRef module; + OwningOpRef reference; + + QCOMeasurementLiftingTest() + : programBuilder(&context), referenceBuilder(&context) {} + + void SetUp() override { + // Register all necessary dialects + DialectRegistry registry; + registry.insert(); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + } + + /** + * @brief Adds the measurementLiftingPass to the current context and runs it. + */ + static LogicalResult runMeasurementLiftingPass(ModuleOp module) { + PassManager pm(module.getContext()); + pm.addPass(createMeasurementLifting()); + pm.addPass(createCanonicalizerPass()); + return pm.run(module); + } + + /** + * @brief Adds the canonicalizerPass to the current context and runs it. + */ + static LogicalResult runCanonicalizerPass(ModuleOp module) { + PassManager pm(module.getContext()); + pm.addPass(createCanonicalizerPass()); + return pm.run(module); + } +}; + +} // namespace + +TEST_F(QCOMeasurementLiftingTest, liftMeasurementOverPositiveControl) { + programBuilder.initialize( + {programBuilder.getI1Type(), programBuilder.getI1Type()}); + auto q0S0 = programBuilder.allocQubit(); + auto q1S0 = programBuilder.allocQubit(); + + auto [q1S1, q0S1] = programBuilder.cx(q1S0, q0S0); + auto [q0S2, q1S2] = programBuilder.ch(q0S1, q1S1); + auto [q0S3, q1S3] = programBuilder.cx(q0S2, q1S2); + + auto [q0S4, c0] = programBuilder.measure(q0S3); + auto [q1S4, c1] = programBuilder.measure(q1S3); + + programBuilder.sink(q0S4); + programBuilder.sink(q1S4); + module = programBuilder.finalize({c0, c1}); + + referenceBuilder.initialize( + {referenceBuilder.getI1Type(), referenceBuilder.getI1Type()}); + auto r0S0 = referenceBuilder.allocQubit(); + auto r1S0 = referenceBuilder.allocQubit(); + + auto [r1S1, r0S1] = referenceBuilder.cx(r1S0, r0S0); + auto [r0S2, cr0] = referenceBuilder.measure(r0S1); + auto [r0S3, r1S2] = referenceBuilder.ch(r0S2, r1S1); + auto [r0S4, r1S3] = referenceBuilder.cx(r0S3, r1S2); + + auto [r1S4, cr1] = referenceBuilder.measure(r1S3); + + referenceBuilder.sink(r0S4); + referenceBuilder.sink(r1S4); + reference = referenceBuilder.finalize({cr0, cr1}); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +} + +TEST_F(QCOMeasurementLiftingTest, liftMeasurementOverOneOfMultipleControls) { + programBuilder.initialize({programBuilder.getI1Type(), + programBuilder.getI1Type(), + programBuilder.getI1Type()}); + auto q0S0 = programBuilder.allocQubit(); + auto q1S0 = programBuilder.allocQubit(); + auto q2S0 = programBuilder.allocQubit(); + + auto [q12_0, q0S1] = + programBuilder.ctrl({q1S0, q2S0}, {q0S0}, [&](const ValueRange target) { + return SmallVector{programBuilder.x(target[0])}; + }); + auto [q12_1, q0S2] = programBuilder.ctrl( + {q12_0[1], q12_0[0]}, q0S1, [&](const ValueRange target) { + return SmallVector{programBuilder.h(target[0])}; + }); + auto [q12_2, q0S3] = programBuilder.ctrl( + {q12_1[1], q12_1[0]}, q0S2, [&](const ValueRange target) { + return SmallVector{programBuilder.x(target[0])}; + }); + + auto [q1S4, c1] = programBuilder.measure(q12_2[0]); + + auto q0S4 = programBuilder.h(q0S3[0]); + auto q2S4 = programBuilder.h(q12_2[1]); + + auto [q0S5, c0] = programBuilder.measure(q0S4); + auto [q2S5, c2] = programBuilder.measure(q2S4); + + programBuilder.sink(q0S5); + programBuilder.sink(q1S4); + programBuilder.sink(q2S5); + + module = programBuilder.finalize({c0, c1, c2}); + + referenceBuilder.initialize({referenceBuilder.getI1Type(), + referenceBuilder.getI1Type(), + referenceBuilder.getI1Type()}); + auto r0S0 = referenceBuilder.allocQubit(); + auto r1S0 = referenceBuilder.allocQubit(); + auto r2S0 = referenceBuilder.allocQubit(); + + auto [r1S1, cr1] = referenceBuilder.measure(r1S0); + + auto [r12_0, r0S1] = + referenceBuilder.ctrl({r1S1, r2S0}, {r0S0}, [&](const ValueRange target) { + return SmallVector{referenceBuilder.x(target[0])}; + }); + auto [r12_1, r0S2] = referenceBuilder.ctrl( + {r12_0[1], r12_0[0]}, r0S1, [&](const ValueRange target) { + return SmallVector{referenceBuilder.h(target[0])}; + }); + auto [r12_2, r0S3] = referenceBuilder.ctrl( + {r12_1[1], r12_1[0]}, r0S2, [&](const ValueRange target) { + return SmallVector{referenceBuilder.x(target[0])}; + }); + + auto r0S4 = referenceBuilder.h(r0S3[0]); + auto r2S4 = referenceBuilder.h(r12_2[1]); + + auto [r0S5, cr0] = referenceBuilder.measure(r0S4); + auto [r2S5, cr2] = referenceBuilder.measure(r2S4); + + referenceBuilder.sink(r0S5); + referenceBuilder.sink(r12_2[0]); + referenceBuilder.sink(r2S5); + + reference = referenceBuilder.finalize({cr0, cr1, cr2}); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +} + +TEST_F(QCOMeasurementLiftingTest, + liftMeasurementMultipleOverOneControlledGate) { + + programBuilder.initialize( + {programBuilder.getI1Type(), programBuilder.getI1Type()}); + auto q0S0 = programBuilder.allocQubit(); + auto q1S0 = programBuilder.allocQubit(); + auto q2S0 = programBuilder.allocQubit(); + + auto [q12_0, q0S1] = + programBuilder.ctrl({q1S0, q2S0}, {q0S0}, [&](const ValueRange target) { + return SmallVector{programBuilder.x(target[0])}; + }); + + auto [q1S1, c1] = programBuilder.measure(q12_0[0]); + auto [q2S1, c2] = programBuilder.measure(q12_0[1]); + + programBuilder.sink(q0S1[0]); + programBuilder.sink(q1S1); + programBuilder.sink(q2S1); + module = programBuilder.finalize({c1, c2}); + + referenceBuilder.initialize( + {referenceBuilder.getI1Type(), referenceBuilder.getI1Type()}); + auto r0S0 = referenceBuilder.allocQubit(); + auto r1S0 = referenceBuilder.allocQubit(); + auto r2S0 = referenceBuilder.allocQubit(); + + auto [r1S1, cr1] = referenceBuilder.measure(r1S0); + auto [r2S1, cr2] = referenceBuilder.measure(r2S0); + + auto [r12_0, r0S1] = + referenceBuilder.ctrl({r1S1, r2S1}, {r0S0}, [&](const ValueRange target) { + return SmallVector{referenceBuilder.x(target[0])}; + }); + + referenceBuilder.sink(r0S1[0]); + referenceBuilder.sink(r12_0[0]); + referenceBuilder.sink(r12_0[1]); + reference = referenceBuilder.finalize({cr1, cr2}); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +} + +TEST_F(QCOMeasurementLiftingTest, + liftMeasurementOverControlledParametrizedGate) { + programBuilder.initialize( + {programBuilder.getI1Type(), programBuilder.getI1Type()}); + auto q0S0 = programBuilder.allocQubit(); + auto q1S0 = programBuilder.allocQubit(); + + auto [q0S1, q1S1] = programBuilder.crx(std::numbers::pi / 2, q0S0, q1S0); + + auto [q0S2, c0] = programBuilder.measure(q0S1); + auto [q1S2, c1] = programBuilder.measure(q1S1); + + programBuilder.sink(q0S2); + programBuilder.sink(q1S2); + module = programBuilder.finalize({c0, c1}); + + referenceBuilder.initialize( + {referenceBuilder.getI1Type(), referenceBuilder.getI1Type()}); + auto r0S0 = referenceBuilder.allocQubit(); + auto r1S0 = referenceBuilder.allocQubit(); + + auto [r0S1, cr0] = referenceBuilder.measure(r0S0); + + auto [r0S2, r1S1] = referenceBuilder.crx(std::numbers::pi / 2, r0S1, r1S0); + + auto [r1S2, cr1] = referenceBuilder.measure(r1S1); + + referenceBuilder.sink(r0S2); + referenceBuilder.sink(r1S2); + reference = referenceBuilder.finalize({cr0, cr1}); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +} + +TEST_F(QCOMeasurementLiftingTest, liftMeasurementOverSingleX) { + + programBuilder.initialize({programBuilder.getI1Type()}); + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.x(q0); + auto [q2, c] = programBuilder.measure(q1); + programBuilder.sink(q2); + module = programBuilder.finalize(c); + + referenceBuilder.initialize({referenceBuilder.getI1Type()}); + auto r0 = referenceBuilder.allocQubit(); + auto trueConstant = referenceBuilder.boolConstant(true); + auto [r1, cr] = referenceBuilder.measure(r0); + + auto xorOp = arith::XOrIOp::create( + referenceBuilder, referenceBuilder.getLoc(), cr, trueConstant); + referenceBuilder.sink(r1); + reference = referenceBuilder.finalize(xorOp.getResult()); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +} + +TEST_F(QCOMeasurementLiftingTest, liftMeasurementOverSingleY) { + programBuilder.initialize({programBuilder.getI1Type()}); + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.y(q0); + auto [q2, c] = programBuilder.measure(q1); + programBuilder.sink(q2); + module = programBuilder.finalize({c}); + + referenceBuilder.initialize({referenceBuilder.getI1Type()}); + auto r0 = referenceBuilder.allocQubit(); + auto trueConstant = referenceBuilder.boolConstant(true); + auto [r1, cr] = referenceBuilder.measure(r0); + auto xorOp = arith::XOrIOp::create( + referenceBuilder, referenceBuilder.getLoc(), cr, trueConstant); + referenceBuilder.sink(r1); + reference = referenceBuilder.finalize({xorOp.getResult()}); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +} + +TEST_F(QCOMeasurementLiftingTest, liftMeasurementOverPhaseGates) { + programBuilder.initialize({programBuilder.getI1Type()}); + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.id(q0); + auto q2 = programBuilder.z(q1); + auto q3 = programBuilder.s(q2); + auto q4 = programBuilder.sdg(q3); + auto q5 = programBuilder.t(q4); + auto q6 = programBuilder.tdg(q5); + auto q7 = programBuilder.p(std::numbers::pi / 2, q6); + auto q8 = programBuilder.rz(std::numbers::pi / 2, q7); + auto [q9, c] = programBuilder.measure(q8); + programBuilder.sink(q9); + module = programBuilder.finalize({c}); + + referenceBuilder.initialize({referenceBuilder.getI1Type()}); + auto r0 = referenceBuilder.allocQubit(); + auto [r1, cr] = referenceBuilder.measure(r0); + referenceBuilder.sink(r1); + reference = referenceBuilder.finalize({cr}); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +} + +TEST_F(QCOMeasurementLiftingTest, liftMeasurementOverMultipleXY) { + programBuilder.initialize({programBuilder.getI1Type()}); + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.x(q0); + auto q2 = programBuilder.y(q1); + auto [q3, c] = programBuilder.measure(q2); + programBuilder.sink(q3); + module = programBuilder.finalize({c}); + + referenceBuilder.initialize({referenceBuilder.getI1Type()}); + auto r0 = referenceBuilder.allocQubit(); + auto [r1, cr] = referenceBuilder.measure(r0); + referenceBuilder.sink(r1); + reference = referenceBuilder.finalize({cr}); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +} + +TEST_F(QCOMeasurementLiftingTest, liftMeasurementOverXAndControlledGates) { + programBuilder.initialize({programBuilder.getI1Type()}); + auto q0S0 = programBuilder.allocQubit(); + auto q1S0 = programBuilder.allocQubit(); + + auto [q0S1, q1S1] = programBuilder.cy(q0S0, q1S0); + auto q0S2 = programBuilder.x(q0S1); + auto [q0S3, q1S2] = programBuilder.cy(q0S2, q1S1); + auto q0S4 = programBuilder.x(q0S3); + + auto [q0S5, c0] = programBuilder.measure(q0S4); + + programBuilder.sink(q0S5); + programBuilder.sink(q1S2); + module = programBuilder.finalize({c0}); + + referenceBuilder.initialize({referenceBuilder.getI1Type()}); + auto r0S0 = referenceBuilder.allocQubit(); + auto r1S0 = referenceBuilder.allocQubit(); + + auto [r0S1, cr0] = referenceBuilder.measure(r0S0); + + auto [r0S2, r1S1] = referenceBuilder.cx(r0S1, r1S0); + auto r0S3 = referenceBuilder.x(r0S2); + auto [r0S4, r1S2] = referenceBuilder.cx(r0S3, r1S1); + + referenceBuilder.sink(r0S4); + referenceBuilder.sink(r1S2); + reference = referenceBuilder.finalize({cr0}); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +} + +TEST_F(QCOMeasurementLiftingTest, liftMeasurementOverDiagonalGateInControl) { + programBuilder.initialize({programBuilder.getI1Type()}); + auto q0S0 = programBuilder.allocQubit(); + auto q1S0 = programBuilder.allocQubit(); + + auto [q0S1, q1S1] = programBuilder.cz(q0S0, q1S0); + + auto [q0S2, c0] = programBuilder.measure(q0S1); + + programBuilder.sink(q0S2); + programBuilder.sink(q1S1); + module = programBuilder.finalize({c0}); + + referenceBuilder.initialize({referenceBuilder.getI1Type()}); + auto r0S0 = referenceBuilder.allocQubit(); + auto r1S0 = referenceBuilder.allocQubit(); + + auto [r0S1, cr0] = referenceBuilder.measure(r0S0); + + referenceBuilder.sink(r0S1); + referenceBuilder.sink(r1S0); + reference = referenceBuilder.finalize({cr0}); + + ASSERT_TRUE(runMeasurementLiftingPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); +}