diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 65239da1b1..fafc3b94be 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -12,6 +12,11 @@

Improvements 🛠

+* 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) diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp index 894cd45d83..aa9ba60a4a 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(); + 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, 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 seen; + std::vector 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 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 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..dc80dfb58e 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp +++ b/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, 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 inputs; double totalCost{0.0}; std::unordered_map basisCounts; + + // TODO: revisit this after testing.. + RuleOrigin origin{RuleOrigin::Default}; }; /** 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_DecompGraphSolverSymbolicOps.cpp b/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolverSymbolicOps.cpp new file mode 100644 index 0000000000..26b09e5b27 --- /dev/null +++ b/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolverSymbolicOps.cpp @@ -0,0 +1,228 @@ +// 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 makeAdjoint and cancels on double application", "[DecompGraph::Core]") +{ + const OperatorNode h{"H", 1, 0, false}; + const OperatorNode adjH = makeAdjoint(h); + + REQUIRE(adjH.name == "H"); + REQUIRE(adjH.adjoint); + REQUIRE(adjH != h); + + // cancel_adjoint: Adjoint(Adjoint(H)) == H + REQUIRE(makeAdjoint(adjH) == h); + REQUIRE_FALSE(makeAdjoint(adjH).adjoint); +} + +TEST_CASE("Test makeAdjointRule", "[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 adj = makeAdjointRule(base); + + REQUIRE(adj.name == "rot_decomp_adjoint"); + REQUIRE(adj.origin == RuleOrigin::AdjointGenerated); + REQUIRE(adj.output == makeAdjoint(rot)); + REQUIRE(adj.output.adjoint); + REQUIRE(adj.inputs.size() == 2); + REQUIRE(adj.inputs[0].op == makeAdjoint(rz)); + REQUIRE(adj.inputs[0].op.adjoint); + REQUIRE(adj.inputs[0].multiplicity == 2); + REQUIRE(adj.inputs[1].op == makeAdjoint(ry)); + REQUIRE(adj.inputs[1].multiplicity == 1); +} + +TEST_CASE("Test DecompositionGraph adjoint 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{{{rz, 1.0}, {ry, 1.0}}}; + const std::vector rules{{"rot_decomp", rot, {{rz, 2}, {ry, 1}}}}; + + // Adjoint(Rot) is a root, so the builder should synthesize its adjoint decomposition + const DecompositionGraph graph({makeAdjoint(rot)}, gateset, rules); + + REQUIRE(graph.getNumRules() == 2); + REQUIRE(graph.hasOperator(makeAdjoint(rot))); + + const auto &adjRules = graph.getAllRulesFor(makeAdjoint(rot)); + REQUIRE(adjRules.size() == 1); + REQUIRE(adjRules[0].name == "rot_decomp_adjoint"); + REQUIRE(adjRules[0].origin == RuleOrigin::AdjointGenerated); + REQUIRE(adjRules[0].output == makeAdjoint(rot)); + REQUIRE(adjRules[0].inputs[0].op == makeAdjoint(rz)); + REQUIRE(adjRules[0].inputs[1].op == makeAdjoint(ry)); +} + +TEST_CASE("Test DecompositionGraph does not synthesize adjoint rules for empty or adjoint rules", + "[DecompGraph::Solver]") +{ + const OperatorNode h{"H", 1, 0, false}; + const OperatorNode adjH = makeAdjoint(h); + + const WeightedGateset gateset{{{h, 1.0}}}; + const std::vector rules{ + {"h_is_basis", h, {}}, // empty rule + {"self_adjoint_H", adjH, {{h, 1}}}, // adjoint output, must not be mirrored!! + }; + + const DecompositionGraph graph({h}, gateset, rules); + + REQUIRE(graph.getNumRules() == 2); + REQUIRE(graph.getAllRulesFor(adjH).size() == 1); + REQUIRE(graph.getAllRulesFor(adjH)[0].name == "self_adjoint_H"); +} + +TEST_CASE("Test Adjoint: self_adjoint (Adjoint(H) -> H)", "[DecompGraph::Solver]") +{ + const OperatorNode h{"H", 1, 0, false}; + const OperatorNode adjH = makeAdjoint(h); + + const WeightedGateset gateset{{{h, 1.0}}}; + const std::vector rules{{"self_adjoint_H", adjH, {{h, 1}}}}; + + const DecompositionGraph graph({adjH}, gateset, rules); + DecompositionSolver solver(graph); + const auto result = solver.solve(); + + REQUIRE(result.find(adjH) != result.end()); + const auto &chosen = result.at(adjH); + REQUIRE_FALSE(chosen.isBasis); + REQUIRE(chosen.ruleName == "self_adjoint_H"); + REQUIRE(chosen.origin == RuleOrigin::Default); + REQUIRE(chosen.totalCost == 1.0); + REQUIRE(chosen.basisCounts.at(h) == 1); + + REQUIRE(graph.getAllRulesFor(adjH).size() == 1); +} + +TEST_CASE("Test Adjoint: adjoint_rotation (Adjoint(RX) -> RX)", "[DecompGraph::Solver]") +{ + const OperatorNode rx{"RX", 1, 1, false}; + const OperatorNode adjRX = makeAdjoint(rx); + + const WeightedGateset gateset{{{rx, 1.0}}}; + const std::vector rules{{"adjoint_rotation_RX", adjRX, {{rx, 1}}}}; + + const DecompositionGraph graph({adjRX}, gateset, rules); + DecompositionSolver solver(graph); + const auto result = solver.solve(); + + const auto &chosen = result.at(adjRX); + REQUIRE(chosen.ruleName == "adjoint_rotation_RX"); + REQUIRE(chosen.totalCost == 1.0); + REQUIRE(chosen.basisCounts.at(rx) == 1); +} + +TEST_CASE("Test Adjoint: multiple rules and the solver should pick the cheapest", + "[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 commonRules{ + {"rot_decomp", rot, {{rz, 2}, {ry, 1}}}, + {"adjoint_rotation_RZ", makeAdjoint(rz), {{rz, 1}}}, + {"adjoint_rotation_RY", makeAdjoint(ry), {{ry, 1}}}, + }; + + SECTION("rot_decomp_adjoint is cheaper") + { + const WeightedGateset gateset{{{rz, 1.0}, {ry, 1.0}, {e, 10.0}}}; + std::vector rules = commonRules; + rules.push_back({"_adjoint_rot", makeAdjoint(rot), {{e, 1}}}); // cost 10 + + const DecompositionGraph graph({makeAdjoint(rot)}, gateset, rules); + + // Both an explicit adjoint rule and the synthesized one exist for Adjoint(Rot). + REQUIRE(graph.getAllRulesFor(makeAdjoint(rot)).size() == 2); + + DecompositionSolver solver(graph); + const auto result = solver.solve(); + const auto &chosen = result.at(makeAdjoint(rot)); + REQUIRE(chosen.ruleName == "rot_decomp_adjoint"); + REQUIRE(chosen.origin == RuleOrigin::AdjointGenerated); + REQUIRE(chosen.totalCost == 3.0); + REQUIRE(chosen.basisCounts.at(rz) == 2); + REQUIRE(chosen.basisCounts.at(ry) == 1); + } + + SECTION("_adjoint_rot is cheaper") + { + const WeightedGateset gateset{{{rz, 1.0}, {ry, 1.0}}}; + std::vector rules = commonRules; + rules.push_back({"_adjoint_rot", makeAdjoint(rot), {{rz, 1}}}); // cost 1 + + const DecompositionGraph graph({makeAdjoint(rot)}, gateset, rules); + DecompositionSolver solver(graph); + const auto result = solver.solve(); + const auto &chosen = result.at(makeAdjoint(rot)); + REQUIRE(chosen.ruleName == "_adjoint_rot"); + REQUIRE(chosen.origin == RuleOrigin::Default); + REQUIRE(chosen.totalCost == 1.0); + } +} + +TEST_CASE("Test Adjoint: adjoint 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{{{a, 1.0}, {b, 1.0}}}; + const std::vector rules{ + {"myop_decomp", myOp, {{a, 1}, {b, 1}}}, + // Define self_adjoint rules so the adjointed produced gates can resolve: + {"self_adjoint_A", makeAdjoint(a), {{a, 1}}}, + {"self_adjoint_B", makeAdjoint(b), {{b, 1}}}, + }; + + const DecompositionGraph graph({makeAdjoint(myOp)}, gateset, rules); + + REQUIRE(graph.getAllRulesFor(makeAdjoint(myOp)).size() == 1); + + DecompositionSolver solver(graph); + const auto result = solver.solve(); + const auto &chosen = result.at(makeAdjoint(myOp)); + REQUIRE(chosen.ruleName == "myop_decomp_adjoint"); + REQUIRE(chosen.origin == RuleOrigin::AdjointGenerated); + REQUIRE(chosen.totalCost == 2.0); + REQUIRE(chosen.basisCounts.at(a) == 1); + REQUIRE(chosen.basisCounts.at(b) == 1); +}