Add merge rotation pass

doc/releases/changelog-dev.md

@@ -209,6 +209,8 @@
 * Samples on lightning.qubit/kokkos can now be seeded with `qjit(seed=...)`.
   [(#1164)](https://github.com/PennyLaneAI/catalyst/pull/1164)
 
+* The compiler pass `-remove-chained-self-inverse` can now also cancel adjoints of arbitrary unitary operations (in addition to the named Hermitian gates).
+  [(#1186)](https://github.com/PennyLaneAI/catalyst/pull/1186)
 
 <h3>Breaking changes</h3>
 

mlir/include/Quantum/Transforms/Passes.h

@@ -28,6 +28,6 @@ std::unique_ptr<mlir::Pass> createAdjointLoweringPass();
 std::unique_ptr<mlir::Pass> createRemoveChainedSelfInversePass();
 std::unique_ptr<mlir::Pass> createAnnotateFunctionPass();
 std::unique_ptr<mlir::Pass> createSplitMultipleTapesPass();
-std::unique_ptr<mlir::Pass> createMergeRotationPass();
+std::unique_ptr<mlir::Pass> createMergeRotationsPass();
 
 } // namespace catalyst

mlir/include/Quantum/Transforms/Passes.td

@@ -110,10 +110,10 @@ def RemoveChainedSelfInversePass : Pass<"remove-chained-self-inverse"> {
     let options = QuantumCircuitTransformationPass.options;
 }
 
-def MergeRotationPass : Pass<"merge-rotation"> {
+def MergeRotationsPass : Pass<"merge-rotations"> {
     let summary = "merge rotation boilerplate words";
 
-    let constructor = "catalyst::createMergeRotationPass()";
+    let constructor = "catalyst::createMergeRotationsPass()";
     let options = !listconcat(
     QuantumCircuitTransformationPass.options,
     [

mlir/include/Quantum/Transforms/Patterns.h

@@ -26,6 +26,7 @@ void populateBufferizationPatterns(mlir::TypeConverter &, mlir::RewritePatternSe
 void populateQIRConversionPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
 void populateAdjointPatterns(mlir::RewritePatternSet &);
 void populateSelfInversePatterns(mlir::RewritePatternSet &);
+void populateMergeRotationsPatterns(mlir::RewritePatternSet &);
 
 } // namespace quantum
 } // namespace catalyst

mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp

@@ -41,13 +41,14 @@ void catalyst::registerAllCatalystPasses()
     mlir::registerPass(catalyst::createInlineNestedModulePass);
     mlir::registerPass(catalyst::createMemrefCopyToLinalgCopyPass);
     mlir::registerPass(catalyst::createMemrefToLLVMWithTBAAPass);
-    mlir::registerPass(catalyst::createMergeRotationPass);
+    mlir::registerPass(catalyst::createMergeRotationsPass);
     mlir::registerPass(catalyst::createMitigationLoweringPass);
     mlir::registerPass(catalyst::createQnodeToAsyncLoweringPass);
     mlir::registerPass(catalyst::createQuantumBufferizationPass);
     mlir::registerPass(catalyst::createQuantumConversionPass);
     mlir::registerPass(catalyst::createRegisterInactiveCallbackPass);
     mlir::registerPass(catalyst::createRemoveChainedSelfInversePass);
+    mlir::registerPass(catalyst::createMergeRotationsPass);
     mlir::registerPass(catalyst::createScatterLoweringPass);
     mlir::registerPass(catalyst::createSplitMultipleTapesPass);
     mlir::registerPass(catalyst::createTestPass);

mlir/lib/Quantum/Transforms/CMakeLists.txt

@@ -14,6 +14,7 @@ file(GLOB SRC
     remove_chained_self_inverse.cpp
     SplitMultipleTapes.cpp
     merge_rotation.cpp
+    MergeRotationsPatterns.cpp
 )
 
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)

mlir/lib/Quantum/Transforms/ChainedSelfInversePatterns.cpp

@@ -14,11 +14,15 @@
 
 #define DEBUG_TYPE "chained-self-inverse"
 
-#include "Quantum/IR/QuantumOps.h"
-#include "Quantum/Transforms/Patterns.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Errc.h"
+
+#include "Quantum/IR/QuantumOps.h"
+#include "Quantum/Transforms/Patterns.h"
+
+#include "VerifyParentGateAnalysis.hpp"
+
 using llvm::dbgs;
 using namespace mlir;
 using namespace catalyst;
@@ -29,7 +33,7 @@ static const mlir::StringSet<> HermitianOps = {"Hadamard", "PauliX", "PauliY", "
 
 namespace {
 
-struct ChainedHadamardOpRewritePattern : public mlir::OpRewritePattern<CustomOp> {
+struct ChainedNamedHermitianOpRewritePattern : public mlir::OpRewritePattern<CustomOp> {
     using mlir::OpRewritePattern<CustomOp>::OpRewritePattern;
 
     /// We simplify consecutive Hermitian quantum gates by removing them.
@@ -41,22 +45,90 @@ struct ChainedHadamardOpRewritePattern : public mlir::OpRewritePattern<CustomOp>
         LLVM_DEBUG(dbgs() << "Simplifying the following operation:\n" << op << "\n");
 
         StringRef OpGateName = op.getGateName();
-        if (!HermitianOps.contains(OpGateName))
+        if (!HermitianOps.contains(OpGateName)) {
             return failure();
+        }
+
+        VerifyParentGateAndNameAnalysis<CustomOp> vpga(op);
+        if (!vpga.getVerifierResult()) {
+            return failure();
+        }
 
+        // Replace uses
         ValueRange InQubits = op.getInQubits();
         auto ParentOp = dyn_cast_or_null<CustomOp>(InQubits[0].getDefiningOp());
-        if (!ParentOp || ParentOp.getGateName() != OpGateName)
+        ValueRange simplifiedVal = ParentOp.getInQubits();
+        rewriter.replaceOp(op, simplifiedVal);
+        return success();
+    }
+};
+
+template <typename OpType>
+struct ChainedUUadjOpRewritePattern : public mlir::OpRewritePattern<OpType> {
+    using mlir::OpRewritePattern<OpType>::OpRewritePattern;
+
+    bool verifyParentGateParams(OpType op, OpType parentOp) const
+    {
+        // Verify that the parent gate has the same parameters
+        ValueRange opParams = op.getAllParams();
+        ValueRange parentOpParams = parentOp.getAllParams();
+
+        if (opParams.size() != parentOpParams.size()) {
+            return false;
+        }
+
+        for (auto [opParam, parentOpParam] : llvm::zip(opParams, parentOpParams)) {
+            if (opParam != parentOpParam) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    bool verifyOneAdjoint(OpType op, OpType parentOp) const
+    {
+        // Verify that exactly one of the neighbouring pair is an adjoint
+        bool opIsAdj = op->hasAttr("adjoint");
+        bool parentIsAdj = parentOp->hasAttr("adjoint");
+        return opIsAdj != parentIsAdj; // "XOR" to check just one true
+    }
+
+    /// Remove generic neighbouring gate pairs of the form
+    /// --- gate --- gate{adjoint} ---
+    /// Conditions:
+    ///  1. Parent gate verification must pass. See VerifyParentGateAnalysis.hpp.
+    ///  2. If there are parameters, both gate must have the same parameters.
+    ///     [This pattern assumes the IR is already processed by CSE]
+    mlir::LogicalResult matchAndRewrite(OpType op, mlir::PatternRewriter &rewriter) const override
+    {
+        LLVM_DEBUG(dbgs() << "Simplifying the following operation:\n" << op << "\n");
+
+        VerifyParentGateAndNameAnalysis<OpType> vpga(op);
+        if (!vpga.getVerifierResult()) {
+            return failure();
+        }
+
+        ValueRange InQubits = op.getInQubits();
+        auto parentOp = dyn_cast_or_null<OpType>(InQubits[0].getDefiningOp());
+
+        if (!verifyParentGateParams(op, parentOp)) {
             return failure();
+        }
 
-        ValueRange ParentOutQubits = ParentOp.getOutQubits();
-        // Check if the input qubits to the current operation match the output qubits of the parent.
-        for (const auto &[Idx, Qubit] : llvm::enumerate(InQubits)) {
-            if (Qubit.getDefiningOp<CustomOp>() != ParentOp || Qubit != ParentOutQubits[Idx])
-                return failure();
+        if (!verifyOneAdjoint(op, parentOp)) {
+            return failure();
+        }
+
+        // Replace uses
+        ValueRange originalNonCtrlQubits = parentOp.getNonCtrlQubitOperands();
+        ValueRange originalCtrlQubits = parentOp.getCtrlQubitOperands();
+        for (const auto &[idx, nonCtrlQubitResult] : llvm::enumerate(op.getNonCtrlQubitResults())) {
+            nonCtrlQubitResult.replaceAllUsesWith(originalNonCtrlQubits[idx]);
+        }
+        for (const auto &[idx, ctrlQubitResult] : llvm::enumerate(op.getCtrlQubitResults())) {
+            ctrlQubitResult.replaceAllUsesWith(originalCtrlQubits[idx]);
         }
-        ValueRange simplifiedVal = ParentOp.getInQubits();
-        rewriter.replaceOp(op, simplifiedVal);
         return success();
     }
 };
@@ -68,7 +140,15 @@ namespace quantum {
 
 void populateSelfInversePatterns(RewritePatternSet &patterns)
 {
-    patterns.add<ChainedHadamardOpRewritePattern>(patterns.getContext(), 1);
+    patterns.add<ChainedNamedHermitianOpRewritePattern>(patterns.getContext(), 1);
+
+    // TODO: better organize the quantum dialect
+    // There is an interface `QuantumGate` for all the unitary gate operations,
+    // but interfaces cannot be accepted by pattern matchers, since pattern
+    // matchers require the target operations to have concrete names in the IR.
+    patterns.add<ChainedUUadjOpRewritePattern<CustomOp>>(patterns.getContext(), 1);
+    patterns.add<ChainedUUadjOpRewritePattern<QubitUnitaryOp>>(patterns.getContext(), 1);
+    patterns.add<ChainedUUadjOpRewritePattern<MultiRZOp>>(patterns.getContext(), 1);
 }
 
 } // namespace quantum

mlir/lib/Quantum/Transforms/MergeRotationsPatterns.cpp

@@ -0,0 +1,131 @@
+// 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.
+
+#define DEBUG_TYPE "merge-rotations"
+
+#include "Quantum/IR/QuantumOps.h"
+#include "Quantum/Transforms/Patterns.h"
+#include "VerifyParentGateAnalysis.hpp"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "llvm/ADT/StringSet.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/Errc.h"
+
+using llvm::dbgs;
+using namespace mlir;
+using namespace catalyst::quantum;
+
+static const mlir::StringSet<> rotationsSet = {"RX",  "RY",  "RZ",  "PhaseShift",           "Rot",
+                                               "CRX", "CRY", "CRZ", "ControlledPhaseShift", "CRot"};
+
+namespace {
+
+struct MergeRotationsRewritePattern : public mlir::OpRewritePattern<CustomOp> {
+    using mlir::OpRewritePattern<CustomOp>::OpRewritePattern;
+
+    mlir::LogicalResult matchAndRewrite(CustomOp op, mlir::PatternRewriter &rewriter) const override
+    {
+        LLVM_DEBUG(dbgs() << "Simplifying the following operation:\n" << op << "\n");
+        auto loc = op.getLoc();
+        StringRef opGateName = op.getGateName();
+
+        if (!rotationsSet.contains(opGateName))
+            return failure();
+        ValueRange inQubits = op.getInQubits();
+        auto parentOp = dyn_cast_or_null<CustomOp>(inQubits[0].getDefiningOp());
+
+        VerifyParentGateAndNameAnalysis<CustomOp> vpga(op);
+        if (!vpga.getVerifierResult()) {
+            return failure();
+        }
+
+        TypeRange outQubitsTypes = op.getOutQubits().getTypes();
+        TypeRange outQubitsCtrlTypes = op.getOutCtrlQubits().getTypes();
+        ValueRange parentInQubits = parentOp.getInQubits();
+        ValueRange parentInCtrlQubits = parentOp.getInCtrlQubits();
+        ValueRange parentInCtrlValues = parentOp.getInCtrlValues();
+
+        auto parentParams = parentOp.getParams();
+        auto params = op.getParams();
+        SmallVector<mlir::Value> sumParams;
+        for (auto [param, parentParam] : llvm::zip(params, parentParams)) {
+            mlir::Value sumParam =
+                rewriter.create<arith::AddFOp>(loc, parentParam, param).getResult();
+            sumParams.push_back(sumParam);
+        };
+        auto mergeOp = rewriter.create<CustomOp>(loc, outQubitsTypes, outQubitsCtrlTypes, sumParams,
+                                                 parentInQubits, opGateName, nullptr,
+                                                 parentInCtrlQubits, parentInCtrlValues);
+
+        op.replaceAllUsesWith(mergeOp);
+        op.erase();
+        parentOp.erase();
+
+        return success();
+    }
+};
+
+struct MergeMultiRZRewritePattern : public mlir::OpRewritePattern<MultiRZOp> {
+    using mlir::OpRewritePattern<MultiRZOp>::OpRewritePattern;
+
+    mlir::LogicalResult matchAndRewrite(MultiRZOp op,
+                                        mlir::PatternRewriter &rewriter) const override
+    {
+        LLVM_DEBUG(dbgs() << "Simplifying the following operation:\n" << op << "\n");
+        auto loc = op.getLoc();
+
+        VerifyParentGateAnalysis<MultiRZOp> vpga(op);
+        if (!vpga.getVerifierResult()) {
+            return failure();
+        }
+
+        ValueRange inQubits = op.getInQubits();
+        auto parentOp = dyn_cast_or_null<MultiRZOp>(inQubits[0].getDefiningOp());
+        if (!parentOp)
+            return failure();
+
+        TypeRange outQubitsTypes = op.getOutQubits().getTypes();
+        TypeRange outQubitsCtrlTypes = op.getOutCtrlQubits().getTypes();
+        ValueRange parentInQubits = parentOp.getInQubits();
+        ValueRange parentInCtrlQubits = parentOp.getInCtrlQubits();
+        ValueRange parentInCtrlValues = parentOp.getInCtrlValues();
+
+        auto parentTheta = parentOp.getTheta();
+        auto theta = op.getTheta();
+
+        mlir::Value sumParam = rewriter.create<arith::AddFOp>(loc, parentTheta, theta).getResult();
+
+        auto mergeOp = rewriter.create<MultiRZOp>(loc, outQubitsTypes, outQubitsCtrlTypes,
+                                                  sumParam, parentInQubits, nullptr,
+                                                  parentInCtrlQubits, parentInCtrlValues);
+        op.replaceAllUsesWith(mergeOp);
+        op.erase();
+        parentOp.erase();
+
+        return success();
+    }
+};
+} // namespace
+
+namespace catalyst {
+namespace quantum {
+
+void populateMergeRotationsPatterns(RewritePatternSet &patterns)
+{
+    patterns.add<MergeRotationsRewritePattern>(patterns.getContext(), 1);
+    patterns.add<MergeMultiRZRewritePattern>(patterns.getContext(), 1);
+}
+
+} // namespace quantum
+} // namespace catalyst