Add support of MCM for decompose pass and some refactoring

.github/workflows/rc_sync.yaml

@@ -28,7 +28,7 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-      
+
       # Import Numpy/Pybind11 because it is imported in setup.py
       - name: Install Numpy and Pybind11
         run: |

doc/releases/changelog-0.13.0.md

@@ -19,7 +19,7 @@
   [(#2002)](https://github.com/PennyLaneAI/catalyst/pull/2002)
 
   Two new functions, ``qml.allocate()`` and ``qml.deallocate()``, [have been added to
-  PennyLane](https://docs.pennylane.ai/en/stable/development/release_notes.html#release-0-43-0) 
+  PennyLane](https://docs.pennylane.ai/en/stable/development/release_notes.html#release-0-43-0)
   to support dynamic wire allocation. With Catalyst, these features can be accessed on
    ``lightning.qubit``, ``lightning.kokkos``, and ``lightning.gpu``.
 
@@ -46,13 +46,13 @@
   ```
 
   In the above program, 2 qubits are allocated during device initialization, and 1
-  additional qubit is allocated inside the circuit with ``qml.allocate(1)``. 
+  additional qubit is allocated inside the circuit with ``qml.allocate(1)``.
 
   For more information on what ``qml.allocate`` and ``qml.deallocate`` do, please consult the
   [PennyLane v0.43 release notes](https://docs.pennylane.ai/en/stable/development/release_notes.html#release-0-43-0).
 
   However, there are some notable differences between the behaviour of these features
-  with ``qjit`` versus without. For details, please see the relevant sections in the 
+  with ``qjit`` versus without. For details, please see the relevant sections in the
   [Catalyst sharp bits page](https://docs.pennylane.ai/projects/catalyst/en/stable/dev/sharp_bits.html#functionality-differences-from-pennylane).
 
 * A new quantum compilation pass that reduces the depth and count of non-Clifford Pauli product
@@ -106,7 +106,7 @@
               qml.H(wires=i)
               qml.T(wires=i)
 
-          return [measure(wires=i) for i in range(n)] 
+          return [measure(wires=i) for i in range(n)]
 
       print(ppm_specs(circuit))
   ```
@@ -116,7 +116,7 @@
   {'circuit_0': {'depth_pi8_ppr': 4, 'depth_ppm': 1, 'logical_qubits': 3, 'max_weight_pi8': 3, 'num_of_ppm': 3, 'pi8_ppr': 6}}
   ```
 
-  After performing the :func:`~.passes.to_ppr`, :func:`~.passes.commute_ppr`, and :func:`~.passes.merge_ppr_ppm`, 
+  After performing the :func:`~.passes.to_ppr`, :func:`~.passes.commute_ppr`, and :func:`~.passes.merge_ppr_ppm`,
   passes, the circuit contains a depth of four of non-Clifford PPRs (`depth_pi8_ppr`). Subsequently applying the
   :func:`~.passes.t_layer_reduction` pass will move PPRs around via commutation, resulting in a circuit with a
   smaller PPR depth of three.
@@ -461,7 +461,7 @@ for example the one-shot mid circuit measurement transform.
   [(#2057)](https://github.com/PennyLaneAI/catalyst/pull/2057)
 
   This pass is part of a bottom-of-stack MBQC execution pathway, with a thin shim between the
-  PPR/PPM layer and MBQC to enable end-to-end compilation on a mocked backend.  Also, in an MBQC gate 
+  PPR/PPM layer and MBQC to enable end-to-end compilation on a mocked backend.  Also, in an MBQC gate
   set, one of the gate `RotXZX` cannot yet be executed on available backends.
 
   ```python
@@ -485,8 +485,8 @@ for example the one-shot mid circuit measurement transform.
 
 <h3>Documentation 📝</h3>
 
-* Typos were fixed and supplemental information was added to the 
-  docstrings for ``ppm_compilaion``, ``to_ppr``, ``commute_ppr``, 
+* Typos were fixed and supplemental information was added to the
+  docstrings for ``ppm_compilaion``, ``to_ppr``, ``commute_ppr``,
   ``ppr_to_ppm``, ``merge_ppr_ppm``, and ``ppm_specs``.
   [(#2050)](https://github.com/PennyLaneAI/catalyst/pull/2050)
 

frontend/test/pytest/from_plxpr/test_capture_integration.py

@@ -1313,6 +1313,38 @@ def circuit(x: float, y: float, z: float):
 
         assert jnp.allclose(circuit(1.5, 2.5, 3.5), capture_result)
 
+    def test_transform_graph_decompose_workflow(self, backend):
+        """Test the integration for a circuit with a 'decompose' graph transform."""
+
+        # Capture enabled
+
+        qml.capture.enable()
+        qml.decomposition.enable_graph()
+
+        @qjit(target="mlir")
+        @partial(qml.transforms.decompose, gate_set=[qml.RX, qml.RY, qml.RZ])
+        @qml.qnode(qml.device(backend, wires=2))
+        def captured_circuit(x: float, y: float, z: float):
+            qml.measure(0)
+            qml.Rot(x, y, z, 0)
+            return qml.expval(qml.PauliZ(0))
+
+        capture_result = captured_circuit(1.5, 2.5, 3.5)
+
+        qml.decomposition.disable_graph()
+        qml.capture.disable()
+
+        # Capture disabled
+        @qjit
+        @partial(qml.transforms.decompose, gate_set=[qml.RX, qml.RY, qml.RZ])
+        @qml.qnode(qml.device(backend, wires=2))
+        def circuit(x: float, y: float, z: float):
+            catalyst.measure(0)
+            qml.Rot(x, y, z, 0)
+            return qml.expval(qml.PauliZ(0))
+
+        assert jnp.allclose(circuit(1.5, 2.5, 3.5), capture_result)
+
     def test_transform_map_wires_workflow(self, backend):
         """Test the integration for a circuit with a 'map_wires' transform."""
 

mlir/lib/Quantum/Transforms/DecomposeLoweringPatterns.cpp

@@ -25,6 +25,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/ValueRange.h"
 
+#include "Quantum/IR/QuantumInterfaces.h"
 #include "Quantum/IR/QuantumOps.h"
 #include "Quantum/Transforms/Patterns.h"
 
@@ -40,6 +41,7 @@ namespace quantum {
 /// - A runtime Value (for dynamic indices computed at runtime)
 /// - An IntegerAttr (for compile-time constant indices)
 /// - Invalid/uninitialized (represented by std::monostate)
+/// And a qreg value to represent the qreg that the index belongs to
 ///
 /// The struct uses std::variant to ensure only one type is active at a time,
 /// preventing invalid states.
@@ -54,17 +56,21 @@ namespace quantum {
 ///       Value idx = dynamicIdx.getValue();   // Get the Value
 ///     }
 ///   }
-struct QubitIndex {
+class QubitIndex {
+  private:
     // use monostate to represent the invalid index
     std::variant<std::monostate, Value, IntegerAttr> index;
+    Value qreg;
 
-    QubitIndex() : index(std::monostate()) {}
-    QubitIndex(Value val) : index(val) {}
-    QubitIndex(IntegerAttr attr) : index(attr) {}
+  public:
+    QubitIndex() : index(std::monostate()), qreg(nullptr) {}
+    QubitIndex(Value val, Value qreg) : index(val), qreg(qreg) {}
+    QubitIndex(IntegerAttr attr, Value qreg) : index(attr), qreg(qreg) {}
 
     bool isValue() const { return std::holds_alternative<Value>(index); }
     bool isAttr() const { return std::holds_alternative<IntegerAttr>(index); }
     operator bool() const { return isValue() || isAttr(); }
+    Value getReg() const { return qreg; }
     Value getValue() const { return isValue() ? std::get<Value>(index) : nullptr; }
     IntegerAttr getAttr() const { return isAttr() ? std::get<IntegerAttr>(index) : nullptr; }
 };
@@ -76,25 +82,16 @@ class OpSignatureAnalyzer {
   public:
     OpSignatureAnalyzer() = delete;
     OpSignatureAnalyzer(CustomOp op, bool enableQregMode)
-        : signature(OpSignature{
-              .params = op.getParams(),
-              .inQubits = op.getInQubits(),
-              .inCtrlQubits = op.getInCtrlQubits(),
-              .inCtrlValues = op.getInCtrlValues(),
-              .outQubits = op.getOutQubits(),
-              .outCtrlQubits = op.getOutCtrlQubits(),
-          })
+        : signature(OpSignature{.params = op.getParams(),
+                                .inQubits = op.getNonCtrlQubitOperands(),
+                                .inCtrlQubits = op.getCtrlQubitOperands(),
+                                .inCtrlValues = op.getCtrlValueOperands(),
+                                .outQubits = op.getNonCtrlQubitResults(),
+                                .outCtrlQubits = op.getCtrlQubitResults()})
     {
         if (!enableQregMode)
             return;
 
-        signature.sourceQreg = getSourceQreg(signature.inQubits.front());
-        if (!signature.sourceQreg) {
-            op.emitError("Cannot get source qreg");
-            isValid = false;
-            return;
-        }
-
         // input wire indices
         for (Value qubit : signature.inQubits) {
             const QubitIndex index = getExtractIndex(qubit);
@@ -117,13 +114,23 @@ class OpSignatureAnalyzer {
             signature.inCtrlWireIndices.emplace_back(index);
         }
 
+        assert((signature.inWireIndices.size() + signature.inCtrlWireIndices.size()) > 0 &&
+               "inWireIndices or inCtrlWireIndices should not be empty");
+
         // Output qubit indices are the same as input qubit indices
         signature.outQubitIndices = signature.inWireIndices;
         signature.outCtrlQubitIndices = signature.inCtrlWireIndices;
     }
 
     operator bool() const { return isValid; }
 
+    Value getUpdatedQreg(PatternRewriter &rewriter, Location loc)
+    {
+        // FIXME: This will cause an issue when the decomposition function has cross-qreg
+        // inputs and outputs. Now, we just assume has only one qreg input, the global one exists.
+        return signature.inWireIndices[0].getReg();
+    }
+
     // Prepare the operands for calling the decomposition function
     // There are two cases:
     // 1. The first input is a qreg, which means the decomposition function is a qreg mode function
@@ -144,15 +151,8 @@ class OpSignatureAnalyzer {
 
         int operandIdx = 0;
         if (isa<quantum::QuregType>(funcInputs[0])) {
-            Value updatedQreg = signature.sourceQreg;
-            for (auto [i, qubit] : llvm::enumerate(signature.inQubits)) {
-                const QubitIndex &index = signature.inWireIndices[i];
-                updatedQreg =
-                    rewriter.create<quantum::InsertOp>(loc, updatedQreg.getType(), updatedQreg,
-                                                       index.getValue(), index.getAttr(), qubit);
-            }
+            operands[operandIdx++] = getUpdatedQreg(rewriter, loc);
 
-            operands[operandIdx++] = updatedQreg;
             if (!signature.params.empty()) {
                 auto [startIdx, endIdx] =
                     findParamTypeRange(funcInputs, signature.params.size(), operandIdx);
@@ -163,16 +163,12 @@ class OpSignatureAnalyzer {
                 }
             }
 
-            if (!signature.inWireIndices.empty()) {
-                operands[operandIdx] = fromTensorOrAsIs(signature.inWireIndices,
-                                                        funcInputs[operandIdx], rewriter, loc);
-                operandIdx++;
-            }
-
-            if (!signature.inCtrlWireIndices.empty()) {
-                operands[operandIdx] = fromTensorOrAsIs(signature.inCtrlWireIndices,
-                                                        funcInputs[operandIdx], rewriter, loc);
-                operandIdx++;
+            for (const auto &indices : {signature.inWireIndices, signature.inCtrlWireIndices}) {
+                if (!indices.empty()) {
+                    operands[operandIdx] =
+                        fromTensorOrAsIs(indices, funcInputs[operandIdx], rewriter, loc);
+                    operandIdx++;
+                }
             }
         }
         else {
@@ -218,18 +214,16 @@ class OpSignatureAnalyzer {
 
         SmallVector<Value> newResults;
         rewriter.setInsertionPointAfter(callOp);
-        for (const QubitIndex &index : signature.outQubitIndices) {
-            auto extractOp = rewriter.create<quantum::ExtractOp>(
-                callOp.getLoc(), rewriter.getType<quantum::QubitType>(), qreg, index.getValue(),
-                index.getAttr());
-            newResults.emplace_back(extractOp.getResult());
-        }
-        for (const QubitIndex &index : signature.outCtrlQubitIndices) {
-            auto extractOp = rewriter.create<quantum::ExtractOp>(
-                callOp.getLoc(), rewriter.getType<quantum::QubitType>(), qreg, index.getValue(),
-                index.getAttr());
-            newResults.emplace_back(extractOp.getResult());
+
+        for (const auto &indices : {signature.outQubitIndices, signature.outCtrlQubitIndices}) {
+            for (const auto &index : indices) {
+                auto extractOp = rewriter.create<quantum::ExtractOp>(
+                    callOp.getLoc(), rewriter.getType<quantum::QubitType>(), qreg, index.getValue(),
+                    index.getAttr());
+                newResults.emplace_back(extractOp.getResult());
+            }
         }
+
         return newResults;
     }
 
@@ -245,7 +239,6 @@ class OpSignatureAnalyzer {
         ValueRange outCtrlQubits;
 
         // Qreg mode specific information
-        Value sourceQreg = nullptr;
         SmallVector<QubitIndex> inWireIndices;
         SmallVector<QubitIndex> inCtrlWireIndices;
         SmallVector<QubitIndex> outQubitIndices;
@@ -333,39 +326,21 @@ class OpSignatureAnalyzer {
         return values.front();
     }
 
-    Value getSourceQreg(Value qubit)
-    {
-        while (qubit) {
-            if (auto extractOp = qubit.getDefiningOp<quantum::ExtractOp>()) {
-                return extractOp.getQreg();
-            }
-
-            if (auto customOp = dyn_cast_or_null<quantum::CustomOp>(qubit.getDefiningOp())) {
-                if (customOp.getQubitOperands().empty()) {
-                    break;
-                }
-                qubit = customOp.getQubitOperands()[0];
-            }
-        }
-
-        return nullptr;
-    }
-
     QubitIndex getExtractIndex(Value qubit)
     {
         while (qubit) {
             if (auto extractOp = qubit.getDefiningOp<quantum::ExtractOp>()) {
                 if (Value idx = extractOp.getIdx()) {
-                    return QubitIndex(idx);
+                    return QubitIndex(idx, extractOp.getQreg());
                 }
                 if (IntegerAttr idxAttr = extractOp.getIdxAttrAttr()) {
-                    return QubitIndex(idxAttr);
+                    return QubitIndex(idxAttr, extractOp.getQreg());
                 }
             }
 
-            if (auto customOp = dyn_cast_or_null<quantum::CustomOp>(qubit.getDefiningOp())) {
-                auto qubitOperands = customOp.getQubitOperands();
-                auto qubitResults = customOp.getQubitResults();
+            if (auto gate = dyn_cast_or_null<quantum::QuantumGate>(qubit.getDefiningOp())) {
+                auto qubitOperands = gate.getQubitOperands();
+                auto qubitResults = gate.getQubitResults();
                 auto it =
                     llvm::find_if(qubitResults, [&](Value result) { return result == qubit; });
 
@@ -377,6 +352,10 @@ class OpSignatureAnalyzer {
                     }
                 }
             }
+            else if (auto measureOp = dyn_cast_or_null<quantum::MeasureOp>(qubit.getDefiningOp())) {
+                qubit = measureOp.getInQubit();
+                continue;
+            }
 
             break;
         }
@@ -394,7 +373,8 @@ struct DecomposeLoweringRewritePattern : public OpRewritePattern<CustomOp> {
     DecomposeLoweringRewritePattern(MLIRContext *context,
                                     const llvm::StringMap<func::FuncOp> &registry,
                                     const llvm::StringSet<llvm::MallocAllocator> &gateSet)
-        : OpRewritePattern(context), decompositionRegistry(registry), targetGateSet(gateSet)
+        : OpRewritePattern<CustomOp>(context), decompositionRegistry(registry),
+          targetGateSet(gateSet)
     {
     }
 
@@ -421,11 +401,12 @@ struct DecomposeLoweringRewritePattern : public OpRewritePattern<CustomOp> {
         assert(decompFunc.getFunctionType().getNumResults() >= 1 &&
                "Decomposition function must have at least one result");
 
+        rewriter.setInsertionPointAfter(op);
+
         auto enableQreg = isa<quantum::QuregType>(decompFunc.getFunctionType().getInput(0));
         auto analyzer = OpSignatureAnalyzer(op, enableQreg);
         assert(analyzer && "Analyzer should be valid");
 
-        rewriter.setInsertionPointAfter(op);
         auto callOperands = analyzer.prepareCallOperands(decompFunc, rewriter, op.getLoc());
         auto callOp =
             rewriter.create<func::CallOp>(op.getLoc(), decompFunc.getFunctionType().getResults(),
@@ -453,4 +434,4 @@ void populateDecomposeLoweringPatterns(RewritePatternSet &patterns,
 }
 
 } // namespace quantum
-} // namespace catalyst
+} // namespace catalyst