Support decomposition to/from Adjoint Ops to the GraphSolver

doc/releases/changelog-dev.md

@@ -12,6 +12,11 @@
 
 <h3>Improvements 🛠</h3>
 
+* Add adjoint support to the decomposition graph solver, enabling `Adjoint(Op)` to be decomposed either
+  via registered adjoint rules or by adjointing the base operator's decomposition rule,
+  with the solver choosing the cheapest.
+  [(#3001)](https://github.com/PennyLaneAI/catalyst/pull/3001)
+
 * Adds a `catalyst::symbolic_array` operation and integrates it with the new `qp.capture.symbolic_array` function.
   [(#2982)](https://github.com/PennyLaneAI/catalyst/pull/2982)
 

mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp

@@ -100,6 +100,7 @@ struct DecompositionGraph::Impl {
           fixedDecomps(std::move(_fixedDecomps)), altDecomps(std::move(_altDecomps))
     {
         materializeRules();
+        generateAdjointRules();
     }
 
     void materializeRules()
@@ -155,6 +156,79 @@ struct DecompositionGraph::Impl {
         rules = std::move(effectiveRules);
     }
 
+    /**
+     * @brief Generate Adjoint decomposition rules
+     *
+     * For a needed adjoint operator `Adjoint(Op)` this synthesizes, from every base decomposition
+     * `Op`, a rule `Adjoint(Op)` so the adjoint operator can be decomposed by adjointing the
+     * decomposition of its base. These coexist with any explicitly registered adjoint rules;
+     * the solver then compares their costs and picks the cheapest.
+     *
+     * @note Rules are synthesized only on-demand: only Adjoint operators that actually appear
+     * in the circuit (as roots) seed the process, and adjointing a decomposition may introduce
+     * new Adjoint(input) operators that require their own synthesized rules, so the process
+     * runs to a fixpoint. We leave the graph untouched when there aren't any adjoint ops.
+     *
+     * Two families of rules are never used as a base in this method:
+     * - rules whose output is already adjoint: we only derive adjoint rules from base
+     *   decompositions.
+     * - empty rules in the graph: these mark basis/target gates, and the adjoint of
+     *   a basis gate is not necessarily available for free, so it must be provided
+     *   explicitly (e.g. a self_adjoint rule) rather than synthesized.
+     *   TODO: we'll revisit this when integrating this with graph-decomposition.
+     */
+    void generateAdjointRules()
+    {
+        // Index valid base decompositions by their (non-adjoint) output.
+        std::unordered_map<OperatorNode, std::vector<RuleNode>, OperatorNodeHash> baseByOutput;
+        for (const auto &rule : rules) {
+            if (!rule.output.adjoint && !rule.isEmpty()) {
+                baseByOutput[rule.output].push_back(rule);
+            }
+        }
+        if (baseByOutput.empty()) {
+            return;
+        }
+
+        // for every Adjoint(Op) in the circuit:
+        std::unordered_set<OperatorNode, OperatorNodeHash> seen;
+        std::vector<OperatorNode> worklist;
+        auto enqueue = [&](const OperatorNode &op) {
+            if (op.adjoint && seen.insert(op).second) {
+                worklist.push_back(op);
+            }
+        };
+        for (const auto &op : operators) {
+            enqueue(op);
+        }
+        for (const auto &rule : rules) {
+            enqueue(rule.output);
+            for (const auto &term : rule.inputs) {
+                enqueue(term.op);
+            }
+        }
+
+        std::vector<RuleNode> generated;
+        while (!worklist.empty()) {
+            const OperatorNode adjOp = worklist.back();
+            worklist.pop_back();
+            const auto it = baseByOutput.find(makeAdjoint(adjOp));
+            if (it == baseByOutput.end()) {
+                continue;
+            }
+            for (const auto &baseRule : it->second) {
+                RuleNode adjRule = makeAdjointRule(baseRule);
+                for (const auto &term : adjRule.inputs) {
+                    enqueue(term.op);
+                }
+                generated.push_back(std::move(adjRule));
+            }
+        }
+        for (auto &rule : generated) {
+            rules.push_back(std::move(rule));
+        }
+    }
+
     void buildGraph()
     {
         // Register all operators

mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.cpp

@@ -51,6 +51,7 @@ ChosenDecompRule DecompositionSolver::evalRule(const RuleNode &rule)
     solution.isBasis = false;
     solution.inputs = rule.inputs;
     solution.op = rule.output;
+    solution.origin = rule.origin;
 
     double total_cost = 0.0;
     for (const auto &input : rule.inputs) {

mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp

@@ -140,8 +140,9 @@ struct RuleTerm {
  * - Default: The default rule for decomposing an operator as defined in the decomposition graph.
  * - Fixed: A fixed rule that cannot be changed or overridden by the solver.
  * - Alternative: An alternative rule that can be used in place of the default rule.
+ * - AdjointGenerated: A rule synthesized by adjointing a base decomposition rule.
  */
-enum class RuleOrigin : uint8_t { Default = 0, Fixed = 1, Alternative = 2 };
+enum class RuleOrigin : uint8_t { Default = 0, Fixed = 1, Alternative = 2, AdjointGenerated = 3 };
 
 /**
  * @brief This represents the decomposition rules in the graph decomposition problem.
@@ -185,6 +186,38 @@ using FixedDecomps = std::unordered_map<OperatorNode, RuleNode, OperatorNodeHash
  */
 using AltDecomps = std::unordered_map<OperatorNode, std::vector<RuleNode>, OperatorNodeHash>;
 
+/**
+ * @brief This returns a copy of the given operator with its adjoint flag.
+ *
+ * Note taht because Adjoint is represented as a boolean flag in the IR,
+ * applying it twice should cancel out (Adjoint(Adjoint(op)) == op).
+ */
+inline OperatorNode makeAdjoint(OperatorNode op)
+{
+    op.adjoint = !op.adjoint;
+    return op;
+}
+
+/**
+ * @brief Constructs the Adjoint decomposition of a base rule.
+ *
+ * Given a rule `output -> {inputs}`, produces `Adjoint(output) -> {Adjoint(input), ...}`
+ * with the same multiplicities: the adjoint of a decomposition is obtained by adjointing
+ * every produced gate (and reversing their order, which does not affect resource/cost counting).
+ */
+inline RuleNode makeAdjointRule(const RuleNode &base)
+{
+    RuleNode adj;
+    adj.name = base.name + "_adjoint";
+    adj.output = makeAdjoint(base.output);
+    adj.origin = RuleOrigin::AdjointGenerated;
+    adj.inputs.reserve(base.inputs.size());
+    for (const auto &term : base.inputs) {
+        adj.inputs.push_back({makeAdjoint(term.op), term.multiplicity});
+    }
+    return adj;
+}
+
 /**
  * @brief This represents the chosen decomposition rule for an operator in
  * the solution of the graph decomposition problem.
@@ -196,6 +229,9 @@ struct ChosenDecompRule {
     std::vector<RuleTerm> inputs;
     double totalCost{0.0};
     std::unordered_map<OperatorNode, std::size_t, OperatorNodeHash> basisCounts;
+
+    // TODO: revisit this after testing..
+    RuleOrigin origin{RuleOrigin::Default};
 };
 
 /**

mlir/unittests/DecompGraphSolver/CMakeLists.txt

@@ -25,6 +25,7 @@ include(Catch)
 add_executable(runner_tests_dgsolver
     Test_DecompGraphCore.cpp
     Test_DecompGraphSolver.cpp
+    Test_DecompGraphSolverSymbolicOps.cpp
 )
 
 target_link_libraries(runner_tests_dgsolver PRIVATE