Create a finite state machine analysis pass to propagate Pauli basis …

…states, and use it to disentangle CNOTs and SWAPs (#1154)

**Context:**
In circuits that mostly involve the six Pauli eigenstates, the effect of
gates can be described by a simple six-state finite state machine. This
allows us to deduce their midcircuit state at compile time.

The knowledge of the snapshot midcircuit states can be useful when:
- Deducing CNOT/Toffoli/SWAP/... effect and decouple them. We decouple
CNOT and SWAP in this PR.
- Replace complicated circuits whose results are |01+-LR> with standard
preparation routines
  - Potentially more

**Description of the Change:**
Added an analysis for propagating simple states (and tests).
Also added CNOT and SWAP decomposition pass `disentangle-CNOT`,
`disentangle-SWAP` (and tests).
Credit to [Ritu Thombre](https://github.com/ritu-thombre99) for the work
on the SWAP gates!

**Benefits:**
More compile time circuit information for potential compile time
transformation/optimization.

**Possible Drawbacks:**
This algorithm is from an external paper. We need to properly cite it. 
https://arxiv.org/abs/2012.07711 (fig.5)

**Related GitHub Issues:** closes #1268 

[sc-74649]
[sc-74650]

---------

Co-authored-by: Ritu Thombre <[email protected]>
Co-authored-by: Ritu Thombre <[email protected]>

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).
+
 <h3>Breaking changes 💔</h3>
 
 * 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.

mlir/include/Quantum/Transforms/Passes.h

@@ -29,6 +29,8 @@ std::unique_ptr<mlir::Pass> createRemoveChainedSelfInversePass();
 std::unique_ptr<mlir::Pass> createAnnotateFunctionPass();
 std::unique_ptr<mlir::Pass> createSplitMultipleTapesPass();
 std::unique_ptr<mlir::Pass> createMergeRotationsPass();
+std::unique_ptr<mlir::Pass> createDisentangleCNOTPass();
+std::unique_ptr<mlir::Pass> createDisentangleSWAPPass();
 std::unique_ptr<mlir::Pass> createIonsDecompositionPass();
 std::unique_ptr<mlir::Pass> createStaticCustomLoweringPass();
 

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

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);

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

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<DisentangleCNOTPass> {
+    using impl::DisentangleCNOTPassBase<DisentangleCNOTPass>::DisentangleCNOTPassBase;
+
+    bool canScheduleOn(RegisteredOperationName opInfo) const override
+    {
+        return opInfo.hasInterface<FunctionOpInterface>();
+    }
+
+    void runOnOperation() override
+    {
+        FunctionOpInterface func = cast<FunctionOpInterface>(getOperation());
+        mlir::IRRewriter builder(func->getContext());
+        Location loc = func->getLoc();
+
+        PropagateSimpleStatesAnalysis &pssa = getAnalysis<PropagateSimpleStatesAnalysis>();
+        llvm::DenseMap<Value, QubitState> 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<quantum::CustomOp>(
+                        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<quantum::CustomOp>(
+                        loc, /*gate_name=*/"PauliZ",
+                        /*in_qubits=*/mlir::ValueRange({controlIn}));
+                    controlOut.replaceAllUsesWith(zgate->getResult(0));
+                    op->erase();
+                    return;
+                }
+            }
+        });
+    }
+};
+
+std::unique_ptr<Pass> createDisentangleCNOTPass()
+{
+    return std::make_unique<DisentangleCNOTPass>();
+}
+
+} // namespace catalyst