Skip to content

Commit 1e784d7

Browse files
authored
Add support for empty rules in graph-decomposition (#2765)
**Context:** Add support for empty rules in graph-decomposition. **Description of the Change:** **Benefits:** ``` python import pennylane as qp from catalyst import qjit from catalyst.jax_primitives import decomposition_rule from catalyst.passes import graph_decomposition @decomposition_rule(op_type=qp.PauliX) def x_to_rx(_wire: int): pass @qjit(capture=True) @graph_decomposition( gate_set={qp.RX, qp.RY}, fixed_decomps={qp.PauliX: x_to_rx}, ) @qp.qnode(qp.device("lightning.qubit", wires=2)) def circuit(): qp.PauliX(0) qp.PauliX(1) qp.RX(0.5, wires=0) # register custom decomposition rules x_to_rx(int) return qp.state() qp.specs(circuit, level="device")().resources.gate_types ``` ``` console {'RX': 1} ``` **Related GitHub Issues:**
1 parent 071c6fb commit 1e784d7

6 files changed

Lines changed: 68 additions & 3 deletions

File tree

doc/releases/changelog-dev.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,7 @@
305305
[(#2568)](https://github.com/PennyLaneAI/catalyst/pull/2568)
306306
[(#2578)](https://github.com/PennyLaneAI/catalyst/pull/2578)
307307
[(#2711)](https://github.com/PennyLaneAI/catalyst/pull/2711)
308+
[(#2765)](https://github.com/PennyLaneAI/catalyst/pull/2765)
308309

309310
The framework is interfaced with a new `graph_decomposition` pass decorator
310311
with key capabilities:

frontend/test/pytest/from_plxpr/test_decompose_transform.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
This module tests the decompose transformation.
1616
"""
1717

18+
# pylint: disable=too-many-lines
19+
1820
from contextlib import nullcontext as does_not_raise
1921
from functools import partial
2022

@@ -318,6 +320,32 @@ def circuit():
318320
resources = qp.specs(circuit, level="device")()["resources"].gate_types
319321
assert resources == expected_resources
320322

323+
def test_empty_rule(self):
324+
"""Test that a decomposition rule with no ops is handled correctly."""
325+
326+
@decomposition_rule(op_type="PauliX")
327+
def empty_decomp(_wire):
328+
pass
329+
330+
@qp.qjit(capture=True)
331+
@graph_decomposition(
332+
gate_set={"PauliY"},
333+
fixed_decomps={"PauliX": empty_decomp},
334+
)
335+
@qp.qnode(qp.device("lightning.qubit", wires=1))
336+
def circuit():
337+
qp.X(0)
338+
qp.Y(0)
339+
340+
# register the empty decomposition rule
341+
empty_decomp(int)
342+
343+
return qp.expval(qp.Z(0))
344+
345+
expected_resources = {"PauliY": 1}
346+
resources = qp.specs(circuit, level="device")()["resources"].gate_types
347+
assert resources == expected_resources
348+
321349
@pytest.mark.xfail(
322350
reason="graph-decomposition supports pre-compiled rules, alt_decomps and fix_decomps"
323351
)

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,12 @@ struct DecompositionGraph::Impl {
182182
// Connect rule vertex to output operator vertex
183183
boost::add_edge(rule_vertex, output_vertex, GraphWeightedEdge{}, graph);
184184

185+
// Empty rules (with no inputs) are effectively just target gates
186+
// and don't need to be connected to input operator vertices
187+
if (rule.isEmpty()) {
188+
continue;
189+
}
190+
185191
// Connect rule vertex to input operator vertices
186192
for (const auto &input : rule.inputs) {
187193
const auto input_id = registerOp(input.op);

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

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,9 +67,6 @@ ChosenDecompRule DecompositionSolver::evalRule(const RuleNode &rule)
6767
}
6868
}
6969

70-
if (rule.inputs.empty() && total_cost == 0.0) {
71-
return invalidRule(solution.op); // invalid rule
72-
}
7370
solution.totalCost = total_cost;
7471
return solution;
7572
}
@@ -81,6 +78,15 @@ ChosenDecompRule DecompositionSolver::bestRule(const OperatorNode &op)
8178
return invalidRule(op); // no valid rules
8279
}
8380

81+
// if there is an empty rule (with no inputs)
82+
// for the given operator, pick this as the
83+
// best rule with zero cost
84+
for (const auto &rule : all_rules) {
85+
if (rule.isEmpty()) {
86+
return evalRule(rule);
87+
}
88+
}
89+
8490
std::optional<ChosenDecompRule> best_rule;
8591

8692
for (const auto &rule : all_rules) {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,8 @@ struct RuleNode {
153153
{
154154
return name == other.name && output == other.output && origin == other.origin;
155155
}
156+
157+
bool isEmpty() const { return inputs.empty(); }
156158
};
157159

158160
/**

mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -533,3 +533,25 @@ TEST_CASE("Test GraphSolver with MultiRZ decompositions", "[DecompGraph::Solver]
533533
REQUIRE(chosen_rule_multiRZ5.ruleName == "multiRZ5_to_rz");
534534
REQUIRE(chosen_rule_multiRZ5.totalCost == 1.0 * 5);
535535
}
536+
537+
TEST_CASE("Test GraphSolver with empty decomposition rules", "[DecompGraph::Solver]")
538+
{
539+
const OperatorNode hadamard{"Hadamard"};
540+
const OperatorNode globalPhase{"GlobalPhase"};
541+
542+
const WeightedGateset gateset{{{globalPhase, 1.0}}};
543+
544+
const std::vector<RuleNode> rules{
545+
{"hadamard_to_globalPhase", hadamard, {}},
546+
};
547+
548+
const DecompositionGraph graph({hadamard}, gateset, rules);
549+
DecompositionSolver solver(graph);
550+
const auto result = solver.solve();
551+
REQUIRE(result.size() == 1);
552+
const auto &chosen_rule = result.at(hadamard);
553+
REQUIRE_FALSE(chosen_rule.isBasis);
554+
REQUIRE(chosen_rule.ruleName == "hadamard_to_globalPhase");
555+
REQUIRE(chosen_rule.inputs.empty());
556+
REQUIRE(chosen_rule.totalCost == 0.0);
557+
}

0 commit comments

Comments
 (0)