diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index d034ca6faf..d47c1db9dd 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -55,6 +55,24 @@
[(#1357)](https://github.com/PennyLaneAI/catalyst/pull/1357)
[(#1385)](https://github.com/PennyLaneAI/catalyst/pull/1385)
+* A new circuit optimization pass, `--disentangle-CNOT`, is available.
+ [(#1154)](https://github.com/PennyLaneAI/catalyst/pull/1154)
+
+ The pass disentangles CNOT gates whenever possible, e.g. when the control bit
+ is known to be in |0>, the pass removes the CNOT. The pass uses a finite state
+ machine to propagate simple one-qubit states, in order to determine
+ the input states to the CNOT.
+
+ The algorithm is taken from [Relaxed Peephole Optimization: A Novel Compiler Optimization for Quantum Circuits, by Ji Liu, Luciano Bello, and Huiyang Zhou](https://arxiv.org/abs/2012.07711).
+
+* A new circuit optimization pass, `--disentangle-SWAP`, is available.
+ [(#1297)](https://github.com/PennyLaneAI/catalyst/pull/1297)
+
+ The pass disentangles SWAP gates whenever possible by using a finite state
+ machine to propagate simple one-qubit states, similar to the `--disentangle-CNOT` pass.
+
+ The algorithm is taken from [Relaxed Peephole Optimization: A Novel Compiler Optimization for Quantum Circuits, by Ji Liu, Luciano Bello, and Huiyang Zhou](https://arxiv.org/abs/2012.07711).
+
Breaking changes 💔
* The `sample` and `counts` measurement primitives now support dynamic shot values across catalyst, although at the PennyLane side, the device shots still is constrained to a static integer literal.
@@ -193,5 +211,6 @@ Mehrdad Malekmohammadi,
William Maxwell
Romain Moyard,
Shuli Shu,
+Ritu Thombre,
Raul Torres,
Paul Haochen Wang.
diff --git a/mlir/include/Quantum/Transforms/Passes.h b/mlir/include/Quantum/Transforms/Passes.h
index d43490d7cd..d733a108d7 100644
--- a/mlir/include/Quantum/Transforms/Passes.h
+++ b/mlir/include/Quantum/Transforms/Passes.h
@@ -29,6 +29,8 @@ std::unique_ptr createRemoveChainedSelfInversePass();
std::unique_ptr createAnnotateFunctionPass();
std::unique_ptr createSplitMultipleTapesPass();
std::unique_ptr createMergeRotationsPass();
+std::unique_ptr createDisentangleCNOTPass();
+std::unique_ptr createDisentangleSWAPPass();
std::unique_ptr createIonsDecompositionPass();
std::unique_ptr createStaticCustomLoweringPass();
diff --git a/mlir/include/Quantum/Transforms/Passes.td b/mlir/include/Quantum/Transforms/Passes.td
index 5a5325adb3..168aed6a00 100644
--- a/mlir/include/Quantum/Transforms/Passes.td
+++ b/mlir/include/Quantum/Transforms/Passes.td
@@ -116,6 +116,28 @@ def MergeRotationsPass : Pass<"merge-rotations"> {
let constructor = "catalyst::createMergeRotationsPass()";
}
+def DisentangleCNOTPass : Pass<"disentangle-CNOT"> {
+ let summary = "Replace a CNOT gate with two single qubit gates whenever possible.";
+
+ let constructor = "catalyst::createDisentangleCNOTPass()";
+ let options = [
+ Option<"EmitFSMStateRemark", "emit-FSM-state-remark",
+ "bool", /*default=*/"false",
+ "Whether to emit the state analysis result from the simple states propagation FSM onto the gate operations.">,
+ ];
+}
+
+def DisentangleSWAPPass : Pass<"disentangle-SWAP"> {
+ let summary = "Replace a SWAP gate with single qubit gates and a shorter SWAPZ gates whenever possible.";
+
+ let constructor = "catalyst::createDisentangleSWAPPass()";
+ let options = [
+ Option<"EmitFSMStateRemark", "emit-FSM-state-remark",
+ "bool", /*default=*/"false",
+ "Whether to emit the state analysis result from the simple states propagation FSM onto the gate operations.">,
+ ];
+}
+
// ----- Quantum circuit transformation passes end ----- //
#endif // QUANTUM_PASSES
diff --git a/mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp b/mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp
index eb5a79feae..442a4a6656 100644
--- a/mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp
+++ b/mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp
@@ -31,6 +31,8 @@ void catalyst::registerAllCatalystPasses()
mlir::registerPass(catalyst::createCopyGlobalMemRefPass);
mlir::registerPass(catalyst::createDetensorizeSCFPass);
mlir::registerPass(catalyst::createDisableAssertionPass);
+ mlir::registerPass(catalyst::createDisentangleCNOTPass);
+ mlir::registerPass(catalyst::createDisentangleSWAPPass);
mlir::registerPass(catalyst::createEmitCatalystPyInterfacePass);
mlir::registerPass(catalyst::createGEPInboundsPass);
mlir::registerPass(catalyst::createGradientBufferizationPass);
diff --git a/mlir/lib/Quantum/Transforms/CMakeLists.txt b/mlir/lib/Quantum/Transforms/CMakeLists.txt
index 2483d4618d..1281b5bba3 100644
--- a/mlir/lib/Quantum/Transforms/CMakeLists.txt
+++ b/mlir/lib/Quantum/Transforms/CMakeLists.txt
@@ -15,6 +15,8 @@ file(GLOB SRC
SplitMultipleTapes.cpp
merge_rotation.cpp
MergeRotationsPatterns.cpp
+ DisentangleSWAP.cpp
+ DisentangleCNOT.cpp
ions_decompositions.cpp
IonsDecompositionPatterns.cpp
static_custom_lowering.cpp
diff --git a/mlir/lib/Quantum/Transforms/DisentangleCNOT.cpp b/mlir/lib/Quantum/Transforms/DisentangleCNOT.cpp
new file mode 100644
index 0000000000..bb5086f11f
--- /dev/null
+++ b/mlir/lib/Quantum/Transforms/DisentangleCNOT.cpp
@@ -0,0 +1,143 @@
+// Copyright 2024 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.
+
+// This algorithm is taken from https://arxiv.org/pdf/2012.07711, table 1
+
+#define DEBUG_TYPE "disentanglecnot"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/Debug.h"
+
+#include "Catalyst/IR/CatalystDialect.h"
+#include "Quantum/IR/QuantumOps.h"
+
+#include "PropagateSimpleStatesAnalysis.hpp"
+
+using namespace mlir;
+using namespace catalyst;
+
+namespace catalyst {
+#define GEN_PASS_DEF_DISENTANGLECNOTPASS
+#define GEN_PASS_DECL_DISENTANGLECNOTPASS
+#include "Quantum/Transforms/Passes.h.inc"
+
+struct DisentangleCNOTPass : public impl::DisentangleCNOTPassBase {
+ using impl::DisentangleCNOTPassBase::DisentangleCNOTPassBase;
+
+ bool canScheduleOn(RegisteredOperationName opInfo) const override
+ {
+ return opInfo.hasInterface();
+ }
+
+ void runOnOperation() override
+ {
+ FunctionOpInterface func = cast(getOperation());
+ mlir::IRRewriter builder(func->getContext());
+ Location loc = func->getLoc();
+
+ PropagateSimpleStatesAnalysis &pssa = getAnalysis();
+ llvm::DenseMap qubitValues = pssa.getQubitValues();
+
+ if (EmitFSMStateRemark) {
+ for (auto it = qubitValues.begin(); it != qubitValues.end(); ++it) {
+ it->first.getDefiningOp()->emitRemark(pssa.QubitState2String(it->second));
+ }
+ }
+
+ func->walk([&](quantum::CustomOp op) {
+ StringRef gate = op.getGateName();
+ if (gate != "CNOT") {
+ return;
+ }
+
+ Value controlIn = op->getOperand(0);
+ Value targetIn = op->getOperand(1);
+ Value controlOut = op->getResult(0);
+ Value targetOut = op->getResult(1);
+
+ // Do nothing if the inputs states are not tracked
+ if (!qubitValues.contains(controlIn) || !qubitValues.contains(targetIn)) {
+ return;
+ }
+
+ // |0> control, always do nothing
+ if (pssa.isZero(qubitValues[controlIn])) {
+ controlOut.replaceAllUsesWith(controlIn);
+ targetOut.replaceAllUsesWith(targetIn);
+ op->erase();
+ return;
+ }
+
+ // |1> control, insert PauliX gate on target
+ if (pssa.isOne(qubitValues[controlIn])) {
+ controlOut.replaceAllUsesWith(controlIn);
+
+ // PauliX on |+-> is unnecessary: they are eigenstates!
+ if ((pssa.isPlus(qubitValues[targetIn])) || (pssa.isMinus(qubitValues[targetIn]))) {
+ targetOut.replaceAllUsesWith(targetIn);
+ op->erase();
+ return;
+ }
+ else {
+ builder.setInsertionPoint(op);
+ quantum::CustomOp xgate = builder.create(
+ loc, /*gate_name=*/"PauliX",
+ /*in_qubits=*/mlir::ValueRange({targetIn}));
+ targetOut.replaceAllUsesWith(xgate->getResult(0));
+ op->erase();
+ return;
+ }
+ }
+
+ // |+> target, always do nothing
+ if (pssa.isPlus(qubitValues[targetIn])) {
+ controlOut.replaceAllUsesWith(controlIn);
+ targetOut.replaceAllUsesWith(targetIn);
+ op->erase();
+ return;
+ }
+
+ // |-> target, insert PauliZ on control
+ if (pssa.isMinus(qubitValues[targetIn])) {
+ targetOut.replaceAllUsesWith(targetIn);
+
+ // PauliZ on |01> is unnecessary: they are eigenstates!
+ if ((pssa.isZero(qubitValues[controlIn])) || (pssa.isOne(qubitValues[controlIn]))) {
+ controlOut.replaceAllUsesWith(controlIn);
+ op->erase();
+ return;
+ }
+ else {
+ builder.setInsertionPoint(op);
+ quantum::CustomOp zgate = builder.create(
+ loc, /*gate_name=*/"PauliZ",
+ /*in_qubits=*/mlir::ValueRange({controlIn}));
+ controlOut.replaceAllUsesWith(zgate->getResult(0));
+ op->erase();
+ return;
+ }
+ }
+ });
+ }
+};
+
+std::unique_ptr createDisentangleCNOTPass()
+{
+ return std::make_unique();
+}
+
+} // namespace catalyst
diff --git a/mlir/lib/Quantum/Transforms/DisentangleSWAP.cpp b/mlir/lib/Quantum/Transforms/DisentangleSWAP.cpp
new file mode 100644
index 0000000000..fe10e3abc4
--- /dev/null
+++ b/mlir/lib/Quantum/Transforms/DisentangleSWAP.cpp
@@ -0,0 +1,507 @@
+// Copyright 2024 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.
+
+// This algorithm is taken from https://arxiv.org/pdf/2012.07711, table 6 (Equivalences for
+// basis-states in SWAP gate)
+
+#define DEBUG_TYPE "disentangleswap"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/Debug.h"
+
+#include "Catalyst/IR/CatalystDialect.h"
+#include "Quantum/IR/QuantumOps.h"
+
+#include "PropagateSimpleStatesAnalysis.hpp"
+
+using namespace mlir;
+using namespace catalyst;
+
+namespace catalyst {
+#define GEN_PASS_DEF_DISENTANGLESWAPPASS
+#define GEN_PASS_DECL_DISENTANGLESWAPPASS
+#include "Quantum/Transforms/Passes.h.inc"
+
+struct DisentangleSWAPPass : public impl::DisentangleSWAPPassBase {
+ using impl::DisentangleSWAPPassBase::DisentangleSWAPPassBase;
+
+ // function to create a single qubit gate with a given name
+ // right after the SWAP is to be erased
+ quantum::CustomOp createSimpleOneBitGate(StringRef gateName, const Value &inQubit,
+ const Value &outQubit, mlir::IRRewriter &builder,
+ Location &loc,
+ const quantum::CustomOp &insert_after_gate)
+ {
+ OpBuilder::InsertionGuard insertionGuard(builder);
+ builder.setInsertionPointAfter(insert_after_gate);
+ quantum::CustomOp newGate =
+ builder.create(loc,
+ /*out_qubits=*/mlir::TypeRange({outQubit.getType()}),
+ /*out_ctrl_qubits=*/mlir::TypeRange(),
+ /*params=*/mlir::ValueRange(),
+ /*in_qubits=*/mlir::ValueRange({inQubit}),
+ /*gate_name=*/gateName,
+ /*adjoint=*/nullptr,
+ /*in_ctrl_qubits=*/mlir::ValueRange(),
+ /*in_ctrl_values=*/mlir::ValueRange());
+
+ return newGate;
+ }
+
+ // above function overloaded to create a single qubit gate with a given name
+ // if multiple gates are to be inserted after SWAP transformation
+ quantum::CustomOp createSimpleOneBitGate(StringRef gateName, const Value &inQubit,
+ mlir::IRRewriter &builder, Location &loc,
+ const quantum::CustomOp &insert_after_gate)
+ {
+ OpBuilder::InsertionGuard insertionGuard(builder);
+ builder.setInsertionPointAfter(insert_after_gate);
+ quantum::CustomOp newGate =
+ builder.create(loc,
+ /*out_qubits=*/mlir::TypeRange({inQubit.getType()}),
+ /*out_ctrl_qubits=*/mlir::TypeRange(),
+ /*params=*/mlir::ValueRange(),
+ /*in_qubits=*/mlir::ValueRange({inQubit}),
+ /*gate_name=*/gateName,
+ /*adjoint=*/nullptr,
+ /*in_ctrl_qubits=*/mlir::ValueRange(),
+ /*in_ctrl_values=*/mlir::ValueRange());
+
+ return newGate;
+ }
+
+ // function to create a two qubit gate with a given name
+ // right after the SWAP is to be erased
+ quantum::CustomOp createSimpleTwoBitGate(StringRef gateName, const Value &controlIn,
+ const Value &targetIn, const Value &controlOut,
+ const Value &targetOut, mlir::IRRewriter &builder,
+ Location &loc,
+ const quantum::CustomOp &insert_after_gate)
+ {
+ OpBuilder::InsertionGuard insertionGuard(builder);
+ builder.setInsertionPointAfter(insert_after_gate);
+ quantum::CustomOp newGate = builder.create(
+ loc,
+ /*out_qubits=*/mlir::TypeRange({controlOut.getType(), targetOut.getType()}),
+ /*out_ctrl_qubits=*/mlir::TypeRange({}),
+ /*params=*/mlir::ValueRange(),
+ /*in_qubits=*/mlir::ValueRange({controlIn, targetIn}),
+ /*gate_name=*/gateName,
+ /*adjoint=*/nullptr,
+ /*in_ctrl_qubits=*/mlir::ValueRange({}),
+ /*in_ctrl_values=*/mlir::ValueRange());
+
+ return newGate;
+ }
+
+ // above function overloaded to create a two qubit gate with a given name
+ // if multiple gates are to be inserted after SWAP transformation
+ quantum::CustomOp createSimpleTwoBitGate(StringRef gateName, const Value &controlIn,
+ const Value &targetIn, mlir::IRRewriter &builder,
+ Location &loc,
+ const quantum::CustomOp &insert_after_gate)
+ {
+ OpBuilder::InsertionGuard insertionGuard(builder);
+ builder.setInsertionPointAfter(insert_after_gate);
+ quantum::CustomOp newGate = builder.create(
+ loc,
+ /*out_qubits=*/mlir::TypeRange({controlIn.getType(), targetIn.getType()}),
+ /*out_ctrl_qubits=*/mlir::TypeRange({}),
+ /*params=*/mlir::ValueRange(),
+ /*in_qubits=*/mlir::ValueRange({controlIn, targetIn}),
+ /*gate_name=*/gateName,
+ /*adjoint=*/nullptr,
+ /*in_ctrl_qubits=*/mlir::ValueRange({}),
+ /*in_ctrl_values=*/mlir::ValueRange());
+
+ return newGate;
+ }
+
+ bool canScheduleOn(RegisteredOperationName opInfo) const override
+ {
+ return opInfo.hasInterface();
+ }
+
+ void runOnOperation() override
+ {
+ FunctionOpInterface func = cast(getOperation());
+ mlir::IRRewriter builder(func->getContext());
+ Location loc = func->getLoc();
+
+ PropagateSimpleStatesAnalysis &pssa = getAnalysis();
+ llvm::DenseMap qubitValues = pssa.getQubitValues();
+
+ func->walk([&](quantum::CustomOp op) {
+ StringRef gate = op.getGateName();
+ if (gate != "SWAP") {
+ return;
+ }
+
+ Value SwapQubit_0_In = op->getOperand(0);
+ Value SwapQubit_1_In = op->getOperand(1);
+ Value SwapQubit_0_Out = op->getResult(0);
+ Value SwapQubit_1_Out = op->getResult(1);
+
+ // first qubit in |0>
+ if (pssa.isZero(qubitValues[SwapQubit_0_In])) {
+ // second qubit in |0>: SWAP(|0>,|0>)
+ if (pssa.isZero(qubitValues[SwapQubit_1_In])) {
+ SwapQubit_0_Out.replaceAllUsesWith(SwapQubit_0_In);
+ SwapQubit_1_Out.replaceAllUsesWith(SwapQubit_1_In);
+ op->erase();
+ return;
+ }
+ // second qubit in |1>: SWAP(|0>,|1>)
+ else if (pssa.isOne(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp xgate_on_0 = createSimpleOneBitGate(
+ "PauliX", SwapQubit_0_In, SwapQubit_0_Out, builder, loc, op);
+ SwapQubit_0_Out.replaceAllUsesWith(xgate_on_0->getResult(0));
+
+ quantum::CustomOp xgate_on_1 = createSimpleOneBitGate(
+ "PauliX", SwapQubit_1_In, SwapQubit_1_Out, builder, loc, xgate_on_0);
+ SwapQubit_1_Out.replaceAllUsesWith(xgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |+>: SWAP(|0>,|+>)
+ else if (pssa.isPlus(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp hgate_on_0 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_0_In, SwapQubit_0_Out, builder, loc, op);
+ SwapQubit_0_Out.replaceAllUsesWith(hgate_on_0->getResult(0));
+
+ quantum::CustomOp hgate_on_1 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_1_In, SwapQubit_1_Out, builder, loc, hgate_on_0);
+ SwapQubit_1_Out.replaceAllUsesWith(hgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |->: SWAP(|0>,|->)
+ else if (pssa.isMinus(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp xgate_on_0 =
+ createSimpleOneBitGate("PauliX", SwapQubit_0_In, builder, loc, op);
+
+ quantum::CustomOp hgate_on_0 = createSimpleOneBitGate(
+ "Hadamard", xgate_on_0->getResult(0), builder, loc, xgate_on_0);
+ SwapQubit_0_Out.replaceAllUsesWith(hgate_on_0->getResult(0));
+
+ quantum::CustomOp hgate_on_1 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_1_In, builder, loc, hgate_on_0);
+
+ quantum::CustomOp xgate_on_1 = createSimpleOneBitGate(
+ "PauliX", hgate_on_1->getResult(0), builder, loc, hgate_on_1);
+ SwapQubit_1_Out.replaceAllUsesWith(xgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in NON_BASIS: SWAP(|0>,|NON_BASIS>)
+ else if (pssa.isOther(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp cnot_on_1_0 = createSimpleTwoBitGate(
+ "CNOT", SwapQubit_1_In, SwapQubit_0_In, builder, loc, op);
+
+ quantum::CustomOp cnot_on_0_1 = createSimpleTwoBitGate(
+ "CNOT", cnot_on_1_0->getResult(1), cnot_on_1_0->getResult(0), builder, loc,
+ cnot_on_1_0);
+
+ SwapQubit_0_Out.replaceAllUsesWith(cnot_on_0_1->getResult(0));
+ SwapQubit_1_Out.replaceAllUsesWith(cnot_on_0_1->getResult(1));
+ op->erase();
+ return;
+ }
+ }
+
+ // first qubit in |1>
+ else if (pssa.isOne(qubitValues[SwapQubit_0_In])) {
+ // second qubit in |0>: SWAP(|1>,|0>)
+ if (pssa.isZero(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp xgate_on_0 = createSimpleOneBitGate(
+ "PauliX", SwapQubit_0_In, SwapQubit_0_Out, builder, loc, op);
+ SwapQubit_0_Out.replaceAllUsesWith(xgate_on_0->getResult(0));
+
+ quantum::CustomOp xgate_on_1 = createSimpleOneBitGate(
+ "PauliX", SwapQubit_1_In, SwapQubit_1_Out, builder, loc, xgate_on_0);
+ SwapQubit_1_Out.replaceAllUsesWith(xgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |1>: SWAP(|1>,|1>)
+ else if (pssa.isOne(qubitValues[SwapQubit_1_In])) {
+ SwapQubit_0_Out.replaceAllUsesWith(SwapQubit_0_In);
+ SwapQubit_1_Out.replaceAllUsesWith(SwapQubit_1_In);
+ op->erase();
+ return;
+ }
+ // second qubit in |+>: SWAP(|1>,|+>)
+ else if (pssa.isPlus(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp xgate_on_0 =
+ createSimpleOneBitGate("PauliX", SwapQubit_0_In, builder, loc, op);
+
+ quantum::CustomOp hgate_on_0 = createSimpleOneBitGate(
+ "Hadamard", xgate_on_0->getResult(0), builder, loc, xgate_on_0);
+ SwapQubit_0_Out.replaceAllUsesWith(hgate_on_0->getResult(0));
+
+ quantum::CustomOp hgate_on_1 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_1_In, builder, loc, hgate_on_0);
+
+ quantum::CustomOp xgate_on_1 = createSimpleOneBitGate(
+ "PauliX", hgate_on_1->getResult(0), builder, loc, hgate_on_1);
+ SwapQubit_1_Out.replaceAllUsesWith(xgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |->: SWAP(|1>,|->)
+ else if (pssa.isMinus(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp hgate_on_0 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_0_In, SwapQubit_0_Out, builder, loc, op);
+ SwapQubit_0_Out.replaceAllUsesWith(hgate_on_0->getResult(0));
+
+ quantum::CustomOp hgate_on_1 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_1_In, SwapQubit_1_Out, builder, loc, hgate_on_0);
+ SwapQubit_1_Out.replaceAllUsesWith(hgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |NON_BASIS>: SWAP(|1>,|NON_BASIS>)
+ else if (pssa.isOther(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp xgate_on_1 =
+ createSimpleOneBitGate("PauliX", SwapQubit_1_In, builder, loc, op);
+
+ quantum::CustomOp cnot_on_1_0 = createSimpleTwoBitGate(
+ "CNOT", xgate_on_1->getResult(0), SwapQubit_0_In, builder, loc, xgate_on_1);
+
+ quantum::CustomOp cnot_on_0_1 = createSimpleTwoBitGate(
+ "CNOT", cnot_on_1_0->getResult(1), cnot_on_1_0->getResult(0), builder, loc,
+ cnot_on_1_0);
+
+ SwapQubit_0_Out.replaceAllUsesWith(cnot_on_0_1->getResult(0));
+ SwapQubit_1_Out.replaceAllUsesWith(cnot_on_0_1->getResult(1));
+ op->erase();
+ return;
+ }
+ }
+
+ // first qubit in |+>
+ else if (pssa.isPlus(qubitValues[SwapQubit_0_In])) {
+ // second qubit in |0>: SWAP(|+>,|0>)
+ if (pssa.isZero(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp hgate_on_0 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_0_In, SwapQubit_0_Out, builder, loc, op);
+ SwapQubit_0_Out.replaceAllUsesWith(hgate_on_0->getResult(0));
+
+ quantum::CustomOp hgate_on_1 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_1_In, SwapQubit_1_Out, builder, loc, hgate_on_0);
+ SwapQubit_1_Out.replaceAllUsesWith(hgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |01>: SWAP(|+>,|1>)
+ else if (pssa.isOne(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp hgate_on_0 =
+ createSimpleOneBitGate("Hadamard", SwapQubit_0_In, builder, loc, op);
+
+ quantum::CustomOp xgate_on_0 = createSimpleOneBitGate(
+ "PauliX", hgate_on_0->getResult(0), builder, loc, hgate_on_0);
+ SwapQubit_0_Out.replaceAllUsesWith(xgate_on_0->getResult(0));
+
+ quantum::CustomOp xgate_on_1 =
+ createSimpleOneBitGate("PauliX", SwapQubit_1_In, builder, loc, xgate_on_0);
+
+ quantum::CustomOp hgate_on_1 = createSimpleOneBitGate(
+ "Hadamard", xgate_on_1->getResult(0), builder, loc, xgate_on_1);
+ SwapQubit_1_Out.replaceAllUsesWith(hgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |+>: SWAP(|+>,|+>)
+ else if (pssa.isPlus(qubitValues[SwapQubit_1_In])) {
+ SwapQubit_0_Out.replaceAllUsesWith(SwapQubit_0_In);
+ SwapQubit_1_Out.replaceAllUsesWith(SwapQubit_1_In);
+ op->erase();
+ return;
+ }
+ // second qubit in |->: SWAP(|+>,|->)
+ else if (pssa.isMinus(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp zgate_on_0 = createSimpleOneBitGate(
+ "PauliZ", SwapQubit_0_In, SwapQubit_0_Out, builder, loc, op);
+ SwapQubit_0_Out.replaceAllUsesWith(zgate_on_0->getResult(0));
+
+ quantum::CustomOp zgate_on_1 = createSimpleOneBitGate(
+ "PauliZ", SwapQubit_1_In, SwapQubit_1_Out, builder, loc, zgate_on_0);
+ SwapQubit_1_Out.replaceAllUsesWith(zgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |NON_BASIS>: SWAP(|+>,|NON_BASIS>)
+ else if (pssa.isOther(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp cnot_on_0_1 = createSimpleTwoBitGate(
+ "CNOT", SwapQubit_0_In, SwapQubit_1_In, builder, loc, op);
+
+ quantum::CustomOp cnot_on_1_0 = createSimpleTwoBitGate(
+ "CNOT", cnot_on_0_1->getResult(1), cnot_on_0_1->getResult(0), builder, loc,
+ cnot_on_0_1);
+
+ SwapQubit_0_Out.replaceAllUsesWith(cnot_on_1_0->getResult(1));
+ SwapQubit_1_Out.replaceAllUsesWith(cnot_on_1_0->getResult(0));
+ op->erase();
+ return;
+ }
+ }
+
+ // first qubit in |->
+ else if (pssa.isMinus(qubitValues[SwapQubit_0_In])) {
+ // second qubit in |0>: SWAP(|->,|0>)
+ if (pssa.isZero(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp hgate_on_0 =
+ createSimpleOneBitGate("Hadamard", SwapQubit_0_In, builder, loc, op);
+
+ quantum::CustomOp xgate_on_0 = createSimpleOneBitGate(
+ "PauliX", hgate_on_0->getResult(0), builder, loc, hgate_on_0);
+ SwapQubit_0_Out.replaceAllUsesWith(xgate_on_0->getResult(0));
+
+ quantum::CustomOp xgate_on_1 =
+ createSimpleOneBitGate("PauliX", SwapQubit_1_In, builder, loc, xgate_on_0);
+
+ quantum::CustomOp hgate_on_1 = createSimpleOneBitGate(
+ "Hadamard", xgate_on_1->getResult(0), builder, loc, xgate_on_1);
+ SwapQubit_1_Out.replaceAllUsesWith(hgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |1>: SWAP(|->,|1>)
+ else if (pssa.isOne(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp hgate_on_0 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_0_In, SwapQubit_0_Out, builder, loc, op);
+ SwapQubit_0_Out.replaceAllUsesWith(hgate_on_0->getResult(0));
+
+ quantum::CustomOp hgate_on_1 = createSimpleOneBitGate(
+ "Hadamard", SwapQubit_1_In, SwapQubit_1_Out, builder, loc, hgate_on_0);
+ SwapQubit_1_Out.replaceAllUsesWith(hgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |+>: SWAP(|->,|+>)
+ else if (pssa.isPlus(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp zgate_on_0 = createSimpleOneBitGate(
+ "PauliZ", SwapQubit_0_In, SwapQubit_0_Out, builder, loc, op);
+ SwapQubit_0_Out.replaceAllUsesWith(zgate_on_0->getResult(0));
+
+ quantum::CustomOp zgate_on_1 = createSimpleOneBitGate(
+ "PauliZ", SwapQubit_1_In, SwapQubit_1_Out, builder, loc, zgate_on_0);
+ SwapQubit_1_Out.replaceAllUsesWith(zgate_on_1->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |->: SWAP(|->,|->)
+ else if (pssa.isMinus(qubitValues[SwapQubit_1_In])) {
+ SwapQubit_0_Out.replaceAllUsesWith(SwapQubit_0_In);
+ SwapQubit_1_Out.replaceAllUsesWith(SwapQubit_1_In);
+ op->erase();
+ return;
+ }
+ // second qubit in |NON_BASIS>: SWAP(|->,|NON_BASIS>)
+ else if (pssa.isOther(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp zgate_on_1 =
+ createSimpleOneBitGate("PauliZ", SwapQubit_1_In, builder, loc, op);
+
+ quantum::CustomOp cnot_on_0_1 = createSimpleTwoBitGate(
+ "CNOT", SwapQubit_0_In, zgate_on_1->getResult(0), builder, loc, zgate_on_1);
+
+ quantum::CustomOp cnot_on_1_0 = createSimpleTwoBitGate(
+ "CNOT", cnot_on_0_1->getResult(1), cnot_on_0_1->getResult(0), builder, loc,
+ cnot_on_0_1);
+
+ SwapQubit_0_Out.replaceAllUsesWith(cnot_on_1_0->getResult(1));
+ SwapQubit_1_Out.replaceAllUsesWith(cnot_on_1_0->getResult(0));
+ op->erase();
+ return;
+ }
+ }
+
+ // first qubit in |NON_BASIS>
+ else if (pssa.isOther(qubitValues[SwapQubit_0_In])) {
+ // second qubit in |0>: SWAP(|NON_BASIS>,|0>)
+ if (pssa.isZero(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp cnot_on_0_1 = createSimpleTwoBitGate(
+ "CNOT", SwapQubit_0_In, SwapQubit_1_In, builder, loc, op);
+
+ quantum::CustomOp cnot_on_1_0 = createSimpleTwoBitGate(
+ "CNOT", cnot_on_0_1->getResult(1), cnot_on_0_1->getResult(0), builder, loc,
+ cnot_on_0_1);
+
+ SwapQubit_0_Out.replaceAllUsesWith(cnot_on_1_0->getResult(1));
+ SwapQubit_1_Out.replaceAllUsesWith(cnot_on_1_0->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |1>: SWAP(|NON_BASIS>,|1>)
+ else if (pssa.isOne(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp xgate_on_0 =
+ createSimpleOneBitGate("PauliX", SwapQubit_0_In, builder, loc, op);
+
+ quantum::CustomOp cnot_on_0_1 = createSimpleTwoBitGate(
+ "CNOT", xgate_on_0->getResult(0), SwapQubit_1_In, builder, loc, xgate_on_0);
+
+ quantum::CustomOp cnot_on_1_0 = createSimpleTwoBitGate(
+ "CNOT", cnot_on_0_1->getResult(1), cnot_on_0_1->getResult(0), builder, loc,
+ cnot_on_0_1);
+
+ SwapQubit_0_Out.replaceAllUsesWith(cnot_on_1_0->getResult(1));
+ SwapQubit_1_Out.replaceAllUsesWith(cnot_on_1_0->getResult(0));
+ op->erase();
+ return;
+ }
+ // second qubit in |+>: SWAP(|NON_BASIS>,|+>)
+ else if (pssa.isPlus(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp cnot_on_1_0 = createSimpleTwoBitGate(
+ "CNOT", SwapQubit_1_In, SwapQubit_0_In, builder, loc, op);
+
+ quantum::CustomOp cnot_on_0_1 = createSimpleTwoBitGate(
+ "CNOT", cnot_on_1_0->getResult(1), cnot_on_1_0->getResult(0), builder, loc,
+ cnot_on_1_0);
+
+ SwapQubit_0_Out.replaceAllUsesWith(cnot_on_0_1->getResult(0));
+ SwapQubit_1_Out.replaceAllUsesWith(cnot_on_0_1->getResult(1));
+ op->erase();
+ return;
+ }
+ // second qubit in |->: SWAP(|NON_BASIS>,|->)
+ else if (pssa.isMinus(qubitValues[SwapQubit_1_In])) {
+ quantum::CustomOp zgate_on_0 =
+ createSimpleOneBitGate("PauliZ", SwapQubit_0_In, builder, loc, op);
+
+ quantum::CustomOp cnot_on_1_0 = createSimpleTwoBitGate(
+ "CNOT", SwapQubit_1_In, zgate_on_0->getResult(0), builder, loc, zgate_on_0);
+
+ quantum::CustomOp cnot_on_0_1 = createSimpleTwoBitGate(
+ "CNOT", cnot_on_1_0->getResult(1), cnot_on_1_0->getResult(0), builder, loc,
+ cnot_on_1_0);
+
+ SwapQubit_0_Out.replaceAllUsesWith(cnot_on_0_1->getResult(0));
+ SwapQubit_1_Out.replaceAllUsesWith(cnot_on_0_1->getResult(1));
+ op->erase();
+ return;
+ }
+ }
+ });
+ }
+};
+
+std::unique_ptr createDisentangleSWAPPass()
+{
+ return std::make_unique();
+}
+
+} // namespace catalyst
diff --git a/mlir/lib/Quantum/Transforms/PropagateSimpleStatesAnalysis.hpp b/mlir/lib/Quantum/Transforms/PropagateSimpleStatesAnalysis.hpp
new file mode 100644
index 0000000000..02efa45246
--- /dev/null
+++ b/mlir/lib/Quantum/Transforms/PropagateSimpleStatesAnalysis.hpp
@@ -0,0 +1,241 @@
+// Copyright 2024 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.
+
+// This algorithm is taken from https://arxiv.org/pdf/2012.07711, figure 5
+
+#pragma once
+
+#include