Support decomposition of Controlled ops to GraphSolver

doc/releases/changelog-dev.md

@@ -12,6 +12,11 @@
 
 <h3>Improvements 🛠</h3>
 
+* Add controlled support to the decomposition graph solver, enabling `C(Op)`
+  to be decomposed either via registered controlled rules or by controlling
+  the base operator's decomposition rule, with the solver choosing the cheapest.
+  [(#3003)](https://github.com/PennyLaneAI/catalyst/pull/3003)
+
 * 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();
+        generateControlledRules();
     }
 
     void materializeRules()
@@ -155,6 +156,98 @@ struct DecompositionGraph::Impl {
         rules = std::move(effectiveRules);
     }
 
+    /**
+     * @brief Operators whose controlled form must be produced by a dedicated rule
+     * rather than by controlling their decomposition, so control-each-gate is
+     * suppressed for them.
+     */
+    static bool isCtrlRuleRequired(const std::string &opName) { return opName == "GlobalPhase"; }
+
+    /**
+     * @brief Index non-empty, uncontrolled decompositions by their output operator.
+     */
+    std::unordered_map<OperatorNode, std::vector<RuleNode>, OperatorNodeHash>
+    indexUncontrolledBaseRules() const
+    {
+        std::unordered_map<OperatorNode, std::vector<RuleNode>, OperatorNodeHash> baseByOutput;
+        for (const auto &rule : rules) {
+            if (rule.output.numControlWires == 0 && !rule.isEmpty()) {
+                baseByOutput[rule.output].push_back(rule);
+            }
+        }
+        return baseByOutput;
+    }
+
+    /**
+     * @brief Generate Controlled decomposition rules.
+     *
+     * For a needed controlled operator `Controlled(Op)` with k control wires this synthesizes,
+     * from every base decomposition `Op`, a rule `Controlled(Op)`with the same k control wires,
+     * so a controlled operator can be decomposed by controlling the decomposition of its base.
+     * These coexist with any explicitly registered controlled rules. The solver then compares
+     * their costs and picks the cheapest rules.
+     *
+     * Rules are synthesized lazily: only controlled operators that actually appear in the problem
+     * seed the process, and controlling a decomposition introduces new `Controlled(Op)` operators
+     * that require their own synthesized rules, so it runs to a fixpoint.
+     *
+     * Only non-empty, uncontrolled base rules are used as a base, and control-each-gate is
+     * suppressed for operators in `isCtrlRuleRequired`. Those operators require dedicated
+     * decomposition rules.
+     */
+    void generateControlledRules()
+    {
+        const auto baseByOutput = indexUncontrolledBaseRules();
+        if (baseByOutput.empty()) {
+            return;
+        }
+
+        // Breadth-first over controlled operators
+        // the controlled inputs each synthesized rule introduces, until no new ones appear:
+        std::unordered_set<OperatorNode, OperatorNodeHash> seen;
+        std::vector<OperatorNode> worklist;
+        auto enqueue = [&](const OperatorNode &op) {
+            if (op.numControlWires > 0 && 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 ctrlOp = worklist.back();
+            worklist.pop_back();
+
+            if (isCtrlRuleRequired(ctrlOp.name)) {
+                continue;
+            }
+
+            const auto it = baseByOutput.find(withoutControls(ctrlOp));
+            if (it == baseByOutput.end()) {
+                continue;
+            }
+            for (const auto &baseRule : it->second) {
+                RuleNode ctrlRule = makeControlledRule(baseRule, ctrlOp.numControlWires);
+                for (const auto &term : ctrlRule.inputs) {
+                    enqueue(term.op);
+                }
+                generated.push_back(std::move(ctrlRule));
+            }
+        }
+
+        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

@@ -52,9 +52,10 @@ namespace DecompGraph::Core {
  */
 struct OperatorNode {
     std::string name;
-    int numWires{-1};
-    int numParams{-1};
+    int64_t numWires{-1};  // -1 = any number of wires
+    int64_t numParams{-1}; // -1 = any number of params
     bool adjoint{false};
+    std::size_t numControlWires{0};
 
     // Optional static arguments for operators that require additional data.
     std::unordered_map<std::string, std::string> staticNamedArgs{};
@@ -75,7 +76,7 @@ struct OperatorNode {
                                        staticNamedArgs == other.staticNamedArgs;
 
         return name == other.name && default_wires && default_params && adjoint == other.adjoint &&
-               static_args_match;
+               static_args_match && numControlWires == other.numControlWires;
     }
     bool operator!=(const OperatorNode &other) const { return !(*this == other); }
 };
@@ -140,8 +141,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.
+ * - ControlGenerated: A rule synthesized by controlling a base decomposition rule.
  */
-enum class RuleOrigin : uint8_t { Default = 0, Fixed = 1, Alternative = 2 };
+enum class RuleOrigin : uint8_t { Default = 0, Fixed, Alternative, ControlGenerated };
 
 /**
  * @brief This represents the decomposition rules in the graph decomposition problem.
@@ -152,7 +154,7 @@ enum class RuleOrigin : uint8_t { Default = 0, Fixed = 1, Alternative = 2 };
  * decomposition rules to break down complex operators into simpler ones that are part of
  * the target gateset.
  *
- * TODO:
+ * @todo
  * - We can add a field for work_wires_required if we want to consider the number of ancillary
  * wires needed for the decomposition, which can be an important factor in resource optimization.
  * - We can also consider adding a field for the decomposition function or a pointer to it,
@@ -185,6 +187,54 @@ 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 wrapped in `numControlWires`
+ * additional controls (Controlled(op)). Control wires accumulate, so applying it
+ * repeatedly yields a multi-controlled operator.
+ */
+inline OperatorNode makeControlled(OperatorNode op, std::size_t numControlWires = 1)
+{
+    op.numControlWires += numControlWires;
+    return op;
+}
+
+/**
+ * @brief This returns a copy of the given operator with all controls removed.
+ */
+inline OperatorNode withoutControls(OperatorNode op)
+{
+    op.numControlWires = 0;
+    return op;
+}
+
+/**
+ * @brief Constructs the Controlled decomposition rule from a base rule.
+ *
+ * Given a rule `output -> {inputs}`, produces `Controlled(output) -> {Controlled(input), ...}`
+ * where every operator gains `numControlWires` control wires: controlling a decomposition means
+ * applying the same controls to each gate it produces.
+ *
+ * The `numControlWires` count is encoded in the rule name so distinct control counts over
+ * the same base rule stay unique. The result is tagged with `RuleOrigin::ControlGenerated`
+ * so later stages can lower it by controlling each gate.
+ *
+ * @note: PennyLane counts `PauliX` flips for zero `control_values`.
+ * Those flips and `control_values` are not supported yet; the cost reflects only the
+ * cost of controlling each produced gate.
+ */
+inline RuleNode makeControlledRule(const RuleNode &base, std::size_t numControlWires)
+{
+    RuleNode ctrl;
+    ctrl.name = base.name + "_controlled_" + std::to_string(numControlWires);
+    ctrl.output = makeControlled(base.output, numControlWires);
+    ctrl.origin = RuleOrigin::ControlGenerated;
+    ctrl.inputs.reserve(base.inputs.size());
+    for (const auto &term : base.inputs) {
+        ctrl.inputs.push_back({makeControlled(term.op, numControlWires), term.multiplicity});
+    }
+    return ctrl;
+}
+
 /**
  * @brief This represents the chosen decomposition rule for an operator in
  * the solution of the graph decomposition problem.
@@ -196,6 +246,7 @@ struct ChosenDecompRule {
     std::vector<RuleTerm> inputs;
     double totalCost{0.0};
     std::unordered_map<OperatorNode, std::size_t, OperatorNodeHash> basisCounts;
+    RuleOrigin origin{RuleOrigin::Default};
 };
 
 /**

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

@@ -40,6 +40,9 @@ static inline auto print_op(const OperatorNode &op) -> std::string
     if (op.adjoint) {
         oss << "[adj]";
     }
+    if (op.numControlWires > 0) {
+        oss << "[c:" << op.numControlWires << "]";
+    }
     if (!op.staticNamedArgs.empty()) {
         std::vector<std::string> keys;
         keys.reserve(op.staticNamedArgs.size());

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

mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp

@@ -562,8 +562,8 @@ TEST_CASE("Test GraphSolver with PauliRot specialized by static argument pauli_w
 {
     // Query: PauliRot[w:1][p:1][pauli_word:X] should match a rule whose output is
     // PauliRot[w:-1][p:-1][pauli_word:X] (wildcards on wires/params, exact match on pauli_word).
-    const OperatorNode pauliRotQuery{"PauliRot", 1, 1, false, {{"pauli_word", "X"}}};
-    const OperatorNode pauliRotRuleOutput{"PauliRot", -1, -1, false, {{"pauli_word", "X"}}};
+    const OperatorNode pauliRotQuery{"PauliRot", 1, 1, false, 0, {{"pauli_word", "X"}}};
+    const OperatorNode pauliRotRuleOutput{"PauliRot", -1, -1, false, 0, {{"pauli_word", "X"}}};
     const OperatorNode hadamard{"Hadamard", 1, 0, false};
     const OperatorNode multiRZ{"MultiRZ", 1, 1, false};
 
@@ -585,16 +585,16 @@ TEST_CASE("Test GraphSolver with PauliRot specialized by static argument pauli_w
     REQUIRE(chosen.basisCounts.at(hadamard) == 2);
     REQUIRE(chosen.basisCounts.at(multiRZ) == 1);
 
-    const OperatorNode pauliRotQueryY{"PauliRot", 1, 1, false, {{"pauli_word", "Y"}}};
+    const OperatorNode pauliRotQueryY{"PauliRot", 1, 1, false, 0, {{"pauli_word", "Y"}}};
     REQUIRE_FALSE(pauliRotQuery == pauliRotQueryY);
     REQUIRE(pauliRotQuery == pauliRotRuleOutput);
 }
 
 TEST_CASE("Test OperatorNode equality with staticNamedArgs", "[DecompGraph::Core]")
 {
-    const OperatorNode pauliRotX{"PauliRot", 1, 1, false, {{"pauli_word", "X"}}};
-    const OperatorNode pauliRotXWildcard{"PauliRot", -1, -1, false, {{"pauli_word", "X"}}};
-    const OperatorNode pauliRotY{"PauliRot", 1, 1, false, {{"pauli_word", "Y"}}};
+    const OperatorNode pauliRotX{"PauliRot", 1, 1, false, 0, {{"pauli_word", "X"}}};
+    const OperatorNode pauliRotXWildcard{"PauliRot", -1, -1, false, 0, {{"pauli_word", "X"}}};
+    const OperatorNode pauliRotY{"PauliRot", 1, 1, false, 0, {{"pauli_word", "Y"}}};
     const OperatorNode pauliRotNoArgs{"PauliRot", 1, 1, false};
 
     REQUIRE(pauliRotX == pauliRotXWildcard);