diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 65239da1b1..332a946ec6 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -12,6 +12,11 @@
Improvements ðŸ›
+* 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)
diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp
index 894cd45d83..583e7cd4b9 100644
--- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp
+++ b/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, OperatorNodeHash>
+ indexUncontrolledBaseRules() const
+ {
+ std::unordered_map, 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 seen;
+ std::vector 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 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
diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.cpp
index c83ccddce5..d05fcb321a 100644
--- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.cpp
+++ b/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) {
diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp
index 9734285415..0d4da31368 100644
--- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp
+++ b/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 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, 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 inputs;
double totalCost{0.0};
std::unordered_map basisCounts;
+ RuleOrigin origin{RuleOrigin::Default};
};
/**
diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGUtils.hpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGUtils.hpp
index ebab329883..f1255d4c45 100644
--- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGUtils.hpp
+++ b/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 keys;
keys.reserve(op.staticNamedArgs.size());
diff --git a/mlir/unittests/DecompGraphSolver/CMakeLists.txt b/mlir/unittests/DecompGraphSolver/CMakeLists.txt
index afa4b9e197..43969b26d6 100644
--- a/mlir/unittests/DecompGraphSolver/CMakeLists.txt
+++ b/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
diff --git a/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp b/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp
index 46a2602791..e5fb579216 100644
--- a/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp
+++ b/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);
diff --git a/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolverSymbolicOps.cpp b/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolverSymbolicOps.cpp
new file mode 100644
index 0000000000..116d596bfc
--- /dev/null
+++ b/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolverSymbolicOps.cpp
@@ -0,0 +1,180 @@
+// Copyright 2026 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.
+
+#include
+#include
+
+#include "DGBuilder.hpp"
+#include "DGSolver.hpp"
+#include "DGTypes.hpp"
+#include "DGUtils.hpp"
+
+#include
+#include
+#include
+#include
+
+using namespace Catch::Matchers;
+using namespace DecompGraph::Core;
+using namespace DecompGraph::Solver;
+
+TEST_CASE("Test makeControlled", "[DecompGraph::Core]")
+{
+ const OperatorNode rx{"RX", 1, 1, false};
+ const OperatorNode crx = makeControlled(rx);
+
+ REQUIRE(crx.name == "RX");
+ REQUIRE(crx.numControlWires == 1);
+ REQUIRE(crx != rx);
+
+ // Controls accumulate into a multi-controlled operator.
+ REQUIRE(makeControlled(crx).numControlWires == 2);
+ REQUIRE(makeControlled(rx, 3).numControlWires == 3);
+}
+
+TEST_CASE("Test makeControlledRule controls the output and every input", "[DecompGraph::Core]")
+{
+ const OperatorNode rot{"Rot", 1, 3, false};
+ const OperatorNode rz{"RZ", 1, 1, false};
+ const OperatorNode ry{"RY", 1, 1, false};
+ const RuleNode base{"rot_decomp", rot, {{rz, 2}, {ry, 1}}};
+
+ const RuleNode ctrl = makeControlledRule(base, 1);
+
+ REQUIRE(ctrl.name == "rot_decomp_controlled_1");
+ REQUIRE(ctrl.origin == RuleOrigin::ControlGenerated);
+ REQUIRE(ctrl.output == makeControlled(rot));
+ REQUIRE(ctrl.output.numControlWires == 1);
+ REQUIRE(ctrl.inputs.size() == 2);
+ REQUIRE(ctrl.inputs[0].op == makeControlled(rz));
+ REQUIRE(ctrl.inputs[0].op.numControlWires == 1);
+ REQUIRE(ctrl.inputs[0].multiplicity == 2);
+ REQUIRE(ctrl.inputs[1].op == makeControlled(ry));
+ REQUIRE(ctrl.inputs[1].multiplicity == 1);
+}
+
+TEST_CASE("Test DecompositionGraph synthesizes controlled rules from base rules",
+ "[DecompGraph::Solver]")
+{
+ const OperatorNode rot{"Rot", 1, 3, false};
+ const OperatorNode rz{"RZ", 1, 1, false};
+ const OperatorNode ry{"RY", 1, 1, false};
+
+ const WeightedGateset gateset{{{makeControlled(rz), 1.0}, {makeControlled(ry), 1.0}}};
+ const std::vector rules{{"rot_decomp", rot, {{rz, 2}, {ry, 1}}}};
+
+ // Controlled(Rot) is a root, so the builder synthesizes its controlled decomposition.
+ const DecompositionGraph graph({makeControlled(rot)}, gateset, rules);
+
+ REQUIRE(graph.getNumRules() == 2);
+ REQUIRE(graph.hasOperator(makeControlled(rot)));
+
+ const auto &ctrlRules = graph.getAllRulesFor(makeControlled(rot));
+ REQUIRE(ctrlRules.size() == 1);
+ REQUIRE(ctrlRules[0].name == "rot_decomp_controlled_1");
+ REQUIRE(ctrlRules[0].origin == RuleOrigin::ControlGenerated);
+ REQUIRE(ctrlRules[0].output == makeControlled(rot));
+ REQUIRE(ctrlRules[0].inputs[0].op == makeControlled(rz));
+ REQUIRE(ctrlRules[0].inputs[1].op == makeControlled(ry));
+}
+
+TEST_CASE("Test Controlled: solver picks the cheaper", "[DecompGraph::Solver]")
+{
+ const OperatorNode rot{"Rot", 1, 3, false};
+ const OperatorNode rz{"RZ", 1, 1, false};
+ const OperatorNode ry{"RY", 1, 1, false};
+ const OperatorNode e{"E", 1, 0, false};
+
+ const std::vector baseRules{{"rot_decomp", rot, {{rz, 2}, {ry, 1}}}};
+
+ SECTION("rot_decomp_controlled_1 is cheaper")
+ {
+ const WeightedGateset gateset{
+ {{makeControlled(rz), 1.0}, {makeControlled(ry), 1.0}, {e, 10.0}}};
+ std::vector rules = baseRules;
+ rules.push_back({"crot_direct", makeControlled(rot), {{e, 1}}}); // cost 10
+
+ const DecompositionGraph graph({makeControlled(rot)}, gateset, rules);
+
+ // Both an explicit controlled rule and the synthesized one exist for Controlled(Rot).
+ REQUIRE(graph.getAllRulesFor(makeControlled(rot)).size() == 2);
+
+ DecompositionSolver solver(graph);
+ const auto result = solver.solve();
+ const auto &chosen = result.at(makeControlled(rot));
+ REQUIRE(chosen.ruleName == "rot_decomp_controlled_1");
+ REQUIRE(chosen.origin == RuleOrigin::ControlGenerated);
+ REQUIRE(chosen.totalCost == 3.0);
+ REQUIRE(chosen.basisCounts.at(makeControlled(rz)) == 2);
+ REQUIRE(chosen.basisCounts.at(makeControlled(ry)) == 1);
+ }
+
+ SECTION("crot_direct is cheaper")
+ {
+ const WeightedGateset gateset{{{makeControlled(rz), 1.0}, {makeControlled(ry), 1.0}}};
+ std::vector rules = baseRules;
+ rules.push_back({"crot_direct", makeControlled(rot), {{makeControlled(rz), 1}}}); // cost 1
+
+ const DecompositionGraph graph({makeControlled(rot)}, gateset, rules);
+ DecompositionSolver solver(graph);
+ const auto result = solver.solve();
+ const auto &chosen = result.at(makeControlled(rot));
+ REQUIRE(chosen.ruleName == "crot_direct");
+ REQUIRE(chosen.origin == RuleOrigin::Default);
+ REQUIRE(chosen.totalCost == 1.0);
+ }
+}
+
+TEST_CASE("Test Controlled: control pushed through a decomposition", "[DecompGraph::Solver]")
+{
+ const OperatorNode myOp{"MyOp", 2, 0, false};
+ const OperatorNode a{"A", 1, 0, false};
+ const OperatorNode b{"B", 1, 0, false};
+
+ const WeightedGateset gateset{{{makeControlled(a), 1.0}, {makeControlled(b), 1.0}}};
+ const std::vector rules{{"myop_decomp", myOp, {{a, 1}, {b, 1}}}};
+
+ const DecompositionGraph graph({makeControlled(myOp)}, gateset, rules);
+
+ REQUIRE(graph.getAllRulesFor(makeControlled(myOp)).size() == 1);
+
+ DecompositionSolver solver(graph);
+ const auto result = solver.solve();
+ const auto &chosen = result.at(makeControlled(myOp));
+ REQUIRE(chosen.ruleName == "myop_decomp_controlled_1");
+ REQUIRE(chosen.origin == RuleOrigin::ControlGenerated);
+ REQUIRE(chosen.totalCost == 2.0);
+ REQUIRE(chosen.basisCounts.at(makeControlled(a)) == 1);
+ REQUIRE(chosen.basisCounts.at(makeControlled(b)) == 1);
+}
+
+TEST_CASE("Test Controlled: suppressed Ctrl rule for special-cased operators",
+ "[DecompGraph::Solver]")
+{
+ const OperatorNode globalPhase{"GlobalPhase", -1, -1, false};
+ const OperatorNode rz{"RZ", 1, 1, false};
+
+ const WeightedGateset gateset{{{makeControlled(rz), 1.0}}};
+ const std::vector rules{
+ {"gp_decomp", globalPhase, {{rz, 1}}},
+ {"cgp_direct", makeControlled(globalPhase), {{makeControlled(rz), 1}}},
+ };
+
+ const DecompositionGraph graph({makeControlled(globalPhase)}, gateset, rules);
+
+ const auto &ctrlRules = graph.getAllRulesFor(makeControlled(globalPhase));
+ REQUIRE(ctrlRules.size() == 1);
+ REQUIRE(ctrlRules[0].name == "cgp_direct");
+ REQUIRE(ctrlRules[0].origin == RuleOrigin::Default);
+}