From 8ea15af524db53d05bec142dad0e3283cecc344f Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Mon, 20 Jul 2026 16:38:50 -0400 Subject: [PATCH 01/36] add pass id support --- .../lit/GraphDecomposition/TestGraphOpId.mlir | 31 +++++++++++++++++++ .../graph_decomposition.cpp | 8 +++++ 2 files changed, 39 insertions(+) create mode 100644 frontend/test/lit/GraphDecomposition/TestGraphOpId.mlir diff --git a/frontend/test/lit/GraphDecomposition/TestGraphOpId.mlir b/frontend/test/lit/GraphDecomposition/TestGraphOpId.mlir new file mode 100644 index 0000000000..8536a5fa85 --- /dev/null +++ b/frontend/test/lit/GraphDecomposition/TestGraphOpId.mlir @@ -0,0 +1,31 @@ +// 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. + +// Test that graph-decomposition succeeds when using graphOpIds + +// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=PauliX=1.0 alt-decomps=Hadamard=my_decomp})' %s | FileCheck %s + +func.func @circuit(%q: !quantum.bit) -> !quantum.bit { + // CHECK-NOT: Hadamard + // CHECK: PauliX + // CHECK: PauliX + %out = quantum.custom "Hadamard"() %q: !quantum.bit + return %out: !quantum.bit +} + +func.func private @my_decomp(%q: !quantum.bit) -> !quantum.bit attributes {target_gate="Hadamard[][1]{}"} { + %q0 = quantum.custom "PauliX"() %q : !quantum.bit + %q1 = quantum.custom "PauliX"() %q0 : !quantum.bit + return %q1 : !quantum.bit +} diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp index 39929153a4..027b54d8fe 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp @@ -453,6 +453,10 @@ struct GraphDecompositionPass : public impl::GraphDecompositionPassBase(op.getOperation())) { + node.id = decompGate.getGraphOpId(); + } + operators.push_back(node); }); } @@ -476,6 +480,10 @@ struct GraphDecompositionPass : public impl::GraphDecompositionPassBase Date: Wed, 22 Jul 2026 12:34:37 -0400 Subject: [PATCH 02/36] migrate pass & solver to ID --- doc/releases/changelog-dev.md | 1 + .../DecompGraphSolver/DGBuilder.cpp | 14 +- .../DecompGraphSolver/DGBuilder.hpp | 1 + .../DecompGraphSolver/DGSolver.cpp | 4 +- .../DecompGraphSolver/DGSolver.hpp | 2 - .../DecompGraphSolver/DGTypes.hpp | 66 ++----- .../DecompGraphSolver/DGUtils.hpp | 1 + .../graph_decomposition.cpp | 22 +-- .../Test_DecompGraphCore.cpp | 139 +++++--------- .../Test_DecompGraphSolver.cpp | 173 +++++++----------- 10 files changed, 143 insertions(+), 280 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 7e66ec4533..72a6ad868a 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -277,6 +277,7 @@ [(#2983)](https://github.com/PennyLaneAI/catalyst/pull/2983) [(#3022)](https://github.com/PennyLaneAI/catalyst/pull/3022) [(#3039)](https://github.com/PennyLaneAI/catalyst/pull/3039) + [(#3046)](https://github.com/PennyLaneAI/catalyst/pull/3046) * The `graph-decomposition` pass eliminates three redundant IR manipulations: the cloning, removal, and re-insertion of user rules. This optimization is particularly diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp index 7d2724941c..97187aa0a3 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp @@ -18,13 +18,23 @@ #include "DGBuilder.hpp" +#include #include #include +#include +#include +#include +#include #include +#include -#include "DGUtils.hpp" +#include "boost/graph/adjacency_list.hpp" +#include "boost/graph/detail/adjacency_list.hpp" +#include "boost/graph/graph_selectors.hpp" +#include "boost/graph/graph_traits.hpp" -#include +#include "DGTypes.hpp" +#include "DGUtils.hpp" using namespace DecompGraph::Core; diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.hpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.hpp index 984dd6a970..56fae1952b 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.hpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.hpp @@ -24,6 +24,7 @@ #pragma once +#include #include #include diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.cpp index c83ccddce5..35b7d760ae 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.cpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.cpp @@ -18,12 +18,14 @@ #include "DGSolver.hpp" -#include #include +#include #include +#include #include #include "DGTypes.hpp" +#include "DGUtils.hpp" using namespace DecompGraph::Core; diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.hpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.hpp index 9e92a32ebd..030481b0b7 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.hpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGSolver.hpp @@ -25,14 +25,12 @@ #pragma once -#include #include #include #include #include "DGBuilder.hpp" #include "DGTypes.hpp" -#include "DGUtils.hpp" namespace DecompGraph::Solver { diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp index db092e74c8..cd4aaeb844 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp @@ -53,39 +53,16 @@ namespace DecompGraph::Core { * when adding support for operators with dynamic numbers of wires/params. */ struct OperatorNode { - std::string name; - int numWires{-1}; - int numParams{-1}; + std::string id; bool adjoint{false}; - // Optional static arguments for operators that require additional data. + // optional params, primarily for debug use + std::string name{""}; + int numWires{-1}; + int numParams{-1}; std::unordered_map staticNamedArgs{}; - std::string id{""}; - bool operator==(const OperatorNode &other) const - { - // id match - if (!id.empty() && !other.id.empty() && id == other.id) { - return true; - } - // legacy fallback if either op is missing ID - - // For equality, we consider numWires and numParams conditionally equal - // if they are not set to -1 (which indicates a wildcard that can match any value). - const bool default_wires = - (numWires == -1 || other.numWires == -1 || numWires == other.numWires); - const bool default_params = - (numParams == -1 || other.numParams == -1 || numParams == other.numParams); - - // Static arguments are optional: if either side has no static args, they - // are treated as matching (wildcard). When both sides provide entries, the maps must - // be equal element-wise for the operators to be considered equivalent. - const bool static_args_match = staticNamedArgs.empty() || other.staticNamedArgs.empty() || - staticNamedArgs == other.staticNamedArgs; - - return name == other.name && default_wires && default_params && adjoint == other.adjoint && - static_args_match; - } + bool operator==(const OperatorNode &other) const { return id == other.id; } bool operator!=(const OperatorNode &other) const { return !(*this == other); } }; @@ -107,11 +84,7 @@ struct OperatorNode { struct OperatorNodeHash { std::size_t operator()(const OperatorNode &node) const { - // prefer id if available - if (!node.id.empty()) { - return std::hash{}(node.id); - } - return std::hash{}(node.name); + return std::hash{}(node.id); } }; @@ -119,33 +92,16 @@ struct OperatorNodeHash { * @brief This represents the weighted target gateset for the graph decomposition problem. */ struct WeightedGateset { + // TODO: using ID here mandates that gatesets specify all legal IDs, rather than generic class + // like "PauliRot". This should be updated to work on generic names std::unordered_map ops; - [[nodiscard]] bool contains(const OperatorNode &op) const - { - // hash match - if (ops.find(op) != ops.end()) { - return true; - } - - // use op-matching if hashes failed (could be id vs name) - for (auto [gatesetOp, cost] : ops) { - if (gatesetOp == op) { - return true; - } - } - - return false; - } + [[nodiscard]] bool contains(const OperatorNode &op) const { return ops.find(op) != ops.end(); } [[nodiscard]] double getCost(const OperatorNode &op) const { auto it = ops.find(op); - if (it != ops.end()) { - return it->second; - } - - return std::numeric_limits::infinity(); + return it != ops.end() ? it->second : std::numeric_limits::infinity(); } }; diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGUtils.hpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGUtils.hpp index e962638a33..27c66217ef 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGUtils.hpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGUtils.hpp @@ -25,6 +25,7 @@ #include #include #include +#include #include #include "DGTypes.hpp" diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp index 027b54d8fe..c43ea88555 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp @@ -422,7 +422,7 @@ struct GraphDecompositionPass : public impl::GraphDecompositionPassBase(op.getOperation())) { - node.name = customOp.getGateName().str(); - } - // Name handling for non-custom ops - else { - std::string name = op->getName().stripDialect().str(); - if (name == "gphase") { - name = "GlobalPhase"; - } - else if (name == "paulirot") { - name = cast(op.getOperation()).getGraphOpId(); - } - node.name = name; - } + node.name = op.getOperatorName(); + node.id = op.getGraphOpId(); if (auto paramOp = llvm::dyn_cast(op.getOperation())) { @@ -453,10 +441,6 @@ struct GraphDecompositionPass : public impl::GraphDecompositionPassBase(op.getOperation())) { - node.id = decompGate.getGraphOpId(); - } - operators.push_back(node); }); } diff --git a/mlir/unittests/DecompGraphSolver/Test_DecompGraphCore.cpp b/mlir/unittests/DecompGraphSolver/Test_DecompGraphCore.cpp index 5d8db73d22..89f02b0049 100644 --- a/mlir/unittests/DecompGraphSolver/Test_DecompGraphCore.cpp +++ b/mlir/unittests/DecompGraphSolver/Test_DecompGraphCore.cpp @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include + #include "DGTypes.hpp" #include @@ -24,77 +26,60 @@ using namespace DecompGraph::Core; TEST_CASE("Test OperatorNode construction", "[DecompGraph::Core]") { - const OperatorNode op1{"H", 1, 0, false}; - const OperatorNode op2{"CNOT", 2, 0, false}; - const OperatorNode op3{"RX", 1, 1, false}; - const OperatorNode op4{"RZ", 1, 1, true}; - - REQUIRE(op1.name == "H"); - REQUIRE(op1.numWires == 1); - REQUIRE(op1.numParams == 0); - REQUIRE(op1.adjoint == false); - - REQUIRE(op2.name == "CNOT"); - REQUIRE(op2.numWires == 2); - REQUIRE(op2.numParams == 0); - REQUIRE(op2.adjoint == false); - - REQUIRE(op3.name == "RX"); - REQUIRE(op3.numWires == 1); - REQUIRE(op3.numParams == 1); - REQUIRE(op3.adjoint == false); - - REQUIRE(op4.name == "RZ"); - REQUIRE(op4.numWires == 1); - REQUIRE(op4.numParams == 1); - REQUIRE(op4.adjoint == true); + const OperatorNode h{"Hadamard[][1]{}"}; + const OperatorNode cnot{"CNOT[][2]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + + REQUIRE(h.id == "Hadamard[][1]{}"); + REQUIRE(cnot.id == "CNOT[][2]{}"); + REQUIRE(rx.id == "RX[f64][1]{}"); + REQUIRE(rz.id == "RZ[f64][1]{}"); } TEST_CASE("Test OperatorNode equality operator", "[DecompGraph::Core]") { - const OperatorNode op1{"H", 1, 0, false}; - const OperatorNode op2{"H", 1, 0, false}; - const OperatorNode op3{"H", 1, 0, true}; - const OperatorNode op4{"CNOT", 2, 0, false}; - - REQUIRE(op1 == op2); - REQUIRE_FALSE(op1 == op3); - REQUIRE_FALSE(op1 == op4); + const OperatorNode h1{"Hadamard[][1]{}"}; + const OperatorNode h2{"Hadamard[][1]{}"}; + const OperatorNode cnot{"CNOT[][2]{}"}; + + REQUIRE(h1 == h2); + REQUIRE(h1 != cnot); } TEST_CASE("Test OperatorNodeHash", "[DecompGraph::Core]") { - const OperatorNode op1{"H", 1, 0, false}; - const OperatorNode op2{"H", 1, 0, false}; - const OperatorNode op3{"H", 1, 0, true}; - const OperatorNode op4{"CNOT", 2, 0, false}; + const OperatorNode h1{"Hadamard[][1]{}"}; + const OperatorNode h2{"Hadamard[][1]{}"}; + const OperatorNode h3{"Hadamard[][1]{}"}; + const OperatorNode cnot{"CNOT[][2]{}"}; const OperatorNodeHash hashFunc; - REQUIRE(hashFunc(op1) == hashFunc(op2)); - REQUIRE(hashFunc(op1) == hashFunc(op3)); - REQUIRE(hashFunc(op1) != hashFunc(op4)); + REQUIRE(hashFunc(h1) == hashFunc(h2)); + REQUIRE(hashFunc(h1) == hashFunc(h3)); + REQUIRE(hashFunc(h1) != hashFunc(cnot)); } TEST_CASE("Test OperatorNode in unordered_map", "[DecompGraph::Core]") { std::unordered_map opMap; - const OperatorNode op1{"H", 1, 0, false}; - const OperatorNode op2{"CNOT", 2, 0, false}; - const OperatorNode op3{"RX", 1, 1, false}; + const OperatorNode h{"Hadamard[][1]{}"}; + const OperatorNode cnot{"CNOT[][2]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; - opMap[op1] = 1.0; - opMap[op2] = 2.0; + opMap[h] = 1.0; + opMap[cnot] = 2.0; - REQUIRE(opMap[op1] == 1.0); - REQUIRE(opMap[op2] == 2.0); - REQUIRE(opMap.find(op3) == opMap.end()); + REQUIRE(opMap[h] == 1.0); + REQUIRE(opMap[cnot] == 2.0); + REQUIRE(opMap.find(rx) == opMap.end()); } TEST_CASE("Test RuleNode construction", "[DecompGraph::Core]") { - const auto h = OperatorNode{"H"}; - const auto rz = OperatorNode{"RZ"}; - const auto rx = OperatorNode{"RX"}; + const OperatorNode h{"Hadamard[][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; const RuleNode h_to_rz_rx_rz{"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}; REQUIRE(h_to_rz_rx_rz.name == "h_to_rz_rx_rz"); @@ -108,9 +93,9 @@ TEST_CASE("Test RuleNode construction", "[DecompGraph::Core]") TEST_CASE("Test WeightedGateset construction and contains", "[DecompGraph::Core]") { - const OperatorNode h{"H"}; - const OperatorNode cnot{"CNOT"}; - const OperatorNode rx{"RX"}; + const OperatorNode h{"Hadamard[][1]{}"}; + const OperatorNode cnot{"CNOT[][2]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; const WeightedGateset gateset{{{h, 1.0}, {cnot, 2.0}}}; @@ -124,9 +109,9 @@ TEST_CASE("Test WeightedGateset construction and contains", "[DecompGraph::Core] TEST_CASE("Test ChosenDecompRule construction", "[DecompGraph::Core]") { - const OperatorNode h{"H"}; - const OperatorNode rz{"RZ"}; - const OperatorNode rx{"RX"}; + const OperatorNode h{"Hadamard[][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; const RuleTerm term1{rz, 2}; const RuleTerm term2{rx, 1}; @@ -146,45 +131,3 @@ TEST_CASE("Test ChosenDecompRule construction", "[DecompGraph::Core]") REQUIRE(chosenRule.basisCounts[rz] == 2); REQUIRE(chosenRule.basisCounts[rx] == 1); } - -TEST_CASE("Test graphOpId Support", "[DecompGraph::Core]") -{ - // comparing an op without an ID should fallback to legacy match - const OperatorNode h{"H"}; - const OperatorNode hId{"H", -1, -1, false, {}, "H[][1]{}"}; - - REQUIRE(h == hId); - - // id + legacy match with non-wildcard params - const OperatorNode x{"X", 1, 0}; - const OperatorNode xId{"X", -1, -1, false, {}, "X[][1]{}"}; - - REQUIRE(x == xId); - - // id nodes should fail legacy match if params differ - const OperatorNode op1{"op", 1, 1}; - const OperatorNode op2{"op", 2, 2, false, {}, "op[f64,f64][2]{}"}; - - REQUIRE_FALSE(op1 == op2); - - // nodes with same ids should match - const OperatorNode pr1{"PauliRot", -1, -1, false, {}, "PauliRot[f64][2]{pauli_word:XX}"}; - const OperatorNode pr2{"PauliRot", -1, -1, false, {}, "PauliRot[f64][2]{pauli_word:XX}"}; - - REQUIRE(pr1 == pr2); - - // nodes with matching ids should ignore other parameters (id is source of truth) - const OperatorNode id1{"name1", 1, 1, false, {}, "sameID"}; - const OperatorNode id2{"name2", 2, 2, true, {}, "sameID"}; - - REQUIRE(id1 == id2); - - // Unit test for `OperatorNodeHash`. Check that the hash function prefers ID over name - const OperatorNode hash1{"name", -1, -1, false, {}, "id"}; - const OperatorNode hash2{"name2", 1, 1, true, {}, "id"}; - const OperatorNode hash3{"name2", 1, 1, true, {}}; - - const OperatorNodeHash hashFunc; - REQUIRE(hashFunc(hash1) == hashFunc(hash2)); - REQUIRE(hashFunc(hash2) != hashFunc(hash3)); -} diff --git a/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp b/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp index 46a2602791..2b2a3216eb 100644 --- a/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp +++ b/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp @@ -13,7 +13,9 @@ // limitations under the License. #include -#include +#include +#include +#include #include "DGBuilder.hpp" #include "DGSolver.hpp" @@ -31,9 +33,9 @@ using namespace DecompGraph::Solver; TEST_CASE("Test DecompositionGraph construction", "[DecompGraph::Solver]") { - const auto h = OperatorNode{"H", 1, 0, false}; - const auto rz = OperatorNode{"RZ", 1, 1, false}; - const auto rx = OperatorNode{"RX", 1, 1, false}; + const auto h = OperatorNode{"H[][1]{}"}; + const auto rz = OperatorNode{"RZ[f64][1]{}"}; + const auto rx = OperatorNode{"RX[f64][1]{}"}; const WeightedGateset gateset{{{rz, 1.0}, {rx, 2.0}}}; @@ -66,8 +68,8 @@ TEST_CASE("Test DecompositionGraph construction", "[DecompGraph::Solver]") TEST_CASE("Test DecompositionSolver solve method with incomplete gates in Gateset", "[DecompGraph::Solver]") { - const auto h = OperatorNode{"H", 1, 0, false}; - const auto h_gateset = OperatorNode{"H"}; + const auto h = OperatorNode{"H[][1]{}"}; + const auto h_gateset = OperatorNode{"H[][1]{}"}; const WeightedGateset gateset{{{h_gateset, 1.0}}}; const std::vector rules{ {"h_to_h", h, {{h, 1}}}, @@ -91,8 +93,8 @@ TEST_CASE("Test DecompositionSolver solve method with incomplete gates in Gatese TEST_CASE("Do not solve for target gates", "[DecompGraph::Solver]") { - const auto h = OperatorNode{"H", 1, 0, false}; - const auto rz = OperatorNode{"RZ", 1, 1, false}; + const auto h = OperatorNode{"H[][1]{}"}; + const auto rz = OperatorNode{"RZ[f64][1]{}"}; const WeightedGateset gateset{{{h, 2.0}, {rz, 1.0}}}; @@ -112,9 +114,9 @@ TEST_CASE("Do not solve for target gates", "[DecompGraph::Solver]") TEST_CASE("Test DecompositionGraph copy and move semantics", "[DecompGraph::Solver]") { - const auto h = OperatorNode{"H", 1, 0, false}; - const auto rz = OperatorNode{"RZ", 1, 1, false}; - const auto rx = OperatorNode{"RX", 1, 1, false}; + const auto h = OperatorNode{"H[][1]{}"}; + const auto rz = OperatorNode{"RZ[f64][1]{}"}; + const auto rx = OperatorNode{"RX[f64][1]{}"}; const WeightedGateset gateset{{{rz, 1.0}, {rx, 2.0}}}; @@ -155,10 +157,10 @@ TEST_CASE("Test DecompositionGraph copy and move semantics", "[DecompGraph::Solv TEST_CASE("Test DecompositionGraph lookup and counting", "[DecompGraph::Solver]") { - const OperatorNode h{"H", 1, 0, false}; - const OperatorNode rz{"RZ", 1, 1, false}; - const OperatorNode rx{"RX", 1, 1, false}; - const OperatorNode ry{"RY", 1, 1, false}; + const OperatorNode h{"H[][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode ry{"RY[f64][1]{}"}; const WeightedGateset gateset{{{rz, 1.0}, {ry, 2.0}, {rx, 3.0}}}; @@ -193,13 +195,13 @@ TEST_CASE("Test DecompositionGraph lookup and counting", "[DecompGraph::Solver]" TEST_CASE("Test the graph construction with realistic ops and multiple rules from PennyLane", "[DecompGraph::Solver]") { - const OperatorNode h{"H", 1, 0, false}; - const OperatorNode rz{"RZ", 1, 1, false}; - const OperatorNode rx{"RX", 1, 1, false}; - const OperatorNode ry{"RY", 1, 1, false}; - const OperatorNode cnot{"CNOT", 2, 0, false}; - const OperatorNode swap{"SWAP", 2, 0, false}; - const OperatorNode customBellOp{"BellOp", 2, 0, false}; + const OperatorNode h{"H[][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode ry{"RY[f64][1]{}"}; + const OperatorNode cnot{"CNOT[][2]{}"}; + const OperatorNode swap{"SWAP[][2]{}"}; + const OperatorNode customBellOp{"BellOp[][2]{}"}; const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}, {cnot, 5.0}}}; @@ -219,10 +221,10 @@ TEST_CASE("Test the graph construction with realistic ops and multiple rules fro TEST_CASE("Test DecompositionSolver with one single operator", "[DecompGraph::Solver]") { - const OperatorNode h{"H", 1, 0, false}; - const OperatorNode rz{"RZ", 1, 1, false}; - const OperatorNode rx{"RX", 1, 1, false}; - const OperatorNode ry{"RY", 1, 1, false}; + const OperatorNode h{"H[][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode ry{"RY[f64][1]{}"}; const WeightedGateset gateset{{{rz, 1.0}, {ry, 2.0}, {rx, 3.0}}}; @@ -258,13 +260,13 @@ TEST_CASE("Test DecompositionSolver with one single operator", "[DecompGraph::So TEST_CASE("Test the graph solver with intermediate ops and multiple rules", "[DecompGraph::Solver]") { - const OperatorNode h{"H", 1, 0, false}; - const OperatorNode rz{"RZ", 1, 1, false}; - const OperatorNode rx{"RX", 1, 1, false}; - const OperatorNode ry{"RY", 1, 1, false}; - const OperatorNode cnot{"CNOT", 2, 0, false}; - const OperatorNode swap{"SWAP", 2, 0, false}; - const OperatorNode customBellOp{"BellOp", 2, 0, false}; + const OperatorNode h{"H[][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode ry{"RY[f64][1]{}"}; + const OperatorNode cnot{"CNOT[][2]{}"}; + const OperatorNode swap{"SWAP[][2]{}"}; + const OperatorNode customBellOp{"BellOp[][2]{}"}; const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}, {cnot, 5.0}}}; @@ -316,8 +318,8 @@ TEST_CASE("Test the graph solver with intermediate ops and multiple rules", "[De TEST_CASE("Test GraphSolveError for unsolvable operator", "[DecompGraph::Solver]") { - const OperatorNode h{"H", 1, 0, false}; - const OperatorNode rz{"RZ", 1, 1, false}; + const OperatorNode h{"H[][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; const WeightedGateset gateset{{{rz, 1.0}}}; @@ -333,7 +335,7 @@ TEST_CASE("Test GraphSolveError for unsolvable operator", "[DecompGraph::Solver] TEST_CASE("Test GraphSolveError for cyclic decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H", 1, 0, false}; + const OperatorNode h{"H[][1]{}"}; const WeightedGateset gateset{}; @@ -349,9 +351,9 @@ TEST_CASE("Test GraphSolveError for cyclic decomposition", "[DecompGraph::Solver TEST_CASE("Test PauliX -> GlobalPhase(1), RX(1) decomposition", "[DecompGraph::Solver]") { - const OperatorNode x{"X"}; - const OperatorNode globalPhase{"GlobalPhase"}; - const OperatorNode rx{"RX"}; + const OperatorNode x{"X[][1]{}"}; + const OperatorNode globalPhase{"GlobalPhase[][]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; const WeightedGateset gateset{{{globalPhase, 1.0}, {rx, 1.0}}}; @@ -379,14 +381,14 @@ TEST_CASE("Test PauliX -> GlobalPhase(1), RX(1) decomposition", "[DecompGraph::S TEST_CASE("Test cyclic decomposition with multiple rules for the same operator", "[DecompGraph::Solver]") { - const OperatorNode hadamard{"Hadamard"}; - const OperatorNode globalPhase{"GlobalPhase"}; - const OperatorNode rx{"RX"}; - const OperatorNode rz{"RZ"}; - const OperatorNode ry{"RY"}; - const OperatorNode changeOpBasis{"ChangeOpBasis"}; - const OperatorNode pauliRot{"PauliRot"}; - const OperatorNode rot{"Rot"}; + const OperatorNode hadamard{"Hadamard[][1]{}"}; + const OperatorNode globalPhase{"GlobalPhase[][]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode ry{"RY[f64][1]{}"}; + const OperatorNode changeOpBasis{"ChangeOpBasis[][2]{}"}; + const OperatorNode pauliRot{"PauliRot[f64][2]{pauli_word:XY}"}; + const OperatorNode rot{"Rot[f64,f64,f64][3]{}"}; const std::vector rules{ {"__builtin__ry_to_rz_cliff", ry, {{changeOpBasis, 1}}}, @@ -417,9 +419,9 @@ TEST_CASE("Test cyclic decomposition with multiple rules for the same operator", TEST_CASE("Test GraphBuilder with fixed decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H"}; - const OperatorNode rz{"RZ"}; - const OperatorNode rx{"RX"}; + const OperatorNode h{"H[][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}}}; @@ -439,9 +441,9 @@ TEST_CASE("Test GraphBuilder with fixed decomposition", "[DecompGraph::Solver]") TEST_CASE("Test GraphBuilder with alternative decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H"}; - const OperatorNode rz{"RZ"}; - const OperatorNode rx{"RX"}; + const OperatorNode h{"H[][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}}}; @@ -459,9 +461,9 @@ TEST_CASE("Test GraphBuilder with alternative decomposition", "[DecompGraph::Sol TEST_CASE("Test GraphSolver with fixed decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H"}; - const OperatorNode rz{"RZ"}; - const OperatorNode rx{"RX"}; + const OperatorNode h{"H[][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; const WeightedGateset gateset{{{rz, 3.0}, {rx, 1.0}}}; @@ -485,9 +487,9 @@ TEST_CASE("Test GraphSolver with fixed decomposition", "[DecompGraph::Solver]") TEST_CASE("Test GraphSolver with alternative decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H"}; - const OperatorNode rz{"RZ"}; - const OperatorNode rx{"RX"}; + const OperatorNode h{"H[][1]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode rx{"RX[f64][1]{}"}; const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}}}; @@ -510,9 +512,9 @@ TEST_CASE("Test GraphSolver with alternative decomposition", "[DecompGraph::Solv TEST_CASE("Test GraphSolver with MultiRZ decompositions", "[DecompGraph::Solver]") { - const OperatorNode multiRZ3{"MultiRZ3"}; - const OperatorNode multiRZ5{"MultiRZ5"}; - const OperatorNode rz{"RZ"}; + const OperatorNode multiRZ3{"MultiRZ[f64][3]{}"}; + const OperatorNode multiRZ5{"MultiRZ[f64][5]{}"}; + const OperatorNode rz{"RZ[f64][1]{}"}; const WeightedGateset gateset{{{rz, 1.0}}}; @@ -537,8 +539,8 @@ TEST_CASE("Test GraphSolver with MultiRZ decompositions", "[DecompGraph::Solver] TEST_CASE("Test GraphSolver with empty decomposition rules", "[DecompGraph::Solver]") { - const OperatorNode hadamard{"Hadamard"}; - const OperatorNode globalPhase{"GlobalPhase"}; + const OperatorNode hadamard{"Hadamard[][1]{}"}; + const OperatorNode globalPhase{"GlobalPhase[][]{}"}; const WeightedGateset gateset{{{globalPhase, 1.0}}}; @@ -557,47 +559,12 @@ TEST_CASE("Test GraphSolver with empty decomposition rules", "[DecompGraph::Solv REQUIRE(chosen_rule.totalCost == 0.0); } -TEST_CASE("Test GraphSolver with PauliRot specialized by static argument pauli_word", - "[DecompGraph::Solver]") -{ - // 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 hadamard{"Hadamard", 1, 0, false}; - const OperatorNode multiRZ{"MultiRZ", 1, 1, false}; - - const WeightedGateset gateset{{{hadamard, 1.0}, {multiRZ, 1.0}}}; - - const std::vector rules{ - {"_pauli_rot_decomposition_X", pauliRotRuleOutput, {{hadamard, 2}, {multiRZ, 1}}}, - }; - - const DecompositionGraph graph({pauliRotQuery}, gateset, rules); - DecompositionSolver solver(graph); - const auto result = solver.solve(); - - REQUIRE(result.find(pauliRotQuery) != result.end()); - const auto &chosen = result.at(pauliRotQuery); - REQUIRE_FALSE(chosen.isBasis); - REQUIRE(chosen.ruleName == "_pauli_rot_decomposition_X"); - REQUIRE(chosen.totalCost == 1.0 * 2 + 1.0 * 1); - REQUIRE(chosen.basisCounts.at(hadamard) == 2); - REQUIRE(chosen.basisCounts.at(multiRZ) == 1); - - const OperatorNode pauliRotQueryY{"PauliRot", 1, 1, false, {{"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 pauliRotNoArgs{"PauliRot", 1, 1, false}; + const OperatorNode pauliRotX{"PauliRot[f64][1]{pauli_word:X}"}; + const OperatorNode pauliRotX2{"PauliRot[f64][1]{pauli_word:X}"}; + const OperatorNode pauliRotY{"PauliRot[f64][1]{pauli_word:Y}"}; - REQUIRE(pauliRotX == pauliRotXWildcard); + REQUIRE(pauliRotX == pauliRotX2); REQUIRE_FALSE(pauliRotX == pauliRotY); - REQUIRE(pauliRotX == pauliRotNoArgs); } From 6a5b47004ef085b5ca67c5df0f9e43b15fb306e0 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 22 Jul 2026 14:34:22 -0400 Subject: [PATCH 03/36] use name for gateset --- .../DecompGraphSolver/DGBuilder.cpp | 9 ++------- .../DecompGraphSolver/DGTypes.hpp | 14 ++++++++------ 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp index 97187aa0a3..a6544c17e8 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGBuilder.cpp @@ -173,11 +173,6 @@ struct DecompositionGraph::Impl { registerOp(op); } - // Register all target gates - for (const auto &[op, _] : gateset.ops) { - registerOp(op); - } - // Register all rules for (RuleId ruleId = 0; ruleId < rules.size(); ruleId++) { const auto &rule = rules[ruleId]; @@ -322,8 +317,8 @@ void DecompositionGraph::showGraph() const // Show target gateset std::cerr << "Target Gateset:\n"; - for (const auto &[op, cost] : impl->gateset.ops) { - std::cerr << " " << print_op(op) << " with cost " << cost << "\n"; + for (const auto &[name, cost] : impl->gateset.ops) { + std::cerr << " " << name << " with cost " << cost << "\n"; } } diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp index cd4aaeb844..a7c800655b 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp @@ -54,10 +54,11 @@ namespace DecompGraph::Core { */ struct OperatorNode { std::string id; + std::string name; // name is required for gateset checking + bool adjoint{false}; // optional params, primarily for debug use - std::string name{""}; int numWires{-1}; int numParams{-1}; std::unordered_map staticNamedArgs{}; @@ -92,15 +93,16 @@ struct OperatorNodeHash { * @brief This represents the weighted target gateset for the graph decomposition problem. */ struct WeightedGateset { - // TODO: using ID here mandates that gatesets specify all legal IDs, rather than generic class - // like "PauliRot". This should be updated to work on generic names - std::unordered_map ops; + std::unordered_map ops; - [[nodiscard]] bool contains(const OperatorNode &op) const { return ops.find(op) != ops.end(); } + [[nodiscard]] bool contains(const OperatorNode &op) const + { + return ops.find(op.name) != ops.end(); + } [[nodiscard]] double getCost(const OperatorNode &op) const { - auto it = ops.find(op); + auto it = ops.find(op.name); return it != ops.end() ? it->second : std::numeric_limits::infinity(); } }; From 72fc4755a8b4cbcaa08a958b62f4bf85eca3ab3a Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Wed, 22 Jul 2026 14:35:57 -0400 Subject: [PATCH 04/36] add names to gateset --- .../Transforms/GraphDecomposition/graph_decomposition.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp index c43ea88555..beb297a981 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp @@ -220,7 +220,7 @@ struct GraphDecompositionPass : public impl::GraphDecompositionPassBase Date: Wed, 22 Jul 2026 14:36:14 -0400 Subject: [PATCH 05/36] add graph solution to debug --- .../Transforms/GraphDecomposition/graph_decomposition.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp index beb297a981..48ce6e108e 100644 --- a/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp +++ b/mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp @@ -60,6 +60,7 @@ #include "DGBuilder.hpp" #include "DGSolver.hpp" #include "DGTypes.hpp" +#include "DGUtils.hpp" #include "DecompUtils.hpp" #define DEBUG_TYPE "graph-decomposition" @@ -139,6 +140,7 @@ struct GraphDecompositionPass : public impl::GraphDecompositionPassBase Date: Wed, 22 Jul 2026 14:36:35 -0400 Subject: [PATCH 06/36] update tests --- .../GraphDecomposition/TestAltDecomps.mlir | 4 +- .../Test_DecompGraphCore.cpp | 53 +++--- .../Test_DecompGraphSolver.cpp | 171 +++++++++--------- 3 files changed, 117 insertions(+), 111 deletions(-) diff --git a/frontend/test/lit/GraphDecomposition/TestAltDecomps.mlir b/frontend/test/lit/GraphDecomposition/TestAltDecomps.mlir index bf0323a4f9..439fb41432 100644 --- a/frontend/test/lit/GraphDecomposition/TestAltDecomps.mlir +++ b/frontend/test/lit/GraphDecomposition/TestAltDecomps.mlir @@ -34,7 +34,7 @@ func.func @circuit() -> !quantum.bit { } // CHECK-LABEL: y_to_ry -func.func @y_to_ry(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="PauliY"} { +func.func @y_to_ry(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="PauliY[][1]{}", resources = { operations = {"RY[f64][1]{}"=1, "GlobalPhase[][]{}"=1}}} { %pi = arith.constant 3.14 : f64 %negpiby2 = arith.constant -1.57 : f64 %q1 = quantum.custom "RY"(%pi) %q0 : !quantum.bit @@ -43,7 +43,7 @@ func.func @y_to_ry(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate=" } // CHECK-LABEL: y_to_x_z -func.func @y_to_x_z(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="PauliY"} { +func.func @y_to_x_z(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="PauliY[][1]{}", resources = { operations = {"PauliX[][1]{}"=1, "PauliZ[][1]{}"=1}}} { %q1 = quantum.custom "PauliX"() %q0 : !quantum.bit %q2 = quantum.custom "PauliZ"() %q1 : !quantum.bit return %q2 : !quantum.bit diff --git a/mlir/unittests/DecompGraphSolver/Test_DecompGraphCore.cpp b/mlir/unittests/DecompGraphSolver/Test_DecompGraphCore.cpp index 89f02b0049..376bb41ce8 100644 --- a/mlir/unittests/DecompGraphSolver/Test_DecompGraphCore.cpp +++ b/mlir/unittests/DecompGraphSolver/Test_DecompGraphCore.cpp @@ -26,22 +26,27 @@ using namespace DecompGraph::Core; TEST_CASE("Test OperatorNode construction", "[DecompGraph::Core]") { - const OperatorNode h{"Hadamard[][1]{}"}; - const OperatorNode cnot{"CNOT[][2]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode h{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode cnot{"CNOT[][2]{}", "CNOT"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; REQUIRE(h.id == "Hadamard[][1]{}"); REQUIRE(cnot.id == "CNOT[][2]{}"); REQUIRE(rx.id == "RX[f64][1]{}"); REQUIRE(rz.id == "RZ[f64][1]{}"); + + REQUIRE(h.name == "Hadamard"); + REQUIRE(cnot.name == "CNOT"); + REQUIRE(rx.name == "RX"); + REQUIRE(rz.name == "RZ"); } TEST_CASE("Test OperatorNode equality operator", "[DecompGraph::Core]") { - const OperatorNode h1{"Hadamard[][1]{}"}; - const OperatorNode h2{"Hadamard[][1]{}"}; - const OperatorNode cnot{"CNOT[][2]{}"}; + const OperatorNode h1{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode h2{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode cnot{"CNOT[][2]{}", "CNOT"}; REQUIRE(h1 == h2); REQUIRE(h1 != cnot); @@ -49,10 +54,10 @@ TEST_CASE("Test OperatorNode equality operator", "[DecompGraph::Core]") TEST_CASE("Test OperatorNodeHash", "[DecompGraph::Core]") { - const OperatorNode h1{"Hadamard[][1]{}"}; - const OperatorNode h2{"Hadamard[][1]{}"}; - const OperatorNode h3{"Hadamard[][1]{}"}; - const OperatorNode cnot{"CNOT[][2]{}"}; + const OperatorNode h1{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode h2{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode h3{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode cnot{"CNOT[][2]{}", "CNOT"}; const OperatorNodeHash hashFunc; REQUIRE(hashFunc(h1) == hashFunc(h2)); @@ -63,9 +68,9 @@ TEST_CASE("Test OperatorNodeHash", "[DecompGraph::Core]") TEST_CASE("Test OperatorNode in unordered_map", "[DecompGraph::Core]") { std::unordered_map opMap; - const OperatorNode h{"Hadamard[][1]{}"}; - const OperatorNode cnot{"CNOT[][2]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode h{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode cnot{"CNOT[][2]{}", "CNOT"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; opMap[h] = 1.0; opMap[cnot] = 2.0; @@ -77,9 +82,9 @@ TEST_CASE("Test OperatorNode in unordered_map", "[DecompGraph::Core]") TEST_CASE("Test RuleNode construction", "[DecompGraph::Core]") { - const OperatorNode h{"Hadamard[][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode h{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; const RuleNode h_to_rz_rx_rz{"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}; REQUIRE(h_to_rz_rx_rz.name == "h_to_rz_rx_rz"); @@ -93,11 +98,11 @@ TEST_CASE("Test RuleNode construction", "[DecompGraph::Core]") TEST_CASE("Test WeightedGateset construction and contains", "[DecompGraph::Core]") { - const OperatorNode h{"Hadamard[][1]{}"}; - const OperatorNode cnot{"CNOT[][2]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode h{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode cnot{"CNOT[][2]{}", "CNOT"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; - const WeightedGateset gateset{{{h, 1.0}, {cnot, 2.0}}}; + const WeightedGateset gateset{{{h.name, 1.0}, {cnot.name, 2.0}}}; REQUIRE(gateset.contains(h)); REQUIRE(gateset.contains(cnot)); @@ -109,9 +114,9 @@ TEST_CASE("Test WeightedGateset construction and contains", "[DecompGraph::Core] TEST_CASE("Test ChosenDecompRule construction", "[DecompGraph::Core]") { - const OperatorNode h{"Hadamard[][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode h{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; const RuleTerm term1{rz, 2}; const RuleTerm term2{rx, 1}; diff --git a/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp b/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp index 2b2a3216eb..25a0216314 100644 --- a/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp +++ b/mlir/unittests/DecompGraphSolver/Test_DecompGraphSolver.cpp @@ -33,11 +33,11 @@ using namespace DecompGraph::Solver; TEST_CASE("Test DecompositionGraph construction", "[DecompGraph::Solver]") { - const auto h = OperatorNode{"H[][1]{}"}; - const auto rz = OperatorNode{"RZ[f64][1]{}"}; - const auto rx = OperatorNode{"RX[f64][1]{}"}; + const auto h = OperatorNode{"H[][1]{}", "Hadamard"}; + const auto rz = OperatorNode{"RZ[f64][1]{}", "RZ"}; + const auto rx = OperatorNode{"RX[f64][1]{}", "RX"}; - const WeightedGateset gateset{{{rz, 1.0}, {rx, 2.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}, {rx.name, 2.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -68,9 +68,9 @@ TEST_CASE("Test DecompositionGraph construction", "[DecompGraph::Solver]") TEST_CASE("Test DecompositionSolver solve method with incomplete gates in Gateset", "[DecompGraph::Solver]") { - const auto h = OperatorNode{"H[][1]{}"}; - const auto h_gateset = OperatorNode{"H[][1]{}"}; - const WeightedGateset gateset{{{h_gateset, 1.0}}}; + const auto h = OperatorNode{"H[][1]{}", "Hadamard"}; + const auto h_gateset = OperatorNode{"H[][1]{}", "Hadamard"}; + const WeightedGateset gateset{{{h_gateset.name, 1.0}}}; const std::vector rules{ {"h_to_h", h, {{h, 1}}}, }; @@ -93,10 +93,10 @@ TEST_CASE("Test DecompositionSolver solve method with incomplete gates in Gatese TEST_CASE("Do not solve for target gates", "[DecompGraph::Solver]") { - const auto h = OperatorNode{"H[][1]{}"}; - const auto rz = OperatorNode{"RZ[f64][1]{}"}; + const auto h = OperatorNode{"H[][1]{}", "Hadamard"}; + const auto rz = OperatorNode{"RZ[f64][1]{}", "RZ"}; - const WeightedGateset gateset{{{h, 2.0}, {rz, 1.0}}}; + const WeightedGateset gateset{{{h.name, 2.0}, {rz.name, 1.0}}}; const std::vector rules{ {"h_to_rz", h, {{rz, 1}}}, @@ -114,11 +114,11 @@ TEST_CASE("Do not solve for target gates", "[DecompGraph::Solver]") TEST_CASE("Test DecompositionGraph copy and move semantics", "[DecompGraph::Solver]") { - const auto h = OperatorNode{"H[][1]{}"}; - const auto rz = OperatorNode{"RZ[f64][1]{}"}; - const auto rx = OperatorNode{"RX[f64][1]{}"}; + const auto h = OperatorNode{"H[][1]{}", "Hadamard"}; + const auto rz = OperatorNode{"RZ[f64][1]{}", "RZ"}; + const auto rx = OperatorNode{"RX[f64][1]{}", "RX"}; - const WeightedGateset gateset{{{rz, 1.0}, {rx, 2.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}, {rx.name, 2.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -157,12 +157,12 @@ TEST_CASE("Test DecompositionGraph copy and move semantics", "[DecompGraph::Solv TEST_CASE("Test DecompositionGraph lookup and counting", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; - const OperatorNode ry{"RY[f64][1]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; + const OperatorNode ry{"RY[f64][1]{}", "RY"}; - const WeightedGateset gateset{{{rz, 1.0}, {ry, 2.0}, {rx, 3.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}, {ry.name, 2.0}, {rx.name, 3.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -181,7 +181,7 @@ TEST_CASE("Test DecompositionGraph lookup and counting", "[DecompGraph::Solver]" for (const auto &rule : graph.getAllRulesFor(h)) { double totalCost = 0.0; for (const auto &input : rule.inputs) { - totalCost += graph.getGateset().ops.at(input.op) * input.multiplicity; + totalCost += graph.getGateset().ops.at(input.op.name) * input.multiplicity; } if (rule.name == "h_to_rz_rx_rz") { REQUIRE(totalCost == 1.0 * 2 + 3.0 * 1); @@ -195,15 +195,15 @@ TEST_CASE("Test DecompositionGraph lookup and counting", "[DecompGraph::Solver]" TEST_CASE("Test the graph construction with realistic ops and multiple rules from PennyLane", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; - const OperatorNode ry{"RY[f64][1]{}"}; - const OperatorNode cnot{"CNOT[][2]{}"}; - const OperatorNode swap{"SWAP[][2]{}"}; - const OperatorNode customBellOp{"BellOp[][2]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; + const OperatorNode ry{"RY[f64][1]{}", "RY"}; + const OperatorNode cnot{"CNOT[][2]{}", "CNOT"}; + const OperatorNode swap{"SWAP[][2]{}", "SWAP"}; + const OperatorNode customBellOp{"BellOp[][2]{}", "BellOp"}; - const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}, {cnot, 5.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}, {rx.name, 3.0}, {cnot.name, 5.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -221,12 +221,12 @@ TEST_CASE("Test the graph construction with realistic ops and multiple rules fro TEST_CASE("Test DecompositionSolver with one single operator", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; - const OperatorNode ry{"RY[f64][1]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; + const OperatorNode ry{"RY[f64][1]{}", "RY"}; - const WeightedGateset gateset{{{rz, 1.0}, {ry, 2.0}, {rx, 3.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}, {ry.name, 2.0}, {rx.name, 3.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -260,15 +260,15 @@ TEST_CASE("Test DecompositionSolver with one single operator", "[DecompGraph::So TEST_CASE("Test the graph solver with intermediate ops and multiple rules", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; - const OperatorNode ry{"RY[f64][1]{}"}; - const OperatorNode cnot{"CNOT[][2]{}"}; - const OperatorNode swap{"SWAP[][2]{}"}; - const OperatorNode customBellOp{"BellOp[][2]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; + const OperatorNode ry{"RY[f64][1]{}", "RY"}; + const OperatorNode cnot{"CNOT[][2]{}", "CNOT"}; + const OperatorNode swap{"SWAP[][2]{}", "SWAP"}; + const OperatorNode customBellOp{"BellOp[][2]{}", "BellOp"}; - const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}, {cnot, 5.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}, {rx.name, 3.0}, {cnot.name, 5.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -318,10 +318,10 @@ TEST_CASE("Test the graph solver with intermediate ops and multiple rules", "[De TEST_CASE("Test GraphSolveError for unsolvable operator", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; - const WeightedGateset gateset{{{rz, 1.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}}}; const std::vector rules{ {"rz_to_rz", rz, {{rz, 1}}}, @@ -335,7 +335,7 @@ TEST_CASE("Test GraphSolveError for unsolvable operator", "[DecompGraph::Solver] TEST_CASE("Test GraphSolveError for cyclic decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; const WeightedGateset gateset{}; @@ -351,11 +351,11 @@ TEST_CASE("Test GraphSolveError for cyclic decomposition", "[DecompGraph::Solver TEST_CASE("Test PauliX -> GlobalPhase(1), RX(1) decomposition", "[DecompGraph::Solver]") { - const OperatorNode x{"X[][1]{}"}; - const OperatorNode globalPhase{"GlobalPhase[][]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode x{"X[][1]{}", "PauliX"}; + const OperatorNode globalPhase{"GlobalPhase[][]{}", "GlobalPhase"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; - const WeightedGateset gateset{{{globalPhase, 1.0}, {rx, 1.0}}}; + const WeightedGateset gateset{{{globalPhase.name, 1.0}, {rx.name, 1.0}}}; const std::vector rules{ {"x_to_globalPhase_rx", x, {{globalPhase, 1}, {rx, 1}}}, @@ -381,14 +381,14 @@ TEST_CASE("Test PauliX -> GlobalPhase(1), RX(1) decomposition", "[DecompGraph::S TEST_CASE("Test cyclic decomposition with multiple rules for the same operator", "[DecompGraph::Solver]") { - const OperatorNode hadamard{"Hadamard[][1]{}"}; - const OperatorNode globalPhase{"GlobalPhase[][]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; - const OperatorNode ry{"RY[f64][1]{}"}; - const OperatorNode changeOpBasis{"ChangeOpBasis[][2]{}"}; - const OperatorNode pauliRot{"PauliRot[f64][2]{pauli_word:XY}"}; - const OperatorNode rot{"Rot[f64,f64,f64][3]{}"}; + const OperatorNode hadamard{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode globalPhase{"GlobalPhase[][]{}", "GlobalPhase"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; + const OperatorNode ry{"RY[f64][1]{}", "RY"}; + const OperatorNode changeOpBasis{"ChangeOpBasis[][2]{}", "ChangeOpBasis"}; + const OperatorNode pauliRot{"PauliRot[f64][2]{pauli_word:XY}", "PauliRot"}; + const OperatorNode rot{"Rot[f64,f64,f64][3]{}", "Rot"}; const std::vector rules{ {"__builtin__ry_to_rz_cliff", ry, {{changeOpBasis, 1}}}, @@ -400,7 +400,7 @@ TEST_CASE("Test cyclic decomposition with multiple rules for the same operator", {"__builtin__hadamard_to_rz_ry", hadamard, {{globalPhase, 1}, {ry, 1}, {rz, 1}}}, }; - const WeightedGateset gateset{{{globalPhase, 1.0}, {rx, 1.0}, {rz, 1.0}}}; + const WeightedGateset gateset{{{globalPhase.name, 1.0}, {rx.name, 1.0}, {rz.name, 1.0}}}; const DecompositionGraph graph({hadamard}, gateset, rules); DecompositionSolver solver(graph); const auto solutions = solver.solve(); @@ -409,7 +409,8 @@ TEST_CASE("Test cyclic decomposition with multiple rules for the same operator", REQUIRE(h_solution.ruleName == "__builtin__hadamard_to_rz_rx"); REQUIRE(h_solution.totalCost == 1.0 * 1 + 1.0 * 1 + 1.0 * 2); - const WeightedGateset gateset2{{{globalPhase, 1.0}, {rx, 1.0}, {rz, 2.0}, {ry, 1.0}}}; + const WeightedGateset gateset2{ + {{globalPhase.name, 1.0}, {rx.name, 1.0}, {rz.name, 2.0}, {ry.name, 1.0}}}; const DecompositionGraph graph2({hadamard}, gateset2, rules); DecompositionSolver solver2(graph2); const auto solutions2 = solver2.solve(); @@ -419,11 +420,11 @@ TEST_CASE("Test cyclic decomposition with multiple rules for the same operator", TEST_CASE("Test GraphBuilder with fixed decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; - const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}, {rx.name, 3.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -441,11 +442,11 @@ TEST_CASE("Test GraphBuilder with fixed decomposition", "[DecompGraph::Solver]") TEST_CASE("Test GraphBuilder with alternative decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; - const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}, {rx.name, 3.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -461,11 +462,11 @@ TEST_CASE("Test GraphBuilder with alternative decomposition", "[DecompGraph::Sol TEST_CASE("Test GraphSolver with fixed decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; - const WeightedGateset gateset{{{rz, 3.0}, {rx, 1.0}}}; + const WeightedGateset gateset{{{rz.name, 3.0}, {rx.name, 1.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -487,11 +488,11 @@ TEST_CASE("Test GraphSolver with fixed decomposition", "[DecompGraph::Solver]") TEST_CASE("Test GraphSolver with alternative decomposition", "[DecompGraph::Solver]") { - const OperatorNode h{"H[][1]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; - const OperatorNode rx{"RX[f64][1]{}"}; + const OperatorNode h{"H[][1]{}", "Hadamard"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; + const OperatorNode rx{"RX[f64][1]{}", "RX"}; - const WeightedGateset gateset{{{rz, 1.0}, {rx, 3.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}, {rx.name, 3.0}}}; const std::vector rules{ {"h_to_rz_rx_rz", h, {{rz, 2}, {rx, 1}}}, @@ -512,11 +513,11 @@ TEST_CASE("Test GraphSolver with alternative decomposition", "[DecompGraph::Solv TEST_CASE("Test GraphSolver with MultiRZ decompositions", "[DecompGraph::Solver]") { - const OperatorNode multiRZ3{"MultiRZ[f64][3]{}"}; - const OperatorNode multiRZ5{"MultiRZ[f64][5]{}"}; - const OperatorNode rz{"RZ[f64][1]{}"}; + const OperatorNode multiRZ3{"MultiRZ[f64][3]{}", "MultiRZ"}; + const OperatorNode multiRZ5{"MultiRZ[f64][5]{}", "MultiRZ"}; + const OperatorNode rz{"RZ[f64][1]{}", "RZ"}; - const WeightedGateset gateset{{{rz, 1.0}}}; + const WeightedGateset gateset{{{rz.name, 1.0}}}; const std::vector rules{ {"multiRZ3_to_rz", multiRZ3, {{rz, 3}}}, @@ -539,10 +540,10 @@ TEST_CASE("Test GraphSolver with MultiRZ decompositions", "[DecompGraph::Solver] TEST_CASE("Test GraphSolver with empty decomposition rules", "[DecompGraph::Solver]") { - const OperatorNode hadamard{"Hadamard[][1]{}"}; - const OperatorNode globalPhase{"GlobalPhase[][]{}"}; + const OperatorNode hadamard{"Hadamard[][1]{}", "Hadamard"}; + const OperatorNode globalPhase{"GlobalPhase[][]{}", "GlobalPhase"}; - const WeightedGateset gateset{{{globalPhase, 1.0}}}; + const WeightedGateset gateset{{{globalPhase.name, 1.0}}}; const std::vector rules{ {"hadamard_to_globalPhase", hadamard, {}}, @@ -561,9 +562,9 @@ TEST_CASE("Test GraphSolver with empty decomposition rules", "[DecompGraph::Solv TEST_CASE("Test OperatorNode equality with staticNamedArgs", "[DecompGraph::Core]") { - const OperatorNode pauliRotX{"PauliRot[f64][1]{pauli_word:X}"}; - const OperatorNode pauliRotX2{"PauliRot[f64][1]{pauli_word:X}"}; - const OperatorNode pauliRotY{"PauliRot[f64][1]{pauli_word:Y}"}; + const OperatorNode pauliRotX{"PauliRot[f64][1]{pauli_word:X}", "PauliRot"}; + const OperatorNode pauliRotX2{"PauliRot[f64][1]{pauli_word:X}", "PauliRot"}; + const OperatorNode pauliRotY{"PauliRot[f64][1]{pauli_word:Y}", "PauliRot"}; REQUIRE(pauliRotX == pauliRotX2); REQUIRE_FALSE(pauliRotX == pauliRotY); From b5672e8dded44f388e4fb137e93e012eb3b4ad77 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 23 Jul 2026 10:25:43 -0400 Subject: [PATCH 07/36] add and centralize decomp utils --- .../precompile_decomposition_rules.py | 195 ++++++++++++++ .../decomposition/python_decompositions.py | 213 ++++++++++++++++ .../decomposition/type_stringify_utils.py | 44 ++++ .../catalyst/device/python_decompositions.py | 90 ------- .../utils/precompile_decomposition_rules.py | 240 ------------------ 5 files changed, 452 insertions(+), 330 deletions(-) create mode 100644 frontend/catalyst/decomposition/precompile_decomposition_rules.py create mode 100644 frontend/catalyst/decomposition/python_decompositions.py create mode 100644 frontend/catalyst/decomposition/type_stringify_utils.py delete mode 100644 frontend/catalyst/device/python_decompositions.py delete mode 100644 frontend/catalyst/utils/precompile_decomposition_rules.py diff --git a/frontend/catalyst/decomposition/precompile_decomposition_rules.py b/frontend/catalyst/decomposition/precompile_decomposition_rules.py new file mode 100644 index 0000000000..c0729d69be --- /dev/null +++ b/frontend/catalyst/decomposition/precompile_decomposition_rules.py @@ -0,0 +1,195 @@ +# 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. + +"""Utilities for AOT compiling PennyLane's decomposition rules to MLIR Bytecode.""" + +from pathlib import Path + +import pennylane as qp +from jax._src.lib.mlir import ir +from pennylane.operation import Operator, Operator2 + +from catalyst.compiler import _quantum_opt +from catalyst.decomposition.python_decompositions import GraphOpID, python_decomposition +from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH + +# TODO: Uncomment dynamic size wires ops once they are supported +COMPILER_OPS_FOR_DECOMPOSITION = { + qp.CNOT, + qp.ControlledPhaseShift, + qp.CRot, + qp.CRX, + qp.CRY, + qp.CRZ, + qp.CSWAP, + qp.CY, + qp.CZ, + qp.H, + # qp.I, + qp.IsingXX, + qp.IsingXY, + qp.IsingYY, + qp.IsingZZ, + qp.SingleExcitation, + qp.SingleExcitationPlus, + qp.SingleExcitationMinus, + qp.DoubleExcitation, + qp.DoubleExcitationPlus, + qp.DoubleExcitationMinus, + qp.ISWAP, + qp.PauliX, + qp.PauliY, + qp.PauliZ, + # qp.PauliRot, + # qp.PauliMeasure, + qp.PhaseShift, + qp.PSWAP, + qp.Rot, + qp.RX, + qp.RY, + qp.RZ, + qp.S, + qp.SWAP, + qp.T, + # qp.Toffoli, // adjoint not supported + qp.U1, + qp.U2, + qp.U3, + # qp.MultiRZ, + # qp.GlobalPhase, +} + + +def get_rule_funcs_from_module(module: ir.Module) -> list[ir.Operation]: + funcOps = [] + + def find_condition(op): + if op.name == "func.func": + if "target_gate" in op.attributes: + old_attr = op.attributes["sym_name"] + op.attributes["sym_name"] = ir.StringAttr.get( + "__builtin_" + old_attr.value.strip('"'), context=old_attr.context + ) + funcOps.append(op) + return ir.WalkResult.SKIP + return ir.WalkResult.ADVANCE + + module.operation.walk(find_condition) + return funcOps + + +def get_rules_from_module(module: ir.Module) -> str: + """ + Parse and modify decomposition rules from a ModuleOp. + + Args: + module: an MLIR module object containing a FuncOp named `rule_wrapper` to be extracted + + Returns: + str: The string representation of any decomposition rules from `module`, pre-pending the + `__builtin_` prefix to their names. + """ + funcOps = get_rule_funcs_from_module(module) + + return "\n".join(str(funcOp) for funcOp in funcOps) if funcOps else "" + + +def get_abstract_args(op_class: type[Operator]) -> list[type]: + """ + Create jax-compatible abstract args for catalyst DecompositionRules that apply to op_class. + + Args: + op_class: operator to create args for. + + Returns: + list: abstract args for DecompositionRules. + """ + # decomposition rule signatures are of the form + # (*op_params, wires, **op_resource_params, **hyperparams) + # see https://github.com/PennyLaneAI/catalyst/pull/2531#discussion_r2949351413 + if isinstance(op_class.ndim_params, tuple) and any(dim > 0 for dim in op_class.ndim_params): + raise ValueError( + f"Cannot generate arguments for {op_class.__name__} with multi-dimensional parameters." + ) + return [float for _ in range(op_class.num_params)] + + +def parse_operator_data(op): + """Parse operator data from an Operator/Operator2 instance.""" + if isinstance(op, Operator2): + # TODO: use real getters here + dynamic_shape = op.getDynamicShape() + wire_lens = op.getWireLens() + static_data = op.getStaticData() + return dynamic_shape, wire_lens, static_data + if issubclass(op, Operator): + # NOTE: handling this the old-fashioned way, remove once Operator2 migration is complete + dynamic_shape = get_abstract_args(op) + num_wires = op.num_wires if op.num_wires else 0 + return dynamic_shape, [num_wires], {} + else: + raise ValueError( + "Only AbstractOperator and CompressedResourceOp types are supported for generating a " + f"graph ID, got {op} of type {type(op)}" + ) + + +def precompile_decomp_rules(decomp_file_path: str = BYTECODE_FILE_PATH): + """ + Compile PennyLane built-in decomposition rules to MLIR Bytecode. + + Intended for use with `make decomp-rules` in catalyst/mlir. + + Args: + decomp_file_path (Path): path to compile rules to. + """ + Path(decomp_file_path).parent.mkdir(parents=True, exist_ok=True) + + # newline to ensure emptystring is never passed + bytecode_lib = "\n" + + with ir.Context(): + # TODO: update this for Operator2, PL will implement a precompilation registry + for op in COMPILER_OPS_FOR_DECOMPOSITION: + # TODO: this fails because GraphOpID only supports Operator2. Remove try/except once + # Operator2 migration is complete + if not issubclass(op, Operator2): + continue + dynamic_data, wire_lens, static_data = parse_operator_data(op) + if static_data: + # we cannot precompile if the rule takes static data + continue + + mlir_rules = python_decomposition( + op.__name__, GraphOpID(op).getID(), dynamic_data, wire_lens, {} + ) + + bytecode_lib += get_rules_from_module(mlir_rules) + "\n" + + bytecode = _quantum_opt( + "--emit-bytecode", + "--canonicalize", + "--convert-to-value-semantics", + "--canonicalize", + "--register-decomp-rule-resource", + stdin=bytecode_lib.encode("utf-8"), + text=None, + ) + + with open(decomp_file_path, "wb") as bytecode_file: + bytecode_file.write(bytecode) + + +if __name__ == "__main__": # pragma: no cover + precompile_decomp_rules() diff --git a/frontend/catalyst/decomposition/python_decompositions.py b/frontend/catalyst/decomposition/python_decompositions.py new file mode 100644 index 0000000000..55b95dc5e9 --- /dev/null +++ b/frontend/catalyst/decomposition/python_decompositions.py @@ -0,0 +1,213 @@ +# 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. + +""" +This module provides infrastructure for compile-time lowering of decomposition rules via python. +""" + +# pylint: disable=protected-access,bare-except + +import warnings + +import jax.numpy as jnp +import pennylane as qp +from jax._src.lib.mlir import ir +from jaxlib.mlir.dialects.builtin import ModuleOp + +from catalyst.decomposition.type_stringify_utils import mlir_stringify_type +from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval + +_MLIR_DTYPES = { + "i1": jnp.bool_, + "i8": jnp.int8, + "i16": jnp.int16, + "i32": jnp.int32, + "i64": jnp.int64, + "f16": jnp.float16, + "f32": jnp.float32, + "f64": jnp.float64, + "complex": jnp.complex64, + "complex": jnp.complex128, +} + + +def get_dummy_values_for_container(container): + """Given a container of python types, replace the types with corresponding dummy values.""" + dummy_args = [] + for dtype in container: + if isinstance(dtype, str): + if dtype in _MLIR_DTYPES: + count = 1 + dtype = _MLIR_DTYPES[dtype] + elif dtype.startswith("tensor"): + # tensor<{number}x{type}> + dtype = dtype.removeprefix("tensor<") + dtype = dtype.remove_suffice(">") + count, dtype = dtype.split("x") + else: + raise ValueError(f"Unknown dtype {dtype}.") + else: + count = 1 + dtype = jnp.dtype(dtype) + + dummy_args.append(jnp.zeros((count,), dtype=dtype)) + + return tuple(dummy_args) + + +class GraphOpID: + """ + Return the graph operator id for the operator2 instance `op`. + + The FuncOp decomposition rules in the returned string satisfy the following requirements: + - Are named `{rule name}_{op graph ID}`. + - Are MLIR representations of the PennyLane decomposition rules associated with the + specified operator. + - Are instantiated with the static data provided, and all other data remains dynamic. + - Are self-contained, and do not contain any device initialization, setup/teardown etc. + - Are compatible with the `decompose-lowering` and `graph-decomposition` passes, meaning + the following: + - Their `target_gate` attribute is set to the provided graph operator ID + - They have a resources attribute containing an operations attribute which maps graph + operator IDs to counts of their occurrences in the rule. + - Their arguments are mappable to the operator they decompose via `decompose-lowering`. + + Note that this function should not be updated without updating the corresponding method on the + DecomposableGate interface in mlir/lib/quantum/IR/QuantumInterfaces.cpp. + """ + + def __init__(self, op: qp.core.Operator2, uid=None): + assert isinstance( + op, qp.core.Operator2 + ), "Graph-based decomposition expects an Operator2 instance" + self.op = op + + self.operator_name = op.name + self.dynamic_shape = self.parse_dynamic_shape() + self.wire_lens = self.parse_wire_lens() + self.static_data = self.parse_static_data() + self.extra_data = uid + + def parse_dynamic_shape(self): + return list(self.op.dynamic_args.values()) + + def parse_wire_lens(self): + return list(map(len, self.op.wire_args.values())) + + def parse_static_data(self): + return { + static_argname: getattr(self.op, static_argname) + for static_argname in self.op.compilable_argnames + } + + def get_operator_name(self): + return self.operator_name + + def get_dynamic_shape_id_format(self): + return f"[{','.join(map(mlir_stringify_type, self.dynamic_shape))}]" + + def get_wire_lens_id_format(self): + return f"[{','.join(map(str, self.wire_lens))}]" + + def get_static_data_id_format(self): + return f"{{{','.join(f'{k}:{v}' for k, v in self.static_data.items())}}}" + + def getID(self): + ID_string = ( + self.get_operator_name() + + self.get_dynamic_shape_id_format() + + self.get_wire_lens_id_format() + + self.get_static_data_id_format() + ) + if self.extra_data: + ID_string += "[" + str(self.extra_data) + "]" + return ID_string + + +def collect_resources_for_op(op_name, static_data): + decomp_rules = list(qp.decomposition.list_decomps(op_name)) + + # map rules to resource resources, in a more generic format + name_to_resource_ids = {} + name_to_resources = {} + for rule in decomp_rules: + # The `compute_resources` function's signature is the same as the Operator2 signature + # for the original op of the rule + resources = rule.compute_resources(**static_data) + name_to_resources[rule.name] = resources.gate_counts + name_to_resource_ids[rule.name] = { + GraphOpID(op).getID(): count for op, count in resources.gate_counts.items() + } + + return name_to_resources, name_to_resource_ids, decomp_rules + + +def python_decomposition(op_name, op_id, dynamic_shape, wire_lens, static_data) -> ModuleOp: + """Python decomposition rule lowering.""" + # TODO update docstring + device = qp.device("null.qubit", wires=sum(wire_lens)) + wires = tuple(jnp.array(range(length), dtype=int) for length in wire_lens) + + _, name_to_resource_ids, decomp_rules = collect_resources_for_op(op_name, static_data) + + def rule_to_subroutine(rule): + def decomp_rule(*args, **kwargs): + rule._impl(*args, **kwargs) + + # keep the frontend name for readability, append target op_id for symbol uniqueness + decomp_rule.__name__ = rule._impl.__name__ + "_" + op_id + + return qp.capture.subroutine(decomp_rule) + + subroutines = [rule_to_subroutine(rule) for rule in decomp_rules] + + @qp.qjit( + target="mlir", + capture=True, + ) + @qp.qnode(device=device) + def circuit(): + for subroutine in subroutines: + subroutine(*get_dummy_values_for_container(dynamic_shape), wires=wires) + + module = circuit.mlir_module + + def update_funcop_attributes(op): + """Update the decomposition rule attributes if op is a decomposition rule. + + For use with module.walk + + This function updates the following attributes: + - Adds the `target_gate` attribute. + - Adds the `resources` attribute. + """ + if op.name == "func.func": + rule_name = ir.StringAttr(op.attributes["sym_name"]).value.removesuffix("_" + op_id) + if rule_name in name_to_resource_ids: + op.attributes["resources"] = get_mlir_attribute_from_pyval( + {"operations": name_to_resource_ids[rule_name]} + ) + op.attributes["target_gate"] = ir.StringAttr.get(op_id) + + return ir.WalkResult.ADVANCE + + with module.context: + module.operation.walk(update_funcop_attributes) + + return module + + +def python_decomposition_wrapper(op_name, op_id, dynamic_shape, wire_lens, static_data) -> str: + """Generic decomposition wrapper.""" + return str(python_decomposition(op_name, op_id, dynamic_shape, wire_lens, static_data)) diff --git a/frontend/catalyst/decomposition/type_stringify_utils.py b/frontend/catalyst/decomposition/type_stringify_utils.py new file mode 100644 index 0000000000..7cd0ab570f --- /dev/null +++ b/frontend/catalyst/decomposition/type_stringify_utils.py @@ -0,0 +1,44 @@ +# 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. + +import jax.numpy as jnp +import pennylane as qp + +from catalyst.utils.exceptions import CompileError + + +def _py_dtype_to_mlir_type_string(python_dtype: type): + match python_dtype: + case jnp.float64: + return "f64" + case _: + raise CompileError("Unknown data type") + + +def _stringify_shaped_type(shape: tuple, dim: int, element_type): + if dim + 1 == len(shape): + inner_content = _py_dtype_to_mlir_type_string(element_type) + else: + inner_content = _stringify_shaped_type(shape, dim + 1, element_type) + length = shape[dim] + return f"[{','.join([inner_content] * length)}]" + + +def mlir_stringify_type(dtype: qp.typing.AbstractArray): + assert isinstance(dtype, qp.typing.AbstractArray) + element_type = dtype.dtype.type + if dtype.shape == (): + return _py_dtype_to_mlir_type_string(element_type) + else: + return _stringify_shaped_type(dtype.shape, 0, element_type) diff --git a/frontend/catalyst/device/python_decompositions.py b/frontend/catalyst/device/python_decompositions.py deleted file mode 100644 index b6ba5d71fe..0000000000 --- a/frontend/catalyst/device/python_decompositions.py +++ /dev/null @@ -1,90 +0,0 @@ -# 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. - -""" -This module provides infrastructure for compile-time lowering of decomposition rules via python. - -Python decomposition wrappers should adhere to the following specifications: - The wrapper: - - Is named `{op name}_decomposition_wrapper`. - - Has a signature identical to the named parameters of the associated PL operator; dynamic - arguments may be unused, but should still be included for compatibility. - - Is able to AOT lower the decomposition rule to MLIR without invoking the compiler, e.g. - using `target="mlir"`, AOT compilation and `QJIT.mlir_module`. - See existing examples for further information. - - Returns a string representation of an MLIR module, containing a FuncOp which represents - the instantiated decomposition rule. - - The FuncOp decomposition rule in the returned string: - - Is named `{op name}_decomp_rule`. - - Is an MLIR representation of the PennyLane decomposition rule associated with the - specified operator. - - Is instantiated with the static data provided, and all other data remains dynamic. - - Is compatible with the `decompose-lowering` pass, i.e. can be mapped to the MLIR operation - it decomposes and inlined. - - Is self-contained, and does not contain any device initialization, setup/teardown etc. -""" - -# pylint: disable=protected-access,bare-except - -import warnings - -import jax.numpy as jnp -import pennylane as qp - - -def python_decomposition_wrapper(op_name, op_id, dynamic_shape, wire_lens, static_data) -> str: - """Generic decomposition wrapper.""" - device = qp.device("null.qubit", wires=sum(wire_lens)) - wires = tuple(jnp.array(range(length), dtype=int) for length in wire_lens) - - def rule_to_subroutine(rule): - def decomp_rule(*params, wires): - rule._impl(*params, *wires, **static_data) - - # TODO remove this once we have unified lowering, we should be able to set target_gate and - # stop relying on function names - decomp_rule.__name__ = op_id + "_" + rule.name - - return qp.capture.subroutine(decomp_rule) - - # let this fail with the standard error message if the op is not found - subroutines = [rule_to_subroutine(rule) for rule in qp.decomposition.list_decomps(op_name)] - - # TODO: not all PL ops have been migrated to the operator 2 format expected by mlir graph decomp - # This means some rules will fail the python callback compilation. - # When migration is complete, remove the try-except. - try: - - @qp.qjit( - target="mlir", - capture=True, - ) - @qp.qnode(device=device) - def circuit(): - for subroutine in subroutines: - # TODO: I know this is dynamic, but we should probably have a better way of handling - # this than hard-coded dummy values. Revisit this when unifying the decomp-rule - # lowering pipeline - subroutine(*[0.5 for _ in dynamic_shape], wires=wires) - - return str(circuit.mlir_module) - except: - warnings.warn( - f"Python decomposition rule compilation failed for operator " - f"'{op_name}' (id: {op_id}); it will be treated as non-decomposable " - f"by the graph solver.", - UserWarning, - ) - return "builtin.module{}" diff --git a/frontend/catalyst/utils/precompile_decomposition_rules.py b/frontend/catalyst/utils/precompile_decomposition_rules.py deleted file mode 100644 index 2ca103944a..0000000000 --- a/frontend/catalyst/utils/precompile_decomposition_rules.py +++ /dev/null @@ -1,240 +0,0 @@ -# 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. - -"""Utilities for AOT compiling PennyLane's decomposition rules to MLIR Bytecode.""" - -import warnings -from pathlib import Path - -import jax -import pennylane as qp -from jax._src.lib.mlir import ir -from pennylane.operation import Operator - -from catalyst.compiler import _quantum_opt -from catalyst.jax_primitives import decomposition_rule -from catalyst.utils.exceptions import CompileError -from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH - -# TODO: Uncomment dynamic size wires ops once they are supported -# FIXME: Use the Gate class instead of this list of compiler ops -# https://github.com/PennyLaneAI/pennylane/pull/8767 -COMPILER_OPS_FOR_DECOMPOSITION = { - qp.CNOT, - qp.ControlledPhaseShift, - qp.CRot, - qp.CRX, - qp.CRY, - qp.CRZ, - qp.CSWAP, - qp.CY, - qp.CZ, - qp.H, - # qp.I, - qp.IsingXX, - qp.IsingXY, - qp.IsingYY, - qp.IsingZZ, - qp.SingleExcitation, - qp.SingleExcitationPlus, - qp.SingleExcitationMinus, - qp.DoubleExcitation, - qp.DoubleExcitationPlus, - qp.DoubleExcitationMinus, - qp.ISWAP, - qp.PauliX, - qp.PauliY, - qp.PauliZ, - # qp.PauliRot, - # qp.PauliMeasure, - qp.PhaseShift, - qp.PSWAP, - qp.Rot, - qp.RX, - qp.RY, - qp.RZ, - qp.S, - qp.SWAP, - qp.T, - # qp.Toffoli, // adjoint not supported - qp.U1, - qp.U2, - qp.U3, - # qp.MultiRZ, - # qp.GlobalPhase, -} - - -def get_abstract_args(op_class: type[Operator]) -> list[type]: - """ - Create jax-compatible abstract args for catalyst DecompositionRules that apply to op_class. - - Args: - op_class: operator to create args for. - - Returns: - list: abstract args for DecompositionRules. - """ - # decomposition rule signatures are of the form - # (*op_params, wires, **op_resource_params, **hyperparams) - # see https://github.com/PennyLaneAI/catalyst/pull/2531#discussion_r2949351413 - if isinstance(op_class.ndim_params, tuple) and any(dim > 0 for dim in op_class.ndim_params): - raise ValueError( - f"Cannot generate arguments for {op_class.__name__} with multi-dimensional parameters." - ) - return [float for _ in range(op_class.num_params)] - - -def get_func_from_circuit(module) -> str | None: - """ - Get the string representation of `rule_wrapper` from module, if it exists. - - Args: - module: an MLIR module object containing a FuncOp named `rule_wrapper` to be extracted - - Returns: - str: string representation of FuncOp named `rule_wrapper` from module - None: if no such FuncOp can be found - """ - decomp_func_op = None - - def find_condition(op): - nonlocal decomp_func_op - if op.name == "func.func": - if ir.StringAttr(op.attributes["sym_name"]).value == "rule_wrapper": - decomp_func_op = op - return ir.WalkResult.INTERRUPT - return ir.WalkResult.ADVANCE - - module.operation.walk(find_condition) - - return str(decomp_func_op) + "\n" if decomp_func_op else None - - -def compile_rule( - op_class, - abstract_args, - op_num_wires, - rule, - dev, -) -> str | None: - """ - Get the string representation of a compiled rule from a python decomposition rule, if possible. - - NOTE: rules with string params are not currently supported. - - Args: - op_class: A PennyLane class subclassing Operation - op_num_wires: the number of wires used by op_class - rule (DecompositionRule): the decomposition rule to be compiled - dev (Device): a device for qjit - - Returns: - str: string representation of the mlir of the decomposition rule. - """ - qp.decomposition.enable_graph() - - # WARNING: do not rename this function, we use it to extract the rule from the compiled - # circuit - @decomposition_rule(is_qreg=True, op_type=op_class.__name__) - def rule_wrapper(*args, wires, **_): - return rule(*args, wires=wires, **_) - - @qp.qjit(capture=True, target="mlir") - @qp.qnode(dev) - def circuit(): - rule_wrapper(*abstract_args, wires=jax.core.ShapedArray((op_num_wires,), int)) - return qp.probs() - - return get_func_from_circuit(circuit.mlir_module) - - -def compile_op_decomp_rules( - op_class: type[Operator], -) -> dict[str, str | None]: - """ - Compile all decomposition rules for op_class. - - Note: the modules include the full circuit IR. - - Args: - op_class (type[Operator]): the op class to compile decomposition rules for. - - Returns: - dict[str, str | None]: decomposition rule names to compiled mlir modules. - """ - op_decomp_rules = qp.decomposition.decomposition_graph.list_decomps(op_class) - - mlir_modules: dict[str, str | None] = {} - - if not hasattr(op_class, "num_wires") or not op_class.num_wires: - warnings.warn( - f"Cannot compile decomposition rules for op {op_class.__name__} with an unknown number " - + "of wires." - ) - return mlir_modules - - dev = qp.device("null.qubit", wires=op_class.num_wires) - - abstract_args = get_abstract_args(op_class) # pylint: disable=protected-access - - for rule in op_decomp_rules: - try: - rule_name = rule._impl.__name__ # pylint: disable=protected-access - mlir_modules[rule_name] = compile_rule( - op_class, abstract_args, op_class.num_wires, rule, dev - ) - except CompileError as e: - warnings.warn(f"Failed to compile {rule_name}: {e}") - except Exception as e: # pylint: disable=broad-exception-caught - warnings.warn(f"Unexpected error while trying to compile {rule_name}: {e}") - finally: - qp.decomposition.disable_graph() - - return mlir_modules - - -def precompile_decomp_rules(decomp_file_path: str = BYTECODE_FILE_PATH): - """ - Compile PennyLane built-in decomposition rules to MLIR Bytecode. - - Intended for use with `make decomp-rules` in catalyst/mlir. - - Args: - decomp_file_path (Path): path to compile rules to. - """ - Path(decomp_file_path).parent.mkdir(parents=True, exist_ok=True) - - mlir_rules = "".join( - str(mlir).replace("@rule_wrapper", f"@__builtin_{name}") - for func in COMPILER_OPS_FOR_DECOMPOSITION - for name, mlir in compile_op_decomp_rules(func).items() - ) - - bytecode = _quantum_opt( - "--emit-bytecode", - "--canonicalize", - "--convert-to-value-semantics", - "--canonicalize", - "--register-decomp-rule-resource", - stdin=mlir_rules.encode("utf-8"), - text=None, - ) - - with open(decomp_file_path, "wb") as bytecode_file: - bytecode_file.write(bytecode) - - -if __name__ == "__main__": # pragma: no cover - precompile_decomp_rules() From fc046700919b70901c013b082cde90a61ebf01d0 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 23 Jul 2026 10:26:01 -0400 Subject: [PATCH 08/36] update makefile for precompiled rules --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 7f7062b5b2..883bb82c04 100644 --- a/Makefile +++ b/Makefile @@ -122,7 +122,7 @@ frontend: # versions of a package with the same version tag (e.g. 0.38-dev0). $(PYTHON) -m pip uninstall -y pennylane $(PYTHON) -m pip install -e . --extra-index-url https://test.pypi.org/simple $(PIP_VERBOSE_FLAG) - $(PYTHON) -m catalyst.utils.precompile_decomposition_rules + $(PYTHON) -m catalyst.decomposition.precompile_decomposition_rules rm -r frontend/pennylane_catalyst.egg-info .PHONY: mlir llvm stablehlo enzyme dialects runtime oqc builtin-decomp-rules From b1c631e8b67a4219e21301b6adbd41796b892481 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 23 Jul 2026 10:34:01 -0400 Subject: [PATCH 09/36] Update frontend/catalyst/decomposition/precompile_decomposition_rules.py Co-authored-by: Paul <79805239+paul0403@users.noreply.github.com> --- .../catalyst/decomposition/precompile_decomposition_rules.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/decomposition/precompile_decomposition_rules.py b/frontend/catalyst/decomposition/precompile_decomposition_rules.py index c0729d69be..ffd7f6c9d7 100644 --- a/frontend/catalyst/decomposition/precompile_decomposition_rules.py +++ b/frontend/catalyst/decomposition/precompile_decomposition_rules.py @@ -162,7 +162,6 @@ def precompile_decomp_rules(decomp_file_path: str = BYTECODE_FILE_PATH): with ir.Context(): # TODO: update this for Operator2, PL will implement a precompilation registry for op in COMPILER_OPS_FOR_DECOMPOSITION: - # TODO: this fails because GraphOpID only supports Operator2. Remove try/except once # Operator2 migration is complete if not issubclass(op, Operator2): continue From 63fc4fe5586bf6282c9a63b6842c0c9d04a61ae6 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 23 Jul 2026 10:34:06 -0400 Subject: [PATCH 10/36] Update frontend/catalyst/decomposition/precompile_decomposition_rules.py Co-authored-by: Paul <79805239+paul0403@users.noreply.github.com> --- .../catalyst/decomposition/precompile_decomposition_rules.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/decomposition/precompile_decomposition_rules.py b/frontend/catalyst/decomposition/precompile_decomposition_rules.py index ffd7f6c9d7..659c7db03b 100644 --- a/frontend/catalyst/decomposition/precompile_decomposition_rules.py +++ b/frontend/catalyst/decomposition/precompile_decomposition_rules.py @@ -162,7 +162,6 @@ def precompile_decomp_rules(decomp_file_path: str = BYTECODE_FILE_PATH): with ir.Context(): # TODO: update this for Operator2, PL will implement a precompilation registry for op in COMPILER_OPS_FOR_DECOMPOSITION: - # Operator2 migration is complete if not issubclass(op, Operator2): continue dynamic_data, wire_lens, static_data = parse_operator_data(op) From be85e2051623409811dd4bf09d6d6e41fc4e688a Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 23 Jul 2026 11:49:48 -0400 Subject: [PATCH 11/36] update module in QPD path --- .../Transforms/QuantumPythonDecompositions/PythonFunction.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Quantum/Transforms/QuantumPythonDecompositions/PythonFunction.cpp b/mlir/lib/Quantum/Transforms/QuantumPythonDecompositions/PythonFunction.cpp index d6e5c40b33..222ab8ef1f 100644 --- a/mlir/lib/Quantum/Transforms/QuantumPythonDecompositions/PythonFunction.cpp +++ b/mlir/lib/Quantum/Transforms/QuantumPythonDecompositions/PythonFunction.cpp @@ -76,7 +76,7 @@ std::string pythonRuleLowering(catalyst::quantum::DecomposableGate op) { QuantumPythonDecompositions::PyInterpreterGuard guard; std::string mlirText = guard.withGil([&] -> std::string { - const char *moduleName = "catalyst.device.python_decompositions"; + const char *moduleName = "catalyst.decomposition.python_decompositions"; const char *functionName = "python_decomposition_wrapper"; try { From 2abc8408f4bc6244f9b0099cc63927bdfd60c1bd Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 23 Jul 2026 12:20:24 -0400 Subject: [PATCH 12/36] remove PL dependency from lit tests --- .../GraphDecomposition/TestAltDecomps.mlir | 28 ++++++++--------- .../test/lit/GraphDecomposition/TestCost.mlir | 11 ++++--- .../GraphDecomposition/TestFailedDecomp.mlir | 6 ++-- .../GraphDecomposition/TestFixedDecomp.mlir | 18 +++++------ .../GraphDecomposition/TestGatesetRXRZ.mlir | 10 +++---- .../GraphDecomposition/TestGatesetRYRZ.mlir | 10 +++---- .../lit/GraphDecomposition/TestGraphOpId.mlir | 16 +++++----- .../lit/GraphDecomposition/TestIdentity.mlir | 30 +++++++++---------- .../GraphDecomposition/TestMultiDecomp.mlir | 22 +++++++------- .../TestMultiDecompUser.mlir | 17 ++++++----- .../lit/GraphDecomposition/test_rules.mlir | 29 ++++++++---------- 11 files changed, 95 insertions(+), 102 deletions(-) diff --git a/frontend/test/lit/GraphDecomposition/TestAltDecomps.mlir b/frontend/test/lit/GraphDecomposition/TestAltDecomps.mlir index 439fb41432..7a1d2462ac 100644 --- a/frontend/test/lit/GraphDecomposition/TestAltDecomps.mlir +++ b/frontend/test/lit/GraphDecomposition/TestAltDecomps.mlir @@ -12,21 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. -// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=RY=1.0,PauliX=3.0,PauliZ=3.0,GlobalPhase=1.0 alt-decomps=PauliY=[y_to_ry,y_to_x_z] bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes RY +// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRY=1.0,testX=3.0,testZ=3.0 alt-decomps=PauliY=[y_to_ry,y_to_x_z] bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes RY -// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=RY=3.0,PauliX=1.0,PauliZ=1.0,GlobalPhase=1.0 alt-decomps=PauliY=[y_to_ry,y_to_x_z] bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes XZ +// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRY=3.0,testX=1.0,testZ=1.0 alt-decomps=PauliY=[y_to_ry,y_to_x_z] bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes XZ func.func @circuit() -> !quantum.bit { %0 = quantum.alloc(2) : !quantum.reg %q = quantum.extract %0[0] : !quantum.reg -> !quantum.bit - // RY-NOT: PauliY - // RY: RY - // RY: gphase + // RY-NOT: test + // RY: testRY - // XZ-NOT: PauliY - // XZ: PauliX - // XZ: PauliZ - %qout = quantum.custom "PauliY"() %q : !quantum.bit + // XZ-NOT: test + // XZ: testX + // XZ: testZ + %qout = quantum.custom "testY"() %q : !quantum.bit // needed to ensure we don't match in the following decomposition rules // CHECK: return @@ -34,17 +33,16 @@ func.func @circuit() -> !quantum.bit { } // CHECK-LABEL: y_to_ry -func.func @y_to_ry(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="PauliY[][1]{}", resources = { operations = {"RY[f64][1]{}"=1, "GlobalPhase[][]{}"=1}}} { +func.func @y_to_ry(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="testY[][1]{}", resources = { operations = {"testRY[f64][1]{}"=1}}} { %pi = arith.constant 3.14 : f64 %negpiby2 = arith.constant -1.57 : f64 - %q1 = quantum.custom "RY"(%pi) %q0 : !quantum.bit - quantum.gphase(%negpiby2) + %q1 = quantum.custom "testRY"(%pi) %q0 : !quantum.bit return %q1 : !quantum.bit } // CHECK-LABEL: y_to_x_z -func.func @y_to_x_z(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="PauliY[][1]{}", resources = { operations = {"PauliX[][1]{}"=1, "PauliZ[][1]{}"=1}}} { - %q1 = quantum.custom "PauliX"() %q0 : !quantum.bit - %q2 = quantum.custom "PauliZ"() %q1 : !quantum.bit +func.func @y_to_x_z(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="testY[][1]{}", resources = { operations = {"testX[][1]{}"=1, "testZ[][1]{}"=1}}} { + %q1 = quantum.custom "testX"() %q0 : !quantum.bit + %q2 = quantum.custom "testZ"() %q1 : !quantum.bit return %q2 : !quantum.bit } diff --git a/frontend/test/lit/GraphDecomposition/TestCost.mlir b/frontend/test/lit/GraphDecomposition/TestCost.mlir index e30f567078..0b52748edc 100644 --- a/frontend/test/lit/GraphDecomposition/TestCost.mlir +++ b/frontend/test/lit/GraphDecomposition/TestCost.mlir @@ -14,16 +14,15 @@ // Test that decomposition chooses cheapest decomposition path -// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=RX=1.0,RY=1.0,RZ=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s +// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRX=1.0,testRY=1.0,testRZ=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s func.func @circuit() -> !quantum.bit { %0 = quantum.alloc(1) : !quantum.reg %q = quantum.extract %0[0] : !quantum.reg -> !quantum.bit - // CHECK-NOT: Hadamard - // CHECK-DAG: RZ - // CHECK-DAG: RY - // CHECK-DAG: gphase - %qout = quantum.custom "Hadamard"() %q : !quantum.bit + // CHECK-NOT: testHadamard + // CHECK-DAG: testRZ + // CHECK-DAG: testRY + %qout = quantum.custom "testHadamard"() %q : !quantum.bit return %qout : !quantum.bit } diff --git a/frontend/test/lit/GraphDecomposition/TestFailedDecomp.mlir b/frontend/test/lit/GraphDecomposition/TestFailedDecomp.mlir index 89777763da..e782aaff74 100644 --- a/frontend/test/lit/GraphDecomposition/TestFailedDecomp.mlir +++ b/frontend/test/lit/GraphDecomposition/TestFailedDecomp.mlir @@ -16,11 +16,9 @@ func.func @circuit(%q0: !quantum.bit) { %pi = arith.constant 3.14 : f64 - %out = quantum.pcphase (%pi, dim : 3) %q0 : !quantum.bit + %out = quantum.custom "failure"() %q0 : !quantum.bit - // CHECK: UserWarning: Python decomposition rule compilation failed for operator 'PCPhase' (id: PCPhase[f64][1]{dim:3}) - // CHECK-SAME: it will be treated as non-decomposable by the graph solver // CHECK: GraphSolverFailedError - // CHECK: Decomposition rule not found for operator 'pcphase + // CHECK: Decomposition rule not found for operator 'id: failure[][1]{}' return } diff --git a/frontend/test/lit/GraphDecomposition/TestFixedDecomp.mlir b/frontend/test/lit/GraphDecomposition/TestFixedDecomp.mlir index 7547150347..264170c2eb 100644 --- a/frontend/test/lit/GraphDecomposition/TestFixedDecomp.mlir +++ b/frontend/test/lit/GraphDecomposition/TestFixedDecomp.mlir @@ -12,23 +12,23 @@ // See the License for the specific language governing permissions and // limitations under the License. -// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=RX=2.0,RY=1.0,RZ=1.0,GlobalPhase=0.0 fixed-decomps=Hadamard=custom_decomp bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s +// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRX=2.0,testRY=1.0,testRZ=1.0 fixed-decomps=testHadamard=custom_decomp bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s func.func @circuit() -> !quantum.bit { %0 = quantum.alloc(1) : !quantum.reg %q = quantum.extract %0[0] : !quantum.reg -> !quantum.bit - // CHECK-NOT: Hadamard" - // CHECK: RX - // CHECK: RZ - // CHECK: RX - %qout = quantum.custom "Hadamard"() %q : !quantum.bit + // CHECK-NOT: testHadamard" + // CHECK: testRX + // CHECK: testRZ + // CHECK: testRX + %qout = quantum.custom "testHadamard"() %q : !quantum.bit return %qout : !quantum.bit } func.func @custom_decomp(%q0 : !quantum.bit) -> !quantum.bit { %cst = arith.constant 1.5707963267948966 : f64 - %q1 = quantum.custom "RX"(%cst) %q0 : !quantum.bit - %q2 = quantum.custom "RZ"(%cst) %q1 : !quantum.bit - %q3 = quantum.custom "RX"(%cst) %q2 : !quantum.bit + %q1 = quantum.custom "testRX"(%cst) %q0 : !quantum.bit + %q2 = quantum.custom "testRZ"(%cst) %q1 : !quantum.bit + %q3 = quantum.custom "testRX"(%cst) %q2 : !quantum.bit return %q3 : !quantum.bit } diff --git a/frontend/test/lit/GraphDecomposition/TestGatesetRXRZ.mlir b/frontend/test/lit/GraphDecomposition/TestGatesetRXRZ.mlir index 5b24777ae1..84fc20d2d2 100644 --- a/frontend/test/lit/GraphDecomposition/TestGatesetRXRZ.mlir +++ b/frontend/test/lit/GraphDecomposition/TestGatesetRXRZ.mlir @@ -12,16 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. -// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=RX=1.0,RZ=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s +// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRX=1.0,testRZ=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s func.func public @circuit() attributes {quantum.node} { %0 = quantum.alloc( 1) : !quantum.reg %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit - %out_qubits = quantum.custom "Hadamard"() %1 : !quantum.bit + %out_qubits = quantum.custom "testHadamard"() %1 : !quantum.bit %2 = quantum.insert %0[ 0], %out_qubits : !quantum.reg, !quantum.bit quantum.dealloc %2 : !quantum.reg - // CHECK-NOT: Hadamard - // CHECK-DAG: RZ - // CHECK-DAG: RX + // CHECK-NOT: testHadamard + // CHECK-DAG: testRZ + // CHECK-DAG: testRX return } diff --git a/frontend/test/lit/GraphDecomposition/TestGatesetRYRZ.mlir b/frontend/test/lit/GraphDecomposition/TestGatesetRYRZ.mlir index f5457d14d0..ea949e1dcb 100644 --- a/frontend/test/lit/GraphDecomposition/TestGatesetRYRZ.mlir +++ b/frontend/test/lit/GraphDecomposition/TestGatesetRYRZ.mlir @@ -12,16 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. -// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=RY=1.0,RZ=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s +// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRY=1.0,testRZ=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s func.func public @circuit() attributes {quantum.node} { %0 = quantum.alloc( 1) : !quantum.reg %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit - %out_qubits = quantum.custom "Hadamard"() %1 : !quantum.bit + %out_qubits = quantum.custom "testHadamard"() %1 : !quantum.bit %2 = quantum.insert %0[ 0], %out_qubits : !quantum.reg, !quantum.bit quantum.dealloc %2 : !quantum.reg - // CHECK-NOT: Hadamard - // CHECK-DAG: RZ - // CHECK-DAG: RY + // CHECK-NOT: testHadamard + // CHECK-DAG: testRZ + // CHECK-DAG: testRY return } diff --git a/frontend/test/lit/GraphDecomposition/TestGraphOpId.mlir b/frontend/test/lit/GraphDecomposition/TestGraphOpId.mlir index 8536a5fa85..3a9d251c50 100644 --- a/frontend/test/lit/GraphDecomposition/TestGraphOpId.mlir +++ b/frontend/test/lit/GraphDecomposition/TestGraphOpId.mlir @@ -14,18 +14,18 @@ // Test that graph-decomposition succeeds when using graphOpIds -// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=PauliX=1.0 alt-decomps=Hadamard=my_decomp})' %s | FileCheck %s +// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=testPauliX=1.0 alt-decomps=testHadamard=my_decomp})' %s | FileCheck %s func.func @circuit(%q: !quantum.bit) -> !quantum.bit { - // CHECK-NOT: Hadamard - // CHECK: PauliX - // CHECK: PauliX - %out = quantum.custom "Hadamard"() %q: !quantum.bit + // CHECK-NOT: testHadamard + // CHECK: testPauliX + // CHECK: testPauliX + %out = quantum.custom "testHadamard"() %q: !quantum.bit return %out: !quantum.bit } -func.func private @my_decomp(%q: !quantum.bit) -> !quantum.bit attributes {target_gate="Hadamard[][1]{}"} { - %q0 = quantum.custom "PauliX"() %q : !quantum.bit - %q1 = quantum.custom "PauliX"() %q0 : !quantum.bit +func.func private @my_decomp(%q: !quantum.bit) -> !quantum.bit attributes {target_gate="testHadamard[][1]{}"} { + %q0 = quantum.custom "testPauliX"() %q : !quantum.bit + %q1 = quantum.custom "testPauliX"() %q0 : !quantum.bit return %q1 : !quantum.bit } diff --git a/frontend/test/lit/GraphDecomposition/TestIdentity.mlir b/frontend/test/lit/GraphDecomposition/TestIdentity.mlir index b27672dbef..3afe7855b8 100644 --- a/frontend/test/lit/GraphDecomposition/TestIdentity.mlir +++ b/frontend/test/lit/GraphDecomposition/TestIdentity.mlir @@ -14,18 +14,18 @@ // Test that decomposition handles circuits that are already decomposed to the given gateset. -// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=Hadamard=1.0,CNOT=1.0 alt-decomps=Hadamard=false_decomp bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s +// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=testHadamard=1.0,testCNOT=1.0 alt-decomps=testHadamard=false_decomp bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s func.func @circuit() -> !quantum.bit { %0 = quantum.alloc(2) : !quantum.reg %q0 = quantum.extract %0[0] : !quantum.reg -> !quantum.bit %q1 = quantum.extract %0[1] : !quantum.reg -> !quantum.bit - // CHECK: Hadamard - // CHECK: Hadamard - // CHECK: CNOT - %q0out = quantum.custom "Hadamard"() %q0 : !quantum.bit - %q1out = quantum.custom "Hadamard"() %q1 : !quantum.bit - %q:2 = quantum.custom "CNOT"() %q0out, %q1out : !quantum.bit, !quantum.bit + // CHECK: testHadamard + // CHECK: testHadamard + // CHECK: testCNOT + %q0out = quantum.custom "testHadamard"() %q0 : !quantum.bit + %q1out = quantum.custom "testHadamard"() %q1 : !quantum.bit + %q:2 = quantum.custom "testCNOT"() %q0out, %q1out : !quantum.bit, !quantum.bit return %q1 : !quantum.bit } @@ -36,17 +36,17 @@ module @test_module { %0 = quantum.alloc(2) : !quantum.reg %q0 = quantum.extract %0[0] : !quantum.reg -> !quantum.bit %q1 = quantum.extract %0[1] : !quantum.reg -> !quantum.bit - // CHECK: Hadamard - // CHECK: Hadamard - // CHECK: CNOT - %q0out = quantum.custom "Hadamard"() %q0 : !quantum.bit - %q1out = quantum.custom "Hadamard"() %q1 : !quantum.bit - %q:2 = quantum.custom "CNOT"() %q0out, %q1out : !quantum.bit, !quantum.bit + // CHECK: testHadamard + // CHECK: testHadamard + // CHECK: testCNOT + %q0out = quantum.custom "testHadamard"() %q0 : !quantum.bit + %q1out = quantum.custom "testHadamard"() %q1 : !quantum.bit + %q:2 = quantum.custom "testCNOT"() %q0out, %q1out : !quantum.bit, !quantum.bit return %q1 : !quantum.bit } - func.func private @false_decomp(%q : !quantum.bit) -> !quantum.bit attributes {target_gate="Hadamard"} { - %qout = quantum.custom "PauliX"() %q : !quantum.bit + func.func private @false_decomp(%q : !quantum.bit) -> !quantum.bit attributes {target_gate="testHadamard[][1]{}"} { + %qout = quantum.custom "testPauliX"() %q : !quantum.bit return %qout : !quantum.bit } } diff --git a/frontend/test/lit/GraphDecomposition/TestMultiDecomp.mlir b/frontend/test/lit/GraphDecomposition/TestMultiDecomp.mlir index 35256498cb..a6f04685de 100644 --- a/frontend/test/lit/GraphDecomposition/TestMultiDecomp.mlir +++ b/frontend/test/lit/GraphDecomposition/TestMultiDecomp.mlir @@ -12,22 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. -// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=RX=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes FIRST +// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRX=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes FIRST -// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=RX=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"},graph-decomposition{gate-set=RZ=1.0,RY=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes SECOND +// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRX=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"},graph-decomposition{gate-set=testRZ=1.0,testRY=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes SECOND func.func @circuit() -> !quantum.bit { %0 = quantum.alloc(2) : !quantum.reg %q = quantum.extract %0[0] : !quantum.reg -> !quantum.bit - // FIRST-NOT: PauliX - // FIRST: RX - - // SECOND-NOT: PauliX - // SECOND-NOT: RX - // SECOND: RZ - // SECOND: RY - // SECOND: RZ - %qout = quantum.custom "PauliX"() %q : !quantum.bit + // FIRST-NOT: testPauliX + // FIRST: testRX + + // SECOND-NOT: testPauliX + // SECOND-NOT: testRX + // SECOND: testRZ + // SECOND: testRY + // SECOND: testRZ + %qout = quantum.custom "testPauliX"() %q : !quantum.bit return %qout : !quantum.bit } diff --git a/frontend/test/lit/GraphDecomposition/TestMultiDecompUser.mlir b/frontend/test/lit/GraphDecomposition/TestMultiDecompUser.mlir index c7a206cb33..5406364d5a 100644 --- a/frontend/test/lit/GraphDecomposition/TestMultiDecompUser.mlir +++ b/frontend/test/lit/GraphDecomposition/TestMultiDecompUser.mlir @@ -12,23 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. -// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module( graph-decomposition{gate-set=Hadamard=1.0 fixed-decomps=PauliX=x_to_h bytecode-rules="%BYTECODE_PATH"}, graph-decomposition{gate-set=PauliX=1.0 fixed-decomps=Hadamard=h_to_x bytecode-rules="%BYTECODE_PATH"}, graph-decomposition{gate-set=Hadamard=1.0 fixed-decomps=PauliX=x_to_h bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes TRIPLE +// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module( graph-decomposition{gate-set=testHadamard=1.0 fixed-decomps=testPauliX=x_to_h bytecode-rules="%BYTECODE_PATH"}, graph-decomposition{gate-set=testPauliX=1.0 fixed-decomps=testHadamard=h_to_x bytecode-rules="%BYTECODE_PATH"}, graph-decomposition{gate-set=testHadamard=1.0 fixed-decomps=testPauliX=x_to_h bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s func.func @circuit() -> !quantum.bit { %0 = quantum.alloc(2) : !quantum.reg %q = quantum.extract %0[0] : !quantum.reg -> !quantum.bit - // TRIPLE-NOT PauliX - // TRIPLE: Hadamard - %qout = quantum.custom "PauliX"() %q : !quantum.bit + // CHECK-NOT testPauliX + // CHECK: testHadamard + %qout = quantum.custom "testPauliX"() %q : !quantum.bit return %qout : !quantum.bit } -func.func @h_to_x(%q : !quantum.bit) -> !quantum.bit attributes {target_gate="Hadamard"} { - %q1 = quantum.custom "PauliX"() %q : !quantum.bit +// CHECK-LABEL: h_to_x +func.func @h_to_x(%q : !quantum.bit) -> !quantum.bit attributes {target_gate="testHadamard[][1]{}"} { + %q1 = quantum.custom "testPauliX"() %q : !quantum.bit return %q1 : !quantum.bit } -func.func @x_to_h(%q : !quantum.bit) -> !quantum.bit attributes {target_gate="PauliX"} { - %q1 = quantum.custom "Hadamard"() %q : !quantum.bit +func.func @x_to_h(%q : !quantum.bit) -> !quantum.bit attributes {target_gate="testPauliX[][1]{}"} { + %q1 = quantum.custom "testHadamard"() %q : !quantum.bit return %q1 : !quantum.bit } diff --git a/frontend/test/lit/GraphDecomposition/test_rules.mlir b/frontend/test/lit/GraphDecomposition/test_rules.mlir index f0040a9ec6..f98dad89ea 100644 --- a/frontend/test/lit/GraphDecomposition/test_rules.mlir +++ b/frontend/test/lit/GraphDecomposition/test_rules.mlir @@ -13,39 +13,36 @@ // limitations under the License. -func.func @__builtin_h_to_rz_ry(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="Hadamard"} { +func.func @__builtin_h_to_rz_ry(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="testHadamard[][1]{}"} { %piby2 = arith.constant 1.57 : f64 %pi = arith.constant 3.14 : f64 %negpiby2 = arith.constant -3.14 : f64 - %q1 = quantum.custom "RZ"(%pi) %q0 : !quantum.bit - %q2 = quantum.custom "RY"(%piby2) %q1 : !quantum.bit - quantum.gphase(%negpiby2) + %q1 = quantum.custom "testRZ"(%pi) %q0 : !quantum.bit + %q2 = quantum.custom "testRY"(%piby2) %q1 : !quantum.bit return %q2 : !quantum.bit } -func.func @__builtin_h_to_rz_rx(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="Hadamard"} { +func.func @__builtin_h_to_rz_rx(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="testHadamard[][1]{}"} { %piby2 = arith.constant 1.57 : f64 %negpiby2 = arith.constant -3.14 : f64 - %q1 = quantum.custom "RZ"(%piby2) %q0 : !quantum.bit - %q2 = quantum.custom "RX"(%piby2) %q1 : !quantum.bit - %q3 = quantum.custom "RZ"(%piby2) %q2 : !quantum.bit - quantum.gphase(%negpiby2) + %q1 = quantum.custom "testRZ"(%piby2) %q0 : !quantum.bit + %q2 = quantum.custom "testRX"(%piby2) %q1 : !quantum.bit + %q3 = quantum.custom "testRZ"(%piby2) %q2 : !quantum.bit return %q3 : !quantum.bit } -func.func @__builtin_x_to_rx(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="PauliX"} { +func.func @__builtin_x_to_rx(%q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="testPauliX[][1]{}"} { %pi = arith.constant 3.14 : f64 %negpiby2 = arith.constant -3.14 : f64 - %q1 = quantum.custom "RX"(%pi) %q0 : !quantum.bit - quantum.gphase(%negpiby2) + %q1 = quantum.custom "testRX"(%pi) %q0 : !quantum.bit return %q1 : !quantum.bit } -func.func @__builtin_rx_to_rz_ry(%angle : f64, %q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="RX"} { +func.func @__builtin_rx_to_rz_ry(%angle : f64, %q0 : !quantum.bit) -> !quantum.bit attributes {target_gate="testRX[f64][1]{}"} { %piby2 = arith.constant 1.57 : f64 %negpiby2 = arith.constant -3.14 : f64 - %q1 = quantum.custom "RZ"(%piby2) %q0 : !quantum.bit - %q2 = quantum.custom "RY"(%angle) %q1 : !quantum.bit - %q3 = quantum.custom "RZ"(%negpiby2) %q2 : !quantum.bit + %q1 = quantum.custom "testRZ"(%piby2) %q0 : !quantum.bit + %q2 = quantum.custom "testRY"(%angle) %q1 : !quantum.bit + %q3 = quantum.custom "testRZ"(%negpiby2) %q2 : !quantum.bit return %q3 : !quantum.bit } From f2007bdf86d27f1510dc5e005c2a6b25ec42a555 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 23 Jul 2026 14:24:24 -0400 Subject: [PATCH 13/36] Factor out dummy op2 test classes --- frontend/test/lit/operator2_dummy_gates.py | 162 +++++++++++++++++ frontend/test/lit/test_operator.py | 168 +++--------------- frontend/test/pytest/operator2_dummy_gates.py | 162 +++++++++++++++++ 3 files changed, 344 insertions(+), 148 deletions(-) create mode 100644 frontend/test/lit/operator2_dummy_gates.py create mode 100644 frontend/test/pytest/operator2_dummy_gates.py diff --git a/frontend/test/lit/operator2_dummy_gates.py b/frontend/test/lit/operator2_dummy_gates.py new file mode 100644 index 0000000000..16c036d4c6 --- /dev/null +++ b/frontend/test/lit/operator2_dummy_gates.py @@ -0,0 +1,162 @@ +# 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. + +"""This file hosts a set of mock pennylane.core.Operator2 subclasses to be used for testing.""" + +# pylint: disable = missing-class-docstring + +import pennylane as qp + + +class NoParams(qp.core.Operator2): + wire_argnames = ("reg",) + + def __init__(self, reg): + super().__init__(reg=reg) + + +class NoParamsCustomOp(qp.core.Operator2): + + def __init__(self, wires): + super().__init__(wires=wires) + + +class SingleParam(qp.core.Operator2): + + dynamic_argnames = ("x",) + wire_argnames = ("reg",) + + def __init__(self, x, reg): + super().__init__(x, reg=reg) + + +class SingleParamCustomOp(qp.core.Operator2): + + dynamic_argnames = ("x",) + + def __init__(self, x, wires): + super().__init__(x, wires=wires) + + +class CompilableData(qp.core.Operator2): + + compilable_argnames = ("a", "b", "thing") + + def __init__(self, a, b, thing, wires): + super().__init__(a=a, b=b, thing=thing, wires=wires) + + +class MultipleRegisters(qp.core.Operator2): + + wire_argnames = ("reg1", "reg2") + + def __init__(self, reg1, reg2): + super().__init__(reg1=reg1, reg2=reg2) + + +class MultiParams(qp.core.Operator2): + + dynamic_argnames = ("a", "b", "c") + wire_argnames = ("reg",) + + def __init__(self, reg, a, b, c): + super().__init__(reg, a, b, c) + + +class MultiParamsCustom(qp.core.Operator2): + + dynamic_argnames = ("a", "b", "c") + + def __init__(self, wires, a, b, c): + super().__init__(wires, a, b, c) + + +class MultiRZ(qp.core.Operator2): + + dynamic_argnames = ("phi",) + + def __init__(self, phi, wires): + super().__init__(phi, wires) + + +class PauliRot(qp.core.Operator2): + + dynamic_argnames = ("phi",) + compilable_argnames = ("pauli_word",) + + def __init__(self, phi, pauli_word, wires): + super().__init__(phi, pauli_word, wires) + + +class GlobalPhase(qp.core.Operator2): + + dynamic_argnames = ("phi",) + wire_argnames = () + + def __init__(self, phi): + super().__init__(phi=phi) + + +class QubitUnitary(qp.core.Operator2): + + dynamic_argnames = ("matrix",) + + def __init__(self, matrix, wires): + super().__init__(matrix, wires) + + +class PCPhase(qp.core.Operator2): + + dynamic_argnames = ("phi",) + compilable_argnames = ("dim",) + + def __init__(self, phi, dim, wires): + super().__init__(phi, dim, wires) + + +class StaticData(qp.core.Operator2): + + static_argnames = ("label",) + wire_argnames = ("reg",) + + def __init__(self, label, reg): + super().__init__(label=label, reg=reg) + + +class HybridWires(qp.core.Operator2): + + hybrid_argnames = ("cwires",) + wire_argnames = ("cwires",) + + def __init__(self, cwires): + super().__init__(cwires=cwires) + + +class HybridNoOpArg(qp.core.Operator2): + + hybrid_argnames = ("angles",) + + def __init__(self, angles, wires): + super().__init__(angles, wires) + + +class HybridOpArg(qp.core.Operator2): + + dynamic_argnames = ("angle",) + hybrid_argnames = ("op",) + wire_argnames = ("cwires",) + static_argnames = ("n_iters",) + + def __init__(self, angle, op, cwires, n_iters=1): + super().__init__(angle, op, cwires, n_iters) diff --git a/frontend/test/lit/test_operator.py b/frontend/test/lit/test_operator.py index a461e23218..096e80ff45 100644 --- a/frontend/test/lit/test_operator.py +++ b/frontend/test/lit/test_operator.py @@ -13,21 +13,31 @@ # limitations under the License. """Tests for operator in Catalyst.""" -# pylint: disable = useless-parent-delegation, missing-function-docstring, missing-class-docstring, line-too-long +# pylint: disable = missing-function-docstring,line-too-long # RUN: %PYTHON %s | FileCheck %s import numpy as np import pennylane as qp - - -class NoParams(qp.core.Operator2): - - # have to use different wire argnames or will in up CustomOp - wire_argnames = ("reg",) - - def __init__(self, reg): - super().__init__(reg=reg) +from operator2_dummy_gates import ( + CompilableData, + GlobalPhase, + HybridNoOpArg, + HybridOpArg, + HybridWires, + MultiParams, + MultiParamsCustom, + MultipleRegisters, + MultiRZ, + NoParams, + NoParamsCustomOp, + PauliRot, + PCPhase, + QubitUnitary, + SingleParam, + SingleParamCustomOp, + StaticData, +) @qp.qjit(target="mlir", capture=True) @@ -143,12 +153,6 @@ def c_adjoint_and_controlled(): print(c_adjoint_and_controlled.mlir) -class NoParamsCustomOp(qp.core.Operator2): - - def __init__(self, wires): - super().__init__(wires=wires) - - @qp.qjit(target="mlir", capture=True) @qp.qnode(qp.device("null.qubit", wires=2)) def c_no_params_custom(): @@ -175,15 +179,6 @@ def c_no_params_custom(): print(c_no_params_custom.mlir) -class SingleParam(qp.core.Operator2): - - dynamic_argnames = ("x",) - wire_argnames = ("reg",) - - def __init__(self, x, reg): - super().__init__(x, reg=reg) - - @qp.qjit(target="mlir", capture=True) @qp.qnode(qp.device("null.qubit", wires=3)) def c_single_param(x: float): @@ -209,14 +204,6 @@ def c_single_param(x: float): print(c_single_param.mlir) -class SingleParamCustomOp(qp.core.Operator2): - - dynamic_argnames = ("x",) - - def __init__(self, x, wires): - super().__init__(x, wires=wires) - - @qp.qjit(target="mlir", capture=True) @qp.qnode(qp.device("null.qubit", wires=3)) def c_single_param_custom(x: float): @@ -237,14 +224,6 @@ def c_single_param_custom(x: float): print(c_single_param_custom.mlir) -class CompilableData(qp.core.Operator2): - - compilable_argnames = ("a", "b", "thing") - - def __init__(self, a, b, thing, wires): - super().__init__(a=a, b=b, thing=thing, wires=wires) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=3)) def c_compilable(): @@ -264,14 +243,6 @@ def c_compilable(): print(c_compilable.mlir) -class MultipleRegisters(qp.core.Operator2): - - wire_argnames = ("reg1", "reg2") - - def __init__(self, reg1, reg2): - super().__init__(reg1=reg1, reg2=reg2) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=5)) def c_multiple_registers(): @@ -301,16 +272,6 @@ def c_multiple_registers(): print(c_multiple_registers.mlir) -class MultiParams(qp.core.Operator2): - - dynamic_argnames = ("a", "b", "c") - wire_argnames = ("reg",) - - # note also having non-standard order with dynamic inputs after wires - def __init__(self, reg, a, b, c): - super().__init__(reg, a, b, c) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=1)) def c_multi_params(): @@ -329,15 +290,6 @@ def c_multi_params(): print(c_multi_params.mlir) -class MultiParamsCustom(qp.core.Operator2): - - dynamic_argnames = ("a", "b", "c") - - # note also having non-standard order with dynamic inputs after wires - def __init__(self, wires, a, b, c): - super().__init__(wires, a, b, c) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=1)) def c_multi_param_custom(): @@ -353,14 +305,6 @@ def c_multi_param_custom(): print(c_multi_param_custom.mlir) -class MultiRZ(qp.core.Operator2): - - dynamic_argnames = ("phi",) - - def __init__(self, phi, wires): - super().__init__(phi, wires) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=4)) def circuit_multirz(x: float): @@ -389,15 +333,6 @@ def circuit_multirz(x: float): print(circuit_multirz.mlir) -class PauliRot(qp.core.Operator2): - - dynamic_argnames = ("phi",) - compilable_argnames = ("pauli_word",) - - def __init__(self, phi, pauli_word, wires): - super().__init__(phi, pauli_word, wires) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=3)) def circuit_paulirot(x: float): @@ -432,15 +367,6 @@ def circuit_paulirot(x: float): print(circuit_paulirot.mlir) -class GlobalPhase(qp.core.Operator2): - - dynamic_argnames = ("phi",) - wire_argnames = () - - def __init__(self, phi): - super().__init__(phi=phi) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=3)) def circuit_gphase(x: float): @@ -459,14 +385,6 @@ def circuit_gphase(x: float): print(circuit_gphase.mlir) -class QubitUnitary(qp.core.Operator2): - - dynamic_argnames = ("matrix",) - - def __init__(self, matrix, wires): - super().__init__(matrix, wires) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("lightning.qubit", wires=3)) def circuit_qubitunitary(): @@ -497,15 +415,6 @@ def circuit_qubitunitary(): print(circuit_qubitunitary.mlir) -class PCPhase(qp.core.Operator2): - - dynamic_argnames = ("phi",) - compilable_argnames = ("dim",) - - def __init__(self, phi, dim, wires): - super().__init__(phi, dim, wires) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("lightning.qubit", wires=3)) def c_pcphase(x: float): @@ -534,15 +443,6 @@ def c_pcphase(x: float): print(c_pcphase.mlir) -class StaticData(qp.core.Operator2): - - static_argnames = ("label",) - wire_argnames = ("reg",) - - def __init__(self, label, reg): - super().__init__(label=label, reg=reg) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=1)) def c_static_data(): @@ -560,15 +460,6 @@ def c_static_data(): print(c_static_data.mlir) -class HybridWires(qp.core.Operator2): - - hybrid_argnames = ("cwires",) - wire_argnames = ("cwires",) - - def __init__(self, cwires): - super().__init__(cwires=cwires) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=2)) def c_hybrid_wires(): @@ -591,14 +482,6 @@ def c_hybrid_wires(): print(c_hybrid_wires.mlir) -class HybridNoOpArg(qp.core.Operator2): - - hybrid_argnames = ("angles",) - - def __init__(self, angles, wires): - super().__init__(angles, wires) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=2)) def c_hybrid_arg_not_op(): @@ -628,17 +511,6 @@ def c_hybrid_arg_not_op(): print(c_hybrid_arg_not_op.mlir) -class HybridOpArg(qp.core.Operator2): - - dynamic_argnames = ("angle",) - hybrid_argnames = ("op",) - wire_argnames = ("cwires",) - static_argnames = ("n_iters",) - - def __init__(self, angle, op, cwires, n_iters=1): - super().__init__(angle, op, cwires, n_iters) - - @qp.qjit(capture=True, target="mlir") @qp.qnode(qp.device("null.qubit", wires=4)) def c_hybrid_op_arg(x: float, y: float): diff --git a/frontend/test/pytest/operator2_dummy_gates.py b/frontend/test/pytest/operator2_dummy_gates.py new file mode 100644 index 0000000000..16c036d4c6 --- /dev/null +++ b/frontend/test/pytest/operator2_dummy_gates.py @@ -0,0 +1,162 @@ +# 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. + +"""This file hosts a set of mock pennylane.core.Operator2 subclasses to be used for testing.""" + +# pylint: disable = missing-class-docstring + +import pennylane as qp + + +class NoParams(qp.core.Operator2): + wire_argnames = ("reg",) + + def __init__(self, reg): + super().__init__(reg=reg) + + +class NoParamsCustomOp(qp.core.Operator2): + + def __init__(self, wires): + super().__init__(wires=wires) + + +class SingleParam(qp.core.Operator2): + + dynamic_argnames = ("x",) + wire_argnames = ("reg",) + + def __init__(self, x, reg): + super().__init__(x, reg=reg) + + +class SingleParamCustomOp(qp.core.Operator2): + + dynamic_argnames = ("x",) + + def __init__(self, x, wires): + super().__init__(x, wires=wires) + + +class CompilableData(qp.core.Operator2): + + compilable_argnames = ("a", "b", "thing") + + def __init__(self, a, b, thing, wires): + super().__init__(a=a, b=b, thing=thing, wires=wires) + + +class MultipleRegisters(qp.core.Operator2): + + wire_argnames = ("reg1", "reg2") + + def __init__(self, reg1, reg2): + super().__init__(reg1=reg1, reg2=reg2) + + +class MultiParams(qp.core.Operator2): + + dynamic_argnames = ("a", "b", "c") + wire_argnames = ("reg",) + + def __init__(self, reg, a, b, c): + super().__init__(reg, a, b, c) + + +class MultiParamsCustom(qp.core.Operator2): + + dynamic_argnames = ("a", "b", "c") + + def __init__(self, wires, a, b, c): + super().__init__(wires, a, b, c) + + +class MultiRZ(qp.core.Operator2): + + dynamic_argnames = ("phi",) + + def __init__(self, phi, wires): + super().__init__(phi, wires) + + +class PauliRot(qp.core.Operator2): + + dynamic_argnames = ("phi",) + compilable_argnames = ("pauli_word",) + + def __init__(self, phi, pauli_word, wires): + super().__init__(phi, pauli_word, wires) + + +class GlobalPhase(qp.core.Operator2): + + dynamic_argnames = ("phi",) + wire_argnames = () + + def __init__(self, phi): + super().__init__(phi=phi) + + +class QubitUnitary(qp.core.Operator2): + + dynamic_argnames = ("matrix",) + + def __init__(self, matrix, wires): + super().__init__(matrix, wires) + + +class PCPhase(qp.core.Operator2): + + dynamic_argnames = ("phi",) + compilable_argnames = ("dim",) + + def __init__(self, phi, dim, wires): + super().__init__(phi, dim, wires) + + +class StaticData(qp.core.Operator2): + + static_argnames = ("label",) + wire_argnames = ("reg",) + + def __init__(self, label, reg): + super().__init__(label=label, reg=reg) + + +class HybridWires(qp.core.Operator2): + + hybrid_argnames = ("cwires",) + wire_argnames = ("cwires",) + + def __init__(self, cwires): + super().__init__(cwires=cwires) + + +class HybridNoOpArg(qp.core.Operator2): + + hybrid_argnames = ("angles",) + + def __init__(self, angles, wires): + super().__init__(angles, wires) + + +class HybridOpArg(qp.core.Operator2): + + dynamic_argnames = ("angle",) + hybrid_argnames = ("op",) + wire_argnames = ("cwires",) + static_argnames = ("n_iters",) + + def __init__(self, angle, op, cwires, n_iters=1): + super().__init__(angle, op, cwires, n_iters) From d5435761655c9ed2a7036e39eb23f592745a9e3e Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 23 Jul 2026 14:25:58 -0400 Subject: [PATCH 14/36] delete tests with old UI --- .../from_plxpr/test_decompose_transform.py | 1191 ----------------- 1 file changed, 1191 deletions(-) delete mode 100644 frontend/test/pytest/from_plxpr/test_decompose_transform.py diff --git a/frontend/test/pytest/from_plxpr/test_decompose_transform.py b/frontend/test/pytest/from_plxpr/test_decompose_transform.py deleted file mode 100644 index e1d20bcaba..0000000000 --- a/frontend/test/pytest/from_plxpr/test_decompose_transform.py +++ /dev/null @@ -1,1191 +0,0 @@ -# Copyright 2025 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. -""" -This module tests the decompose transformation. -""" - -# pylint: disable=too-many-lines, too-many-public-methods - -from contextlib import nullcontext as does_not_raise -from functools import partial - -import numpy as np -import pennylane as qp -import pytest -from jax.core import ShapedArray -from pennylane.decomposition import controlled_resource_rep -from pennylane.exceptions import DecompositionError, DecompositionWarning -from pennylane.typing import TensorLike -from pennylane.wires import WiresLike -from pennylane_lightning.lightning_qubit.lightning_qubit import ( - stopping_condition as lightning_stopping_condition, -) - -from catalyst import CompileError -from catalyst.jax_primitives import decomposition_rule -from catalyst.passes import graph_decomposition - - -def _normalize_gate_types(gate_types): - """ - TODO: Remove this function once PennyLane tape-based resource counting specs format - is unified with the updated resource tracking specs format. - - Normalize gate type names by stripping NullQubit suffixes (e.g. 'PauliRot-Phi-w4') - back to the base name ('PauliRot') and summing counts for matching base names. - """ - result = {} - for k, v in gate_types.items(): - base = k.split("-")[0] - result[base] = result.get(base, 0) + v - return result - - -class TestGraphDecomposition: - """Test the graph-decomposition built-in transform.""" - - @pytest.mark.parametrize("weight", [1, 1.0], ids=["int", "float"]) - def test_gateset_with_weights(self, weight): - """Tests that a gate_set with weights works correctly. - - Regression test for https://github.com/PennyLaneAI/catalyst/issues/2766.""" - - @qp.qjit(capture=True) - @graph_decomposition(gate_set={qp.RX: weight, qp.RY: weight, qp.RZ: weight}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit(x, y, z): - qp.Rot(x, y, z, wires=0) - return qp.expval(qp.Z(0)) - - x, y, z = 0.0, qp.numpy.pi, 0.0 - - assert qp.math.allclose([-1], circuit(x, y, z)) - - expected_resources = {"RY": 1, "RZ": 2} - resources = qp.specs(circuit, level="device")(x, y, z)["resources"].gate_types - assert resources == expected_resources - - def test_with_precompiled_rule(self): - """Test graph-decomposition with precompiled rules are handled correctly.""" - - @qp.qjit(capture=True) - @graph_decomposition(gate_set=[qp.RX, qp.RY, qp.RZ]) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit(x, y, z): - qp.Rot(x, y, z, wires=0) - return qp.expval(qp.Z(0)) - - x = 0.5 - y = 0.3 - z = 0.2 - - assert qp.math.allclose([0.9553364891256059], circuit(x, y, z)) - - expected_resources = {"RY": 1, "RZ": 2} - resources = qp.specs(circuit, level="device")(x, y, z)["resources"].gate_types - assert resources == expected_resources - - def test_decompose_multi_qubit_gates_precompiled(self): - """Test that multi-qubit gates are decomposed correctly using precompiled rules.""" - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={"RY", "RX", "CNOT", "Hadamard", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - def circuit(): - qp.SingleExcitation(0.5, wires=[0, 1]) - qp.SingleExcitationPlus(0.5, wires=[0, 1]) - qp.SingleExcitationMinus(0.5, wires=[0, 1]) - qp.DoubleExcitation(0.5, wires=[0, 1, 2, 3]) - return qp.expval(qp.Z(0)) - - expected_resources = {"GlobalPhase": 6, "RX": 6, "RY": 30, "CNOT": 24, "Hadamard": 12} - resources = qp.specs(circuit, level="device")()["resources"].gate_types - assert resources == expected_resources - - def test_alt_decomps(self): - """Test the conversion of a circuit with a custom decomposition.""" - - @decomposition_rule(op_type=qp.CNOT) - def my_cnot(wires, **__): - qp.H(wires=wires[1]) - qp.CZ(wires=wires) - qp.H(wires=wires[1]) - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={"H", "CZ", "GlobalPhase"}, - alt_decomps={qp.CNOT: [my_cnot]}, - _builtin_rule_path=None, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - qp.H(0) - qp.CNOT(wires=[0, 1]) - - # register custom decomposition rules - my_cnot(ShapedArray((2,), int)) - - return qp.state() - - expected_resources = {"CZ": 1, "Hadamard": 3} - resources = qp.specs(circuit, level="device")()["resources"].gate_types - assert resources == expected_resources - - def test_fixed_rules(self): - """Test the decompose lowering pass with custom decomposition rules.""" - - @decomposition_rule(op_type=qp.RY) - def rz_rx(phi, wires: WiresLike, **__): - """Decomposition of RY gate using RZ and RX gates.""" - qp.RZ(-np.pi / 2, wires=wires) - qp.RX(phi, wires=wires) - qp.RZ(np.pi / 2, wires=wires) - - @decomposition_rule(op_type=qp.Rot) - def rz_ry_rz(phi, theta, omega, wires: WiresLike, **__): - """Decomposition of Rot gate using RZ and RY gates.""" - qp.RZ(phi, wires=wires) - qp.RY(theta, wires=wires) - qp.RZ(omega, wires=wires) - - @decomposition_rule(op_type=qp.PauliY) - def ry_gp(wires: WiresLike, **__): - """Decomposition of PauliY gate using RY and GlobalPhase gates.""" - qp.RY(np.pi, wires=wires) - qp.GlobalPhase(-np.pi / 2, wires=wires) - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={"RX", "RZ", "GlobalPhase"}, - fixed_decomps={ - qp.RY: rz_rx, - qp.Rot: rz_ry_rz, - qp.PauliY: ry_gp, - }, - _builtin_rule_path=None, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - def circuit(): - qp.RY(0.5, wires=0) - qp.Rot(0.2, 0.3, 0.4, wires=1) - qp.PauliY(wires=2) - qp.Rot(0.2, 0.3, 0.4, wires=2) - qp.RY(0.5, wires=1) - - # register custom decomposition rules - rz_rx(float, int) - rz_ry_rz(float, float, float, int) - ry_gp(int) - - return qp.expval(qp.Z(0)) - - expected_resources = {"GlobalPhase": 1, "RX": 5, "RZ": 14} - resources = qp.specs(circuit, level="device")()["resources"].gate_types - assert resources == expected_resources - - def test_multi_passes(self): - """Test the graph_decomposition pass with other passes.""" - - @qp.qjit(capture=True) - @qp.transforms.merge_rotations - @graph_decomposition( - gate_set={"RZ", "RY", "CNOT", "GlobalPhase"}, - ) - @qp.transforms.cancel_inverses - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit(): - qp.PauliX(0) - qp.PauliX(0) - qp.RX(0.1, wires=0) - return qp.expval(qp.PauliX(0)) - - expected_resources = {"RY": 1, "RZ": 2} - resources = qp.specs(circuit, level="device")()["resources"].gate_types - assert resources == expected_resources - - def test_multi_graph_decomposition(self): - """Test that multiple graph-decomposition builtin transforms can be applied.""" - - @decomposition_rule(op_type=qp.PauliX) - def x_to_rx(wire: int): - qp.RX(np.pi, wire) - - @decomposition_rule(op_type=qp.PauliY) - def y_to_ry(wire: int): - qp.RY(np.pi, wire) - - @decomposition_rule(op_type=qp.Hadamard) - def h_to_rx_ry(wire: int): - qp.RX(np.pi / 2, wire) - qp.RY(np.pi / 2, wire) - - @qp.qjit(capture=True) - @graph_decomposition(gate_set={qp.Rot}) - @qp.transforms.merge_rotations - @graph_decomposition( - gate_set={qp.RX, qp.RY}, - fixed_decomps={qp.PauliX: x_to_rx, qp.PauliY: y_to_ry}, - alt_decomps={qp.H: [h_to_rx_ry]}, - ) - @qp.transforms.cancel_inverses - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(x: float, y: float): - qp.H(0) - qp.H(0) - qp.RX(x, wires=0) - qp.PauliX(0) - qp.RY(y, wires=0) - qp.PauliY(0) - qp.RY(x + y, wires=0) - - # register custom decomposition rules - x_to_rx(int) - y_to_ry(int) - h_to_rx_ry(int) - - return qp.state() - - expected_resources = {"Rot": 2} - resources = qp.specs(circuit, level="device")(1.23, 4.56)["resources"].gate_types - assert resources == expected_resources - - @pytest.mark.xfail( - reason="only quantum.custom gates are currently supported with graph_decomposition" - ) - def test_multirz(self): - """Test that TensorLike parameters in MultiRZ are handled correctly in rules.""" - - @graph_decomposition(op_type="MultiRZ") - def custom_multirz(params: TensorLike, wires: WiresLike, **__): - qp.CNOT(wires=(wires[2], wires[1])) - qp.CNOT(wires=(wires[1], wires[0])) - qp.RZ(params, wires=wires[0]) - qp.CNOT(wires=(wires[1], wires[0])) - qp.CNOT(wires=(wires[2], wires[1])) - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={qp.RY, qp.RX, qp.CNOT}, - fixed_decomps={qp.MultiRZ: custom_multirz}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3), shots=1000) - def circuit(x, y): - qp.MultiRZ(x + y, wires=[0, 1, 2]) - - # register custom decomposition rules - custom_multirz(TensorLike, [int, int, int]) - - return qp.expval(qp.Z(0)) - - expected_resources = {"RX": 1, "RY": 2, "CNOT": 4} - resources = qp.specs(circuit, level="device")(0.5, 0.3)["resources"].gate_types - assert resources == expected_resources - - def test_with_subroutine(self): - """Test that decompositions can happen inside subroutines.""" - - @qp.templates.Subroutine - def f(x, wires): - qp.IsingXX(x, wires) - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set=qp.gate_sets.ROTATIONS_PLUS_CNOT, - ) - @qp.qnode(qp.device("lightning.qubit", wires=5)) - def circuit(): - f(0.5, (0, 1)) - f(1.2, (2, 3)) - return qp.expval(qp.Z(0)), qp.expval(qp.Z(2)) - - resources = qp.specs(circuit, level="device")().resources.gate_types - assert resources == {"RX": 2, "CNOT": 4} - - r1, r2 = circuit() - assert qp.math.allclose(r1, np.cos(0.5)) - assert qp.math.allclose(r2, np.cos(1.2)) - - def test_ftqc_rotxzx(self): - """Test qp.ftqc.RotXZX with alt_decomps.""" - - @decomposition_rule(op_type="RotXZX") - def _xzx_decompose(phi, theta, omega, wires, **__): - qp.RX(phi, wires=wires) - qp.RZ(theta, wires=wires) - qp.RX(omega, wires=wires) - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={"CNOT", "GlobalPhase", "RX", "RZ", "PauliRot"}, - alt_decomps={qp.ftqc.RotXZX: [_xzx_decompose]}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - qp.ftqc.RotXZX(0.5, 0.3, 0.7, wires=0) - - _xzx_decompose(float, float, float, int) - return qp.expval(qp.X(0)) - - expected_resources = {"RX": 2, "RZ": 1} - resources = qp.specs(circuit, level="device")()["resources"].gate_types - assert resources == expected_resources - - def test_empty_rule(self): - """Test that a decomposition rule with no ops is handled correctly.""" - - @decomposition_rule(op_type="PauliX") - def empty_decomp(_wire): - pass - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={"PauliY"}, - fixed_decomps={"PauliX": empty_decomp}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit(): - qp.X(0) - qp.Y(0) - - # register the empty decomposition rule - empty_decomp(int) - - return qp.expval(qp.Z(0)) - - expected_resources = {"PauliY": 1} - resources = qp.specs(circuit, level="device")()["resources"].gate_types - assert resources == expected_resources - - @pytest.mark.xfail( - reason="graph-decomposition supports pre-compiled rules, alt_decomps and fix_decomps" - ) - def test_ftqc_custom_ops(self): - """Test that ftqc Ops cannot be decomposed without defining rules.""" - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={"CNOT", "GlobalPhase", "RX", "RZ", "PauliRot"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - qp.ftqc.RotXZX(0.5, 0.3, 0.7, wires=0) - return qp.expval(qp.X(0)) - - expected_resources = {"RX": 2, "RZ": 1} - resources = qp.specs(circuit, level="device")()["resources"].gate_types - assert resources == expected_resources - - @pytest.mark.xfail(reason="graph-decomposition does not yet support adjoint or ctrl operations") - def test_adjoint(self): - """Test the graph_decomposition pass with adjoint operations.""" - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={"RY", "RX", "CZ", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - def circuit(): - qp.adjoint(qp.Hadamard(wires=2)) - qp.adjoint(qp.CNOT(wires=[0, 1])) - qp.adjoint(qp.RX(0.5, wires=3)) - qp.adjoint(qp.Toffoli(wires=[0, 1, 2])) - return qp.expval(qp.Z(0)) - - expected_resources = {"GlobalPhase": 24, "CZ": 7, "RX": 25, "RY": 65} - resources = qp.specs(circuit, level="device")()["resources"].gate_types - assert resources == expected_resources - - @pytest.mark.xfail(reason="graph-decomposition does not yet support adjoint or ctrl operations") - def test_ctrl(self): - """Test the graph_decomposition pass with controlled operations.""" - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={"RX", "RZ", "H", "CZ", "PauliRot"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - qp.ctrl(qp.Hadamard(wires=1), 0) - qp.ctrl(qp.RY, control=0)(0.5, 1) - qp.ctrl(qp.PauliX, control=0)(1) - return qp.expval(qp.Z(0)) - - expected_resources = {"RX": 1, "RZ": 2, "H": 2, "CZ": 1} - resources = qp.specs(circuit, level="device")()["resources"].gate_types - assert resources == expected_resources - - @pytest.mark.xfail(reason="graph-decomposition does not yet support work wires") - def test_work_wires(self): - """Test that graph decomposition supports work_wires.""" - - @decomposition_rule(op_type=qp.CRX) - def my_decomp(angle, wires, **_): - def true_func(): - qp.CNOT(wires) - - with qp.allocate(2, state="any", restored=True) as w: - qp.H(w[0]) - qp.H(w[0]) - qp.X(w[1]) - qp.X(w[1]) - return - - def false_func(): - with qp.allocate(1, state="any", restored=False) as w: - qp.H(w) - - m = qp.measure(wires[0]) - qp.cond(m, qp.CNOT)(wires) - return - - qp.cond(angle > 1.2, true_func, false_func)() - - @qp.qjit(capture=True) - @graph_decomposition( - gate_set={qp.CNOT, qp.H, qp.X, "Conditional", "MidMeasure"}, - fixed_decomps={qp.CRX: my_decomp}, - num_work_wires=7, - ) - @qp.qnode(qp.device("lightning.qubit", wires=9)) - def circuit(): - qp.CRX(1.7, wires=[0, 1]) - qp.CRX(-7.2, wires=[0, 1]) - return qp.state() - - def test_non_custom_op(self): - """Test that the graph correctly registers non-custom ops.""" - - with pytest.raises(CompileError): - - @qp.qjit - @graph_decomposition(gate_set={qp.X}) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(x: float, y: float): # pylint: disable=unused-argument - qp.PauliRot(0.1, "ZZ", wires=[0, 1]) - return qp.state() - - circuit() - - def test_paulirot_python_decomp(self): - """Test that paulirot is successfully decomposed by compile-time lowered rules.""" - - def circuit(): - qp.PauliRot(0.3, "YXZ", [0, 1, 2]) - return qp.state() - - qnode = qp.QNode(circuit, qp.device("null.qubit", wires=3)) - - without_qjit = qnode() - - with_qjit = qp.qjit( - graph_decomposition(qnode, gate_set={qp.H, qp.RX, qp.MultiRZ, qp.GlobalPhase}) - )() - - assert np.allclose(without_qjit, with_qjit) - - -class TestPlxPRDecomposition: - """Test the PLxPR-based graph-based decomposition integration with from_plxpr.""" - - @pytest.mark.usefixtures("use_capture_dgraph") - def test_with_multiple_decomps_transforms(self): - """Test that a circuit with multiple decompositions and transforms can be converted.""" - - @qp.qjit(target="mlir") - @partial( - qp.transforms.decompose, - gate_set={"RX", "RY"}, - ) - @partial( - qp.transforms.decompose, - gate_set={"NOT", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=0)) - def circuit(x): - qp.GlobalPhase(x) - return qp.expval(qp.PauliX(0)) - - with pytest.raises( - NotImplementedError, match="Multiple decomposition transforms are not yet supported." - ): - circuit(0.2) - - @pytest.mark.usefixtures("use_capture_dgraph") - def test_fallback_warnings(self): - """Test the fallback to legacy decomposition system with warnings.""" - - @qp.qjit - @partial(qp.transforms.decompose, gate_set={qp.GlobalPhase}) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(x): - qp.Hadamard(x) - return qp.state() - - # TODO: RZ/RX warnings should not be raised, remove (PL issue #8885) - with pytest.warns(UserWarning, match="Falling back to the legacy decomposition system"): - with pytest.warns( - DecompositionWarning, match="unable to find a decomposition for {'Hadamard'}" - ): - with pytest.warns(UserWarning, match="Operator RX does not define"): - with pytest.warns(UserWarning, match="Operator RZ does not define"): - circuit(0) - - def test_decompose_lowering_on_empty_circuit(self): - """Test that the decompose lowering pass works on an empty circuit.""" - qp.decomposition.enable_graph() - - @partial( - qp.transforms.decompose, - gate_set={"RX", "RY"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - return qp.expval(qp.X(0)) - - without_qjit = circuit() - - with_qjit = qp.qjit(circuit, capture=True) - - assert qp.math.allclose(without_qjit, with_qjit()) - - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_alt_decomps(self): - """Test the conversion of a circuit with a custom decomposition.""" - qp.decomposition.enable_graph() - - @qp.register_resources({qp.H: 2, qp.CZ: 1}) - def my_cnot(wires, **__): - qp.H(wires=wires[1]) - qp.CZ(wires=wires) - qp.H(wires=wires[1]) - - @partial( - qp.transforms.decompose, - gate_set={"H", "CZ", "GlobalPhase"}, - alt_decomps={qp.CNOT: [my_cnot]}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - qp.H(0) - qp.CNOT(wires=[0, 1]) - return qp.state() - - qjited_circuit = qp.qjit(circuit, capture=True) - - expected = np.array([1, 0, 0, 1]) / np.sqrt(2) - assert qp.math.allclose(qjited_circuit(), expected) - - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - resources = qp.specs(qjited_circuit, level="device")()["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_fixed_rules(self): - """Test the decompose lowering pass with custom decomposition rules.""" - qp.decomposition.enable_graph() - - @qp.register_resources({qp.RZ: 2, qp.RX: 1}) - def rz_rx(phi, wires: WiresLike, **__): - """Decomposition of RY gate using RZ and RX gates.""" - qp.RZ(-np.pi / 2, wires=wires) - qp.RX(phi, wires=wires) - qp.RZ(np.pi / 2, wires=wires) - - @qp.register_resources({qp.RZ: 2, qp.RY: 1}) - def rz_ry_rz(phi, theta, omega, wires: WiresLike, **__): - """Decomposition of Rot gate using RZ and RY gates.""" - qp.RZ(phi, wires=wires) - qp.RY(theta, wires=wires) - qp.RZ(omega, wires=wires) - - @qp.register_resources({qp.RY: 1, qp.GlobalPhase: 1}) - def ry_gp(wires: WiresLike, **__): - """Decomposition of PauliY gate using RY and GlobalPhase gates.""" - qp.RY(np.pi, wires=wires) - qp.GlobalPhase(-np.pi / 2, wires=wires) - - @partial( - qp.transforms.decompose, - gate_set={"RX", "RZ", "GlobalPhase"}, - fixed_decomps={ - qp.RY: rz_rx, - qp.Rot: rz_ry_rz, - qp.PauliY: ry_gp, - }, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - def circuit(): - qp.RY(0.5, wires=0) - qp.Rot(0.2, 0.3, 0.4, wires=1) - qp.PauliY(wires=2) - qp.Rot(0.2, 0.3, 0.4, wires=2) - qp.RY(0.5, wires=1) - qp.PauliX(wires=0) - return qp.expval(qp.Z(0)) - - without_qjit = circuit() - - with_qjit = qp.qjit(circuit, capture=True) - - assert qp.math.allclose(without_qjit, with_qjit()) - - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_tensorlike(self): - """Test that TensorLike parameters are handled correctly in rules.""" - qp.decomposition.enable_graph() - - @qp.register_resources({qp.RZ: 1, qp.CNOT: 4}) - def custom_multirz(params: TensorLike, wires: WiresLike, **__): - qp.CNOT(wires=(wires[2], wires[1])) - qp.CNOT(wires=(wires[1], wires[0])) - qp.RZ(params, wires=wires[0]) - qp.CNOT(wires=(wires[1], wires[0])) - qp.CNOT(wires=(wires[2], wires[1])) - - @partial( - qp.transforms.decompose, - gate_set={"RY", "RX", qp.CNOT}, - fixed_decomps={qp.MultiRZ: custom_multirz}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3), shots=1000) - def circuit(x, y): - qp.MultiRZ(x + y, wires=[0, 1, 2]) - return qp.expval(qp.Z(0)) - - x = 0.5 - y = 0.3 - - without_qjit = circuit(x, y) - - with_qjit = qp.qjit(circuit, capture=True) - - assert qp.math.allclose(without_qjit, with_qjit(x, y)) - expected_resources = qp.specs(circuit, level="device")(x, y)["resources"].gate_types - resources = qp.specs(with_qjit, level="device")(x, y)["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_inordered_params(self): - """Test that unordered parameters in rules are handled correctly.""" - - qp.decomposition.enable_graph() - - @partial(qp.transforms.decompose, gate_set=[qp.RX, qp.RY, qp.RZ]) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit(x, y, z): - qp.Rot(x, y, z, wires=0) - return qp.expval(qp.Z(0)) - - x = 0.5 - y = 0.3 - z = 0.2 - - without_qjit = circuit(x, y, z) - - with_qjit = qp.qjit(circuit, capture=True) - - assert qp.math.allclose(without_qjit, with_qjit(x, y, z)) - - expected_resources = qp.specs(circuit, level="device")(x, y, z)["resources"].gate_types - resources = qp.specs(with_qjit, level="device")(x, y, z)["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_decompose_with_stopping_condition(self): - """Test that decompose with stopping_condition uses plxpr decomposition correctly. - - When stopping_condition is passed to qp.transforms.decompose, from_plxpr uses - the plxpr decompose path (no graph), passing stopping_condition to the transform. - This test ensures that path compiles and produces correct results. - """ - qp.decomposition.enable_graph() - - def stopping_condition(op): - return op.name == "MultiRZ" - - @partial( - qp.transforms.decompose, - gate_set=[qp.RX, qp.RY, qp.RZ], - stopping_condition=stopping_condition, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(x, y, z): - qp.Rot(x, y, z, wires=0) - qp.MultiRZ(0.5, wires=[0, 1]) - return qp.expval(qp.PauliZ(0)) - - x, y, z = 0.5, 0.3, 0.2 - without_qjit = circuit(x, y, z) - - with_qjit = qp.qjit(circuit, capture=True) - assert qp.math.allclose(without_qjit, with_qjit(x, y, z)) - - expected_resources = qp.specs(circuit, level="device")(x, y, z)["resources"].gate_types - resources = qp.specs(with_qjit, level="device")(x, y, z)["resources"].gate_types - assert "MultiRZ" in resources - assert "MultiRZ" in expected_resources - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_decompose_with_lightning_stopping_condition(self): - """Test that decompose with stopping_condition using Lightning's stopping condition.""" - qp.decomposition.enable_graph() - device = qp.device("null.qubit", wires=4) - - @partial( - qp.transforms.decompose, - gate_set=[qp.CNOT, qp.PauliZ], - stopping_condition=lightning_stopping_condition, - ) - @qp.qnode(device) - def circuit(x): - qp.PauliRot(x, "XYZZ", wires=[0, 1, 2, 3]) - qp.StatePrep(np.array([1, 0, 0, 0]), wires=range(2)) - return qp.expval(qp.PauliZ(0)) - - x = 0.5 - without_qjit = circuit(x) - - with_qjit = qp.qjit(circuit, capture=True) - assert qp.math.allclose(without_qjit, with_qjit(x)) - - expected_resources = qp.specs(circuit, level="device")(x)["resources"].gate_types - resources = qp.specs(with_qjit, level="device")(x)["resources"].gate_types - assert any(k.startswith("PauliRot") for k in expected_resources) - assert any(k.startswith("PauliRot") for k in resources) - assert not any(k.startswith("StatePrep") for k in expected_resources) - assert not any(k.startswith("StatePrep") for k in resources) - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - @pytest.mark.skip( - reason="inconsistent type and error msg across gcc/clang on arm/x86 for undefined symbols" - ) - def test_gateset_with_rotxzx(self): - """Test the runtime raises an error if RotXZX is not decomposed.""" - qp.decomposition.enable_graph() - - @partial( - qp.transforms.decompose, - gate_set={qp.ftqc.RotXZX}, - ) - @qp.qnode(qp.device("null.qubit", wires=2)) - def circuit(): - qp.ftqc.RotXZX(0.5, 0.3, 0.7, wires=0) - return qp.expval(qp.X(0)) - - with pytest.raises( - OSError, - match="undefined symbol", # ___catalyst__qis__RotXZX - ): - qp.qjit(circuit, capture=True)() - qp.decomposition.disable_graph() - - def test_ftqc_rotxzx(self): - """Test that FTQC RotXZX decomposition works with from_plxpr.""" - qp.decomposition.enable_graph() - - @partial( - qp.transforms.decompose, - gate_set={"CNOT", "GlobalPhase", "RX", "RZ", "PauliRot"}, - ) - @qp.qnode(qp.device("null.qubit", wires=2)) - def circuit(): - qp.ftqc.RotXZX(0.5, 0.3, 0.7, wires=0) - qp.ctrl(qp.ftqc.RotXZX(0.4, 0.2, 0.6, wires=1), control=0) - return qp.expval(qp.X(0)) - - without_qjit = circuit() - - with_qjit = qp.qjit(circuit, capture=True) - - assert qp.math.allclose(without_qjit, with_qjit()) - - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - @pytest.mark.xfail(reason="unstable global phase numbers", strict=False) - def test_multirz(self): - """Test that multirz decomposition works with from_plxpr.""" - qp.decomposition.enable_graph() - - @partial( - qp.transforms.decompose, - gate_set={"X", "Y", "Z", "S", "H", "CNOT", "RZ", "Rot", "GlobalPhase"}, - ) - @qp.qnode(qp.device("null.qubit", wires=4)) - def circuit(): - qp.Hadamard(0) - qp.ctrl(qp.MultiRZ(0.345, wires=[1, 2]), control=0) - qp.adjoint(qp.MultiRZ(0.25, wires=[1, 2])) - qp.MultiRZ(0.5, wires=[0, 1]) - qp.MultiRZ(0.5, wires=[0]) - qp.MultiRZ(0.5, wires=[0, 1, 3]) - return qp.expval(qp.X(0)) - - with_qjit = qp.qjit(circuit, capture=True) - - result_with_qjit = with_qjit() - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - - result_without_qjit = circuit() - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - assert qp.math.allclose(result_without_qjit, result_with_qjit) - qp.decomposition.disable_graph() - - def test_gphase(self): - """Test that the decompose lowering pass works with GlobalPhase.""" - qp.decomposition.enable_graph() - - @partial( - qp.transforms.decompose, - gate_set={"RX", "RY", "GlobalPhase"}, - ) - @qp.qnode(qp.device("null.qubit", wires=1)) - def circuit(): - qp.GlobalPhase(0.5) - qp.ctrl(qp.GlobalPhase, control=0)(0.3) - qp.ctrl(qp.GlobalPhase, control=0)(phi=0.3, wires=[1, 2]) - return qp.expval(qp.Z(0)) - - without_qjit = circuit() - with_qjit = qp.qjit(circuit, capture=True) - - assert qp.math.allclose(without_qjit, with_qjit()) - - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_multi_qubits(self): - """Test that the decompose lowering pass works with multi-qubit gates.""" - qp.decomposition.enable_graph() - - @partial( - qp.transforms.decompose, - gate_set={"RY", "RX", "CNOT", "Hadamard", "GlobalPhase"}, - ) - @qp.qnode(qp.device("null.qubit", wires=4)) - def circuit(): - qp.SingleExcitation(0.5, wires=[0, 1]) - qp.SingleExcitationPlus(0.5, wires=[0, 1]) - qp.SingleExcitationMinus(0.5, wires=[0, 1]) - qp.DoubleExcitation(0.5, wires=[0, 1, 2, 3]) - return qp.expval(qp.Z(0)) - - without_qjit = circuit() - - with_qjit = qp.qjit(circuit, capture=True) - assert qp.math.allclose(without_qjit, with_qjit()) - - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_adjoint(self): - """Test the decompose lowering pass with adjoint operations.""" - qp.decomposition.enable_graph() - - @partial( - qp.transforms.decompose, - gate_set={"RY", "RX", "CZ", "GlobalPhase"}, - ) - @qp.qnode(qp.device("null.qubit", wires=4)) - def circuit(): - qp.adjoint(qp.Hadamard(wires=2)) - qp.adjoint(qp.CNOT(wires=[0, 1])) - qp.adjoint(qp.RX(0.5, wires=3)) - qp.adjoint(qp.Toffoli(wires=[0, 1, 2])) - return qp.expval(qp.Z(0)) - - without_qjit = circuit() - - with_qjit = qp.qjit(circuit, capture=True) - - assert qp.math.allclose(without_qjit, with_qjit()) - - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_ctrl(self): - """Test the decompose lowering pass with controlled operations.""" - qp.decomposition.enable_graph() - - @partial( - qp.transforms.decompose, - gate_set={"RX", "RZ", "H", "CZ", "PauliRot"}, - ) - @qp.qnode(qp.device("null.qubit", wires=2)) - def circuit(): - qp.ctrl(qp.Hadamard(wires=1), 0) - qp.ctrl(qp.RY, control=0)(0.5, 1) - qp.ctrl(qp.PauliX, control=0)(1) - return qp.expval(qp.Z(0)) - - without_qjit = circuit() - - with_qjit = qp.qjit(circuit, capture=True) - - assert qp.math.allclose(without_qjit, with_qjit()) - - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - def test_template_qft(self): - """Test the decompose lowering pass with the QFT template.""" - qp.decomposition.enable_graph() - - @partial( - qp.transforms.decompose, - gate_set={"RX", "RY", "CNOT", "GlobalPhase", "PauliRot"}, - ) - @qp.qnode(qp.device("null.qubit", wires=4)) - def circuit(): - qp.QFT(wires=[0, 1, 2, 3]) - return qp.expval(qp.Z(0)) - - with_qjit = qp.qjit(circuit, capture=True) - result_with_qjit = with_qjit() - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - - result_without_qjit = circuit() - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - assert qp.math.allclose(result_without_qjit, result_with_qjit) - qp.decomposition.disable_graph() - - def test_multi_passes(self): - """Test the decompose lowering pass with multiple passes.""" - qp.decomposition.enable_graph() - - @qp.transforms.merge_rotations - @qp.transforms.cancel_inverses - @partial( - qp.transforms.decompose, - gate_set=frozenset({"RZ", "RY", "CNOT", "GlobalPhase"}), - ) - @qp.qnode(qp.device("null.qubit", wires=1)) - def circuit(): - qp.PauliX(0) - qp.PauliX(0) - qp.RX(0.1, wires=0) - return qp.expval(qp.PauliX(0)) - - without_qjit = circuit() - - with_qjit = qp.qjit(circuit, capture=True) - - assert qp.math.allclose(without_qjit, with_qjit()) - - expected_resources = qp.specs(circuit, level="device")()["resources"].gate_types - resources = qp.specs(with_qjit, level="device")()["resources"].gate_types - assert _normalize_gate_types(resources) == _normalize_gate_types(expected_resources) - qp.decomposition.disable_graph() - - @pytest.mark.parametrize( - "num_work_wires,expectation", - [ - (0, pytest.raises(DecompositionError)), - (2, pytest.raises(DecompositionError)), - (3, does_not_raise()), - (7, does_not_raise()), - ], - ) - def test_work_wires(self, num_work_wires, expectation): - """ - Test that graph decomposition raises the correct exception when given an insufficient - number of work wires, and passes otherwise. - """ - qp.decomposition.enable_graph() - - @qp.register_resources( - {qp.CNOT: 3, qp.H: 1, qp.X: 1, qp.ops.op_math.Conditional: 2}, - work_wires={ - "borrowed": 2, - "garbage": 1, - }, - ) - def my_decomp(angle, wires, **_): - def true_func(): - qp.CNOT(wires) - - with qp.allocate(2, state="any", restored=True) as w: - qp.H(w[0]) - qp.H(w[0]) - qp.X(w[1]) - qp.X(w[1]) - - return - - def false_func(): - with qp.allocate(1, state="any", restored=False) as w: - qp.H(w) - - m = qp.measure(wires[0]) - - qp.cond(m, qp.CNOT)(wires) - - return - - qp.cond(angle > 1.2, true_func, false_func)() - - with expectation: - - @qp.qjit(capture=True) - @partial( - qp.transforms.decompose, - gate_set={qp.CNOT, qp.H, qp.X, "Conditional", "MidMeasure"}, - fixed_decomps={qp.CRX: my_decomp}, - num_work_wires=num_work_wires, - ) - @qp.qnode(qp.device("null.qubit", wires=9)) - def circuit(): - qp.CRX(1.7, wires=[0, 1]) - qp.CRX(-7.2, wires=[0, 1]) - return qp.state() - - qp.decomposition.disable_graph() - - def test_decomp_inside_subroutine(self): - """Test that decompositions can happen inside subroutines.""" - qp.decomposition.enable_graph() - - @qp.templates.Subroutine - def f(x, wires): - qp.IsingXX(x, wires) - - @qp.qjit(capture=True) - @qp.decompose(gate_set=qp.gate_sets.ROTATIONS_PLUS_CNOT) - @qp.qnode(qp.device("lightning.qubit", wires=5)) - def c(): - f(0.5, (0, 1)) - f(1.2, (2, 3)) - return qp.expval(qp.Z(0)), qp.expval(qp.Z(2)) - - resources = qp.specs(c, level="device")().resources.gate_types - assert resources == {"RX": 2, "CNOT": 4} - - r1, r2 = c() - assert qp.math.allclose(r1, np.cos(0.5)) - assert qp.math.allclose(r2, np.cos(1.2)) - qp.decomposition.disable_graph() - - def test_unknown_op_in_solution_raises(self): - """An op that ends up in the decomposition graph solution but is - neither in the captured circuit nor in ``COMPILER_OPS_FOR_DECOMPOSITION`` - and is not a symbolic op must raise a clear ValueError so the user knows the - wire count cannot be inferred. - """ - - with qp.decomposition.local_decomps(): - - class _UnknownOp(qp.operation.Operation): - num_wires = 1 - num_params = 0 - name = "_UnknownOp" - - def _unknown_resources(): - return {qp.resource_rep(qp.PauliX): 1} - - @qp.register_resources(_unknown_resources) - def _unknown_decomp(wires): - qp.PauliX(wires) - - qp.add_decomps(_UnknownOp, _unknown_decomp) - - def _rx_resources(): - return {qp.resource_rep(_UnknownOp): 1} - - @qp.register_resources(_rx_resources) - def _rx_decomp(_, wires): - _UnknownOp(wires=wires) - - qp.decomposition.enable_graph() - - @qp.qjit(capture=True) - @qp.decompose( - gate_set={"PauliX"}, - fixed_decomps={qp.RX: _rx_decomp}, - ) - @qp.qnode(qp.device("null.qubit", wires=1)) - def f(phi): - qp.RX(phi, 0) - return qp.state() - - try: - with pytest.raises( - ValueError, - match=r"Could not capture _UnknownOp without the number of wires\.", - ): - f(0.5) - finally: - qp.decomposition.disable_graph() - - def test_symbolic_controlled_op_is_skipped(self): - """Symbolic Controlled ops produced by ``qml.ctrl`` must be skipped when - iterating the decomposition graph solution. - """ - qp.decomposition.enable_graph() - - def _resources(): - return { - controlled_resource_rep( - qp.BasisEmbedding, {"num_wires": 1}, num_control_wires=1 - ): 1, - qp.resource_rep(qp.GlobalPhase): 1, - } - - @qp.register_resources(_resources) - def my_rz(phi, wires): - qp.GlobalPhase(phi / 2) - qp.ctrl(qp.BasisEmbedding, control=wires)([1], wires=[1]) - - @qp.qjit(capture=True) - @qp.decompose( - gate_set={"CNOT", "PauliX", "GlobalPhase", "MultiControlledX"}, - fixed_decomps={qp.RZ: my_rz}, - ) - @qp.qnode(qp.device("null.qubit", wires=2)) - def f(phi): - qp.RZ(phi, 0) - return qp.state() - - resources = qp.specs(f, level="device")(0.123).resources.gate_types - assert resources == {"BasisState": 1, "GlobalPhase": 1} - qp.decomposition.disable_graph() - - -if __name__ == "__main__": - pytest.main(["-x", __file__]) From 474eec9963915ab8d02e78b52e797c857d88c07a Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 23 Jul 2026 14:37:26 -0400 Subject: [PATCH 15/36] . --- frontend/test/pytest/test_QPD.py | 2 +- frontend/test/pytest/test_precompile_decomp_rules.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/test_QPD.py b/frontend/test/pytest/test_QPD.py index a4986d9d7d..3c73473740 100644 --- a/frontend/test/pytest/test_QPD.py +++ b/frontend/test/pytest/test_QPD.py @@ -17,7 +17,7 @@ import pennylane as qp import pytest -from catalyst.device.python_decompositions import python_decomposition_wrapper +from catalyst.decomposition.python_decompositions import python_decomposition_wrapper class TestQPD: diff --git a/frontend/test/pytest/test_precompile_decomp_rules.py b/frontend/test/pytest/test_precompile_decomp_rules.py index f2e78327d6..92e80c72fa 100644 --- a/frontend/test/pytest/test_precompile_decomp_rules.py +++ b/frontend/test/pytest/test_precompile_decomp_rules.py @@ -20,7 +20,7 @@ import pytest from catalyst.compiler import _quantum_opt -from catalyst.utils.precompile_decomposition_rules import ( +from catalyst.decomposition.precompile_decomposition_rules import ( compile_op_decomp_rules, get_abstract_args, precompile_decomp_rules, From 06d783611fbfed9cb194868ffab3f27d5aafb740 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 23 Jul 2026 14:41:10 -0400 Subject: [PATCH 16/36] add empty test file --- frontend/test/pytest/test_decomposition.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 frontend/test/pytest/test_decomposition.py diff --git a/frontend/test/pytest/test_decomposition.py b/frontend/test/pytest/test_decomposition.py new file mode 100644 index 0000000000..fb8fe117eb --- /dev/null +++ b/frontend/test/pytest/test_decomposition.py @@ -0,0 +1,13 @@ +# 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. From ac41517cf9e43f3c0f9e5270642b17ba6d5199a5 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 23 Jul 2026 14:46:15 -0400 Subject: [PATCH 17/36] new empty lit test file --- frontend/test/lit/_old_test_decomposition.py | 1863 ++++++++++++++++++ frontend/test/lit/test_decomposition.py | 1852 +---------------- 2 files changed, 1864 insertions(+), 1851 deletions(-) create mode 100644 frontend/test/lit/_old_test_decomposition.py diff --git a/frontend/test/lit/_old_test_decomposition.py b/frontend/test/lit/_old_test_decomposition.py new file mode 100644 index 0000000000..e8f9448e7c --- /dev/null +++ b/frontend/test/lit/_old_test_decomposition.py @@ -0,0 +1,1863 @@ +# Copyright 2022-2025 Xanadu Quantum Technologies Inc. +import os +import pathlib +import platform +from copy import deepcopy +from functools import partial + +import jax +import numpy as np +import pennylane as qp +from pennylane.devices.capabilities import OperatorProperties +from pennylane.typing import TensorLike +from pennylane.wires import WiresLike + +from catalyst import measure, qjit +from catalyst.compiler import get_lib_path +from catalyst.device import get_device_capabilities +from catalyst.jax_primitives import decomposition_rule +from catalyst.passes import graph_decomposition + +# 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. + +# RUN: %PYTHON %s | FileCheck %s +# pylint: disable=line-too-long +# pylint: disable=too-many-lines + + +# Helper to skip tests that fail due to PauliRot type annotation issue +# TODO: Remove this once PennyLane fixes the PauliRot decomposition type annotations +def skip_if_pauli_rot_issue(test_func): + """Wrapper to skip tests that fail due to PauliRot type annotation issues.""" + + def wrapper(): + try: + test_func() + except (ValueError, IndexError) as e: + error_msg = str(e) + if ( + "Unsupported type annotation None for parameter pauli_word" in error_msg + or "Unsupported type annotation for parameter pauli_word" in error_msg + or "index is out of bounds for axis" in error_msg + ): + print(f"# SKIPPED {test_func.__name__}: PauliRot type annotation issue") + else: + raise + + return wrapper + + +TEST_PATH = os.path.dirname(__file__) +CONFIG_CUSTOM_DEVICE = pathlib.Path(f"{TEST_PATH}/../custom_device/custom_device.toml") + + +def get_custom_device_without(num_wires, discards=frozenset(), force_matrix=frozenset()): + """Generate a custom device without gates in discards.""" + + class CustomDevice(qp.devices.Device): + """Custom Gate Set Device""" + + name = "Custom Device" + config_filepath = CONFIG_CUSTOM_DEVICE + + _to_matrix_ops = {} + + def __init__(self, wires=None): + super().__init__(wires=wires) + self.qjit_capabilities = deepcopy(get_device_capabilities(self)) + for gate in discards: + self.qjit_capabilities.operations.pop(gate, None) + for gate in force_matrix: + self.qjit_capabilities.operations.pop(gate, None) + self._to_matrix_ops[gate] = OperatorProperties(False, False, False) + + def apply(self, operations, **kwargs): + """Unused""" + raise RuntimeError("Only C/C++ interface is defined") + + @staticmethod + def get_c_interface(): + """Returns a tuple consisting of the device name, and + the location to the shared object with the C/C++ device implementation. + """ + system_extension = ".dylib" if platform.system() == "Darwin" else ".so" + lib_path = ( + get_lib_path("runtime", "RUNTIME_LIB_DIR") + "/librtd_null_qubit" + system_extension + ) + return "NullQubit", lib_path + + def execute(self, circuits, execution_config): + """Execution.""" + return circuits, execution_config + + return CustomDevice(wires=num_wires) + + +def test_decompose_multicontrolledx(): + """Test decomposition of MultiControlledX as an aliased gate.""" + dev = get_custom_device_without(5, discards={"MultiControlledX"}) + + @qjit(target="mlir") + @qp.qnode(dev) + # CHECK-LABEL: @jit_decompose_multicontrolled_x1 + def decompose_multicontrolled_x1(theta: float): + qp.RX(theta, wires=[0]) + # CHECK-NOT: name = "MultiControlledX" + # CHECK: quantum.custom "PauliX"() {{%[a-zA-Z0-9_]+}} ctrls({{%[a-zA-Z0-9_]+}}, {{%[a-zA-Z0-9_]+}}, {{%[a-zA-Z0-9_]+}}) + # CHECK-NOT: name = "MultiControlledX" + qp.MultiControlledX(wires=[0, 1, 2, 3]) + return qp.state() + + print(decompose_multicontrolled_x1.mlir) + + +test_decompose_multicontrolledx() + + +def test_decompose_rot(): + """Test decomposition of Rot gate.""" + dev = get_custom_device_without(1, discards={"Rot", "C(Rot)"}) + + @qjit(target="mlir") + @qp.qnode(dev) + # CHECK-LABEL: @jit_decompose_rot + def decompose_rot(phi: float, theta: float, omega: float): + # CHECK-NOT: name = "Rot" + # CHECK: [[phi:%.+]] = tensor.extract %arg0 + # CHECK-NOT: name = "Rot" + # CHECK: {{%.+}} = quantum.custom "RZ"([[phi]]) + # CHECK-NOT: name = "Rot" + # CHECK: [[theta:%.+]] = tensor.extract %arg1 + # CHECK-NOT: name = "Rot" + # CHECK: {{%.+}} = quantum.custom "RY"([[theta]]) + # CHECK-NOT: name = "Rot" + # CHECK: [[omega:%.+]] = tensor.extract %arg2 + # CHECK-NOT: name = "Rot" + # CHECK: {{%.+}} = quantum.custom "RZ"([[omega]]) + # CHECK-NOT: name = "Rot" + qp.Rot(phi, theta, omega, wires=0) + return measure(wires=0) + + print(decompose_rot.mlir) + + +test_decompose_rot() + + +def test_decompose_s(): + """Test decomposition of S gate.""" + dev = get_custom_device_without(1, discards={"S", "C(S)"}) + + @qjit(target="mlir") + @qp.qnode(dev) + # CHECK-LABEL: @jit_decompose_s + def decompose_s(): + # CHECK-NOT: name="S" + # CHECK: [[pi_div_2:%.+]] = arith.constant 1.57079{{.+}} : f64 + # CHECK-NOT: name = "S" + # CHECK: {{%.+}} = quantum.custom "PhaseShift"([[pi_div_2]]) + # CHECK-NOT: name = "S" + qp.S(wires=0) + return measure(wires=0) + + print(decompose_s.mlir) + + +test_decompose_s() + + +def test_decompose_qubitunitary(): + """Test decomposition of QubitUnitary""" + dev = get_custom_device_without(1, discards={"QubitUnitary"}) + + @qjit(target="mlir") + @qp.qnode(dev) + # CHECK-LABEL: @jit_decompose_qubit_unitary + def decompose_qubit_unitary(U: jax.core.ShapedArray([2, 2], float)): + # CHECK-NOT: name = "QubitUnitary" + # CHECK: quantum.custom "RZ" + # CHECK: quantum.custom "RY" + # CHECK: quantum.custom "RZ" + # CHECK-NOT: name = "QubitUnitary" + qp.QubitUnitary(U, wires=0) + return measure(wires=0) + + print(decompose_qubit_unitary.mlir) + + +test_decompose_qubitunitary() + + +def test_decompose_singleexcitation(): + """ + Test that single excitation is not decomposed. + """ + dev = get_custom_device_without(2) + + @qjit(target="mlir") + @qp.qnode(dev) + # CHECK-LABEL: @jit_decompose_singleexcitation + def decompose_singleexcitation(theta: float): + # CHECK: quantum.custom "SingleExcitation" + + qp.SingleExcitation(theta, wires=[0, 1]) + return measure(wires=0) + + print(decompose_singleexcitation.mlir) + + +test_decompose_singleexcitation() + + +def test_decompose_doubleexcitation(): + """ + Test that Double excitation is not decomposed. + """ + dev = get_custom_device_without(4) + + @qjit(target="mlir") + @qp.qnode(dev) + # CHECK-LABEL: @jit_decompose_doubleexcitation + def decompose_doubleexcitation(theta: float): + # CHECK: quantum.custom "DoubleExcitation" + + qp.DoubleExcitation(theta, wires=[0, 1, 2, 3]) + return measure(wires=0) + + print(decompose_doubleexcitation.mlir) + + +test_decompose_doubleexcitation() + + +def test_decompose_singleexcitationplus(): + """ + Test decomposition of single excitation plus. + See + https://github.com/PennyLaneAI/pennylane/blob/main/pennylane/ops/qubit/qchem_ops.py + for the decomposition of qp.SingleExcitationPlus + """ + dev = get_custom_device_without(2, discards={"SingleExcitationPlus", "C(SingleExcitationPlus)"}) + + @qjit(target="mlir") + @qp.qnode(dev) + # CHECK-LABEL: @jit_decompose_singleexcitationplus + def decompose_singleexcitationplus(theta: float): + # CHECK-NOT: "SingleExcitationPlus" + # CHECK: quantum.custom "Hadamard" + # CHECK: quantum.custom "CNOT" + # CHECK: quantum.custom "RY" + # CHECK: quantum.custom "RY" + # CHECK: quantum.custom "CY" + # CHECK: quantum.custom "S" + # CHECK: quantum.custom "Hadamard" + # CHECK: quantum.custom "RZ" + # CHECK: quantum.custom "CNOT" + # CHECK: quantum.gphase + + qp.SingleExcitationPlus(theta, wires=[0, 1]) + return measure(wires=0) + + print(decompose_singleexcitationplus.mlir) + + +test_decompose_singleexcitationplus() + + +def test_decompose_to_matrix(): + """Test decomposition of QubitUnitary""" + dev = get_custom_device_without(1, force_matrix={"PauliY"}) + + @qjit(target="mlir") + @qp.qnode(dev) + # CHECK-LABEL: @jit_decompose_to_matrix + def decompose_to_matrix(): + # CHECK: quantum.custom "PauliX" + qp.PauliX(wires=0) + # CHECK: quantum.unitary + qp.PauliY(wires=0) + # CHECK: quantum.custom "PauliZ" + qp.PauliZ(wires=0) + return measure(wires=0) + + print(decompose_to_matrix.mlir) + + +test_decompose_to_matrix() + + +def test_decomposition_rule_lowering(): + """Test that decomposition rules are lowered to private functions.""" + + @decomposition_rule(is_qreg=True) + def my_decomp(): + return + + @qp.qjit(capture=True) + @qp.qnode(qp.device("null.qubit", wires=1)) + def circuit(): + # CHECK-LABEL: func.func private @my_decomp + my_decomp() + return + + print(circuit.mlir) + + +test_decomposition_rule_lowering() + + +def test_decomposition_rule_wire_param(): + """Test decomposition rule with passing a parameter that is a wire/integer""" + + @decomposition_rule(is_qreg=False) + def Hadamard0(wire: WiresLike): + qp.Hadamard(wire) + + @qp.qjit(capture=True) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK-LABEL: @circuit + def circuit(_: float): + # CHECK: @circuit([[ARG0:%.+]] + # CHECK: [[QREG:%.+]] = qref.alloc + Hadamard0(int) + return qp.probs() + + # CHECK: @Hadamard0([[QBIT:%.+]]: !qref.bit) + # CHECK-NEXT: qref.custom "Hadamard"() [[QBIT]] : !qref.bit + # CHECK-NEXT: return + + print(circuit.mlir) + + +test_decomposition_rule_wire_param() + + +def test_decomposition_rule_gate_param_param(): + """Test decomposition rule with passing a regular parameter""" + + @decomposition_rule(is_qreg=False, num_params=1) + def RX_on_wire_0(param: TensorLike, w0: WiresLike): + qp.RX(param, wires=w0) + + @qp.qjit(capture=True) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK: module @circuit_2 + def circuit_2(_: float): + RX_on_wire_0(float, int) + return qp.probs() + + # CHECK: @RX_on_wire_0([[PARAM_TENSOR:%.+]]: tensor, [[QUBIT:%.+]]: !qref.bit) + # CHECK-NEXT: [[PARAM:%.+]] = tensor.extract [[PARAM_TENSOR]][] : tensor + # CHECK-NEXT: qref.custom "RX"([[PARAM]]) [[QUBIT]] : !qref.bit + # CHECK-NEXT: return + print(circuit_2.mlir) + + +test_decomposition_rule_gate_param_param() + + +def test_multiple_decomposition_rules(): + """Test with multiple decomposition rules""" + + @decomposition_rule + def identity(): ... + + @decomposition_rule(is_qreg=True) + def all_wires_rx(param: TensorLike, w0: WiresLike, w1: WiresLike, w2: WiresLike): + qp.RX(param, wires=w0) + qp.RX(param, wires=w1) + qp.RX(param, wires=w2) + + @qp.qjit(capture=True) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + def circuit_3(_: float): + # CHECK: [[QREG:%.+]] = qref.alloc + # CHECK-NEXT: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit + # CHECK-NEXT: qref.custom "Hadamard"() [[QUBIT]] : !qref.bit + # CHECK-NEXT: qref.compbasis(qreg [[QREG]] : !qref.reg<1>) : !quantum.obs + identity() + all_wires_rx(float, int, int, int) + qp.Hadamard(0) + return qp.probs() + + # CHECK-LABEL: @identity + # CHECK-LABEL: @all_wires_rx + + print(circuit_3.mlir) + + +test_multiple_decomposition_rules() + + +def test_decomposition_rule_shaped_wires(): + """Test decomposition rule with passing a shaped array of wires""" + + @decomposition_rule(is_qreg=True) + def shaped_wires_rule(param: TensorLike, wires: WiresLike): + qp.RX(param, wires=wires[0]) + qp.RX(param, wires=wires[1]) + qp.RX(param, wires=wires[2]) + + @qp.qjit(capture=True) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + def circuit_4(_: float): + # CHECK: module @circuit_4 + shaped_wires_rule(float, jax.core.ShapedArray((3,), int)) + qp.Hadamard(0) + return qp.probs() + + # CHECK: @shaped_wires_rule([[QREG:%.+]]: !qref.reg<1>, [[PARAM_TENSOR:%.+]]: tensor, [[QUBITS:%.+]]: tensor<3xi64>) + # CHECK-NEXT: [[IDX_0:%.+]] = stablehlo.slice [[QUBITS]] [0:1] : (tensor<3xi64>) -> tensor<1xi64> + # CHECK-NEXT: [[RIDX_0:%.+]] = stablehlo.reshape [[IDX_0]] : (tensor<1xi64>) -> tensor + # CHECK-NEXT: [[EXTRACTED:%.+]] = tensor.extract [[RIDX_0]][] : tensor + # CHECK-NEXT: [[QUBIT:%.+]] = qref.get [[QREG]][[[EXTRACTED]]] : !qref.reg<1>, i64 -> !qref.bit + # CHECK-NEXT: [[EXTRACTED_0:%.+]] = tensor.extract [[PARAM_TENSOR]][] : tensor + # CHECK-NEXT: qref.custom "RX"([[EXTRACTED_0]]) [[QUBIT]] : !qref.bit + + print(circuit_4.mlir) + + +test_decomposition_rule_shaped_wires() + + +def test_decomposition_rule_expanded_wires(): + """Test decomposition rule with passing expanding wires as a Python list""" + + def shaped_wires_rule(param: TensorLike, wires: WiresLike): + qp.RX(param, wires=wires[0]) + qp.RX(param, wires=wires[1]) + qp.RX(param, wires=wires[2]) + + @decomposition_rule(is_qreg=False, num_params=1) + def expanded_wires_rule(param: TensorLike, w1, w2, w3): + shaped_wires_rule(param, [w1, w2, w3]) + + @qp.qjit(capture=True) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + def circuit_5(_: float): + # CHECK: module @circuit_5 + expanded_wires_rule(float, int, int, int) + qp.Hadamard(0) + return qp.probs() + + # CHECK-LABEL: @expanded_wires_rule(%arg0: tensor, %arg1: !qref.bit, %arg2: !qref.bit, %arg3: !qref.bit) + + print(circuit_5.mlir) + + +test_decomposition_rule_expanded_wires() + + +def test_decomposition_rule_with_cond(): + """Test decomposition rule with a conditional path""" + + @decomposition_rule(is_qreg=True) + def cond_RX(param: TensorLike, w0: WiresLike): + + def true_path(): + qp.RX(param, wires=w0) + + def false_path(): ... + + qp.cond(param != 0.0, true_path, false_path)() + + @qp.qjit(autograph=False, capture=True) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + def circuit_6(): + # CHECK: module @circuit_6 + cond_RX(float, jax.core.ShapedArray((1,), int)) + return qp.probs() + + # CHECK: @cond_RX([[QREG:%.+]]: !qref.reg<1>, [[PARAM_TENSOR:%.+]]: tensor, [[QUBITS:%.+]]: tensor<1xi64>) + # CHECK-NEXT: [[ZERO:%.+]] = stablehlo.constant dense<0.000000e+00> : tensor + # CHECK-NEXT: [[COND_TENSOR:%.+]] = stablehlo.compare NE, [[PARAM_TENSOR]], [[ZERO]], FLOAT : (tensor, tensor) -> tensor + # CHECK-NEXT: [[COND:%.+]] = tensor.extract [[COND_TENSOR]][] : tensor + # CHECK-NEXT: scf.if [[COND]] + # CHECK-DAG: [[QUBIT:%.+]] = qref.get [[QREG]][%extracted_0] : !qref.reg<1>, i64 -> !qref.bit + # CHECK-DAG: [[PARAM:%.+]] = tensor.extract [[PARAM_TENSOR]][] : tensor + # CHECK: qref.custom "RX"([[PARAM]]) [[QUBIT]] : !qref.bit + # CHECK: return + + print(circuit_6.mlir) + + +test_decomposition_rule_with_cond() + + +def test_decomposition_rule_caller(): + """Test decomposition rules with a caller""" + + @decomposition_rule(is_qreg=True) + def rule_op1_decomp(_: TensorLike, wires: WiresLike): + qp.Hadamard(wires=wires[0]) + qp.Hadamard(wires=[1]) + + @decomposition_rule(is_qreg=True) + def rule_op2_decomp(param: TensorLike, wires: WiresLike): + qp.RX(param, wires=wires[0]) + + def decomps_caller(param: TensorLike, wires: WiresLike): + rule_op1_decomp(param, wires) + rule_op2_decomp(param, wires) + + @qp.qjit(autograph=False, capture=True) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK: module @circuit_7 + def circuit_7(): + # CHECK: [[QREG:%.+]] = qref.alloc + # CHECK: qref.compbasis(qreg [[QREG]] : !qref.reg<1>) : !quantum.obs + decomps_caller(float, jax.core.ShapedArray((2,), int)) + return qp.probs() + + # CHECK-LABEL: @rule_op1_decomp(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<2xi64>) + # CHECK-LABEL: @rule_op2_decomp(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<2xi64>) + print(circuit_7.mlir) + + +test_decomposition_rule_caller() + + +def test_decompose_gateset_without_graph(): + """Test the decompose transform to a target gate set without the graph decomposition.""" + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={"RX", "RZ"}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK-LABEL: @circuit_8() -> tensor attributes {diff_method = "adjoint", llvm.linkage = #llvm.linkage, quantum.node} + def circuit_8(): + return qp.expval(qp.Z(0)) + + print(circuit_8.mlir) + + +test_decompose_gateset_without_graph() + + +def test_decompose_gateset_with_graph(): + """Test the decompose transform to a target gate set with the graph decomposition.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={"RX"}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK-LABEL: @simple_circuit_9() -> tensor attributes {decompose_gatesets + def simple_circuit_9(): + return qp.expval(qp.Z(0)) + + print(simple_circuit_9.mlir) + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={"RX", "RZ"}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" + # CHECK-LABEL: @circuit_9() -> tensor attributes {decompose_gatesets + def circuit_9(): + return qp.expval(qp.Z(0)) + + print(circuit_9.mlir) + + qp.decomposition.disable_graph() + + +test_decompose_gateset_with_graph() + + +def test_decompose_gateset_operator_with_graph(): + """Test the decompose transform to a target gate set with the graph decomposition.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={qp.RX}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK-LABEL: @simple_circuit_10() -> tensor attributes {decompose_gatesets + def simple_circuit_10(): + return qp.expval(qp.Z(0)) + + print(simple_circuit_10.mlir) + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={qp.RX, qp.RZ, "PauliZ", qp.PauliX, qp.Hadamard}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK-LABEL: @circuit_10() -> tensor attributes {decompose_gatesets + def circuit_10(): + return qp.expval(qp.Z(0)) + + print(circuit_10.mlir) + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={qp.RX, qp.RZ, qp.PauliZ, qp.PauliX, qp.Hadamard}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" + # CHECK-LABEL: @circuit_11() -> tensor attributes {decompose_gatesets + def circuit_11(): + return qp.expval(qp.Z(0)) + + print(circuit_11.mlir) + + qp.decomposition.disable_graph() + + +test_decompose_gateset_operator_with_graph() + + +def test_decompose_gateset_with_rotxzx(): + """Test the decompose transform with a custom operator with the graph decomposition.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={"RotXZX"}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK-LABEL: @simple_circuit_12() -> tensor attributes {decompose_gatesets + def simple_circuit_12(): + return qp.expval(qp.Z(0)) + + print(simple_circuit_12.mlir) + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={qp.ftqc.RotXZX}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" + # CHECK-LABEL: @circuit_12() -> tensor attributes {decompose_gatesets + def circuit_12(): + return qp.expval(qp.Z(0)) + + print(circuit_12.mlir) + + qp.decomposition.disable_graph() + + +test_decompose_gateset_with_rotxzx() + + +def test_decomposition_rule_name(): + """Test the name of the decomposition rule is not updated with circuit instantiation.""" + + qp.decomposition.enable_graph() + + @decomposition_rule + def _ry_to_rz_rx(phi, wires: WiresLike, **__): + """Decomposition of RY gate using RZ and RX gates.""" + qp.RZ(-np.pi / 2, wires=wires) + qp.RX(phi, wires=wires) + qp.RZ(np.pi / 2, wires=wires) + + @decomposition_rule + def _rot_to_rz_ry_rz(phi, theta, omega, wires: WiresLike, **__): + """Decomposition of Rot gate using RZ and RY gates.""" + qp.RZ(phi, wires=wires) + qp.RY(theta, wires=wires) + qp.RZ(omega, wires=wires) + + @decomposition_rule + def _u2_phaseshift_rot_decomposition(phi, delta, wires, **__): + """Decomposition of U2 gate using Rot and PhaseShift gates.""" + pi_half = qp.math.ones_like(delta) * (np.pi / 2) + qp.Rot(delta, pi_half, -delta, wires=wires) + qp.PhaseShift(delta, wires=wires) + qp.PhaseShift(phi, wires=wires) + + @decomposition_rule + def _xzx_decompose(phi, theta, omega, wires, **__): + """Decomposition of Rot gate using RX and RZ gates in XZX format.""" + qp.RX(phi, wires=wires) + qp.RZ(theta, wires=wires) + qp.RX(omega, wires=wires) + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={"RX", "RZ", "PhaseShift"}) + @qp.qnode(qp.device("lightning.qubit", wires=3)) + # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" + # CHECK-LABEL: @circuit_13() -> tensor attributes {decompose_gatesets + def circuit_13(): + _ry_to_rz_rx(float, int) + _rot_to_rz_ry_rz(float, float, float, int) + _u2_phaseshift_rot_decomposition(float, float, int) + _xzx_decompose(float, float, float, int) + return qp.expval(qp.Z(0)) + + # CHECK-LABEL: @_ry_to_rz_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor) + # CHECK-LABEL: @_rot_to_rz_ry_rz(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor) + # CHECK-LABEL: @_u2_phaseshift_rot_decomposition(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor) + # CHECK-LABEL: @_xzx_decompose(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor) + print(circuit_13.mlir) + + qp.decomposition.disable_graph() + + +test_decomposition_rule_name() + + +def test_decomposition_rule_name_update(): + """Test the name of the decomposition rule is updated in the MLIR output.""" + + qp.decomposition.enable_graph() + + @qp.register_resources({qp.RZ: 2, qp.RX: 1}) + def rz_rx(phi, wires: WiresLike, **__): + """Decomposition of RY gate using RZ and RX gates.""" + qp.RZ(-np.pi / 2, wires=wires) + qp.RX(phi, wires=wires) + qp.RZ(np.pi / 2, wires=wires) + + @qp.register_resources({qp.RZ: 2, qp.RY: 1}) + def rz_ry_rz(phi, theta, omega, wires: WiresLike, **__): + """Decomposition of Rot gate using RZ and RY gates.""" + qp.RZ(phi, wires=wires) + qp.RY(theta, wires=wires) + qp.RZ(omega, wires=wires) + + @qp.register_resources({qp.RY: 1, qp.GlobalPhase: 1}) + def ry_gp(wires: WiresLike, **__): + """Decomposition of PauliY gate using RY and GlobalPhase gates.""" + qp.RY(np.pi, wires=wires) + qp.GlobalPhase(-np.pi / 2, wires=wires) + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RX", "RZ", "GlobalPhase"}, + fixed_decomps={ + qp.RY: rz_rx, + qp.Rot: rz_ry_rz, + qp.PauliY: ry_gp, + }, + ) + @qp.qnode(qp.device("lightning.qubit", wires=3)) + # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" + # CHECK-LABEL: @circuit_14() -> tensor attributes {decompose_gatesets + def circuit_14(): + qp.RY(0.5, wires=0) + qp.Rot(0.1, 0.2, 0.3, wires=1) + qp.PauliY(wires=2) + return qp.expval(qp.Z(0)) + + # CHECK-DAG: @rz_ry_rz(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) + # CHECK-DAG: @rz_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) + # CHECK-DAG: @ry_gp(%arg0: !qref.reg<3>, %arg1: tensor<1xi64>) + print(circuit_14.mlir) + + qp.decomposition.disable_graph() + + +test_decomposition_rule_name_update() + + +def test_decomposition_inside_subroutine(): + """Test that operators inside subroutines can be decomposed.""" + + qp.decomposition.enable_graph() + + @qp.templates.Subroutine + def f(x, wires): + qp.IsingXX(x, wires) + + @qp.qjit(capture=True, target="mlir") + @qp.decompose(gate_set=qp.gate_sets.ROTATIONS_PLUS_CNOT) + @qp.qnode(qp.device("lightning.qubit", wires=5)) + # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" + def subroutine_circuit(): + # CHECK-DAG: [[FIRST_CONST:%.+]] = stablehlo.constant dense<5.000000e-01> : tensor + # CHECK-DAG: [[SECOND_CONST:%.+]] = stablehlo.constant dense<1.200000e+00> : tensor + + # CHECK: [[QREG:%.+]] = qref.alloc + # CHECK: call @f([[QREG]], [[FIRST_CONST]], {{%.+}}) : (!qref.reg<5>, tensor, tensor<2xi64>) + # CHECK: call @f([[QREG]], [[SECOND_CONST]], {{%.+}}) : (!qref.reg<5>, tensor, tensor<2xi64>) + + f(0.5, (0, 1)) + f(1.2, (2, 3)) + return qp.probs(wires=0) + + # CHECK-DAG: @_isingxx_to_cnot_rx_cnot(%arg0: !qref.reg<5>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) + print(subroutine_circuit.mlir) + qp.decomposition.disable_graph() + + +test_decomposition_inside_subroutine() + + +def test_decomposition_rule_name_update_multi_qubits(): + """Test the name of the decomposition rule with multi-qubit gates.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RY", "RX", "CNOT", "Hadamard", "GlobalPhase"}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=4)) + # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" + # CHECK-LABEL: @circuit_15() -> tensor attributes {decompose_gatesets + def circuit_15(): + qp.SingleExcitation(0.5, wires=[0, 1]) + qp.SingleExcitationPlus(0.5, wires=[0, 1]) + qp.SingleExcitationMinus(0.5, wires=[0, 1]) + qp.DoubleExcitation(0.5, wires=[0, 1, 2, 3]) + return qp.expval(qp.Z(0)) + + # CHECK-DAG: @_cry(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CRY"} + # CHECK-DAG: @_s_phaseshift(%arg0: !qref.reg<4>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "S"} + # CHECK-DAG: @_phaseshift_to_rz_gp(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PhaseShift"} + # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} + # CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} + # CHECK-DAG: @_doublexcit(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<4xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 4 : i64, target_gate = "DoubleExcitation"} + # CHECK-DAG: @_single_excitation_decomp(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "SingleExcitation"} + print(circuit_15.mlir) + + qp.decomposition.disable_graph() + + +skip_if_pauli_rot_issue(test_decomposition_rule_name_update_multi_qubits)() + + +def test_decomposition_rule_name_adjoint(): + """Test decomposition rule with qp.adjoint.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RY", "RX", "CZ", "GlobalPhase", "Adjoint(SingleExcitation)"}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=4)) + # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" + def circuit_16(x: float): + # CHECK: qref.adjoint { + # CHECK: qref.adjoint { + # CHECK: qref.adjoint { + # CHECK: qref.adjoint { + qp.adjoint(qp.CNOT)(wires=[0, 1]) + qp.adjoint(qp.Hadamard)(wires=2) + qp.adjoint(qp.RZ)(0.5, wires=3) + qp.adjoint(qp.SingleExcitation)(0.1, wires=[0, 1]) + qp.adjoint(qp.SingleExcitation(x, wires=[0, 1])) + return qp.expval(qp.Z(0)) + + # CHECK-DAG: @_single_excitation_decomp(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "SingleExcitation"} + # CHECK-DAG: @_hadamard_to_rz_ry(%arg0: !qref.reg<4>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Hadamard"} + # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} + # CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} + # CHECK-DAG: @_cnot_to_cz_h(%arg0: !qref.reg<4>, %arg1: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CNOT"} + print(circuit_16.mlir) + + qp.decomposition.disable_graph() + + +skip_if_pauli_rot_issue(test_decomposition_rule_name_adjoint)() + + +# TODO: Reenable this once the underlying non-determinism issue is resolved +def test_decomposition_rule_name_ctrl(): + """Test decomposition rule with qp.ctrl.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RX", "RZ", "H", "CZ"}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=2)) + # SKIP-CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" + # SKIP-CHECK{LITERAL}: @circuit_17() -> tensor attributes {decompose_gatesets + def circuit_17(): + # SKIP-CHECK: %out_qubits:2 = quantum.custom "CRY"(%cst) %1, %2 : !quantum.bit, !quantum.bit + # SKIP-CHECK-NEXT: %out_qubits_0:2 = quantum.custom "CNOT"() %out_qubits#0, %out_qubits#1 : !quantum.bit, !quantum.bit + qp.ctrl(qp.RY, control=0)(0.5, 1) + qp.ctrl(qp.PauliX, control=0)(1) + return qp.expval(qp.Z(0)) + + # SKIP-CHECK-DAG: @_cnot_to_cz_h(%arg0: !quantum.reg, %arg1: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CNOT"} + # SKIP-CHECK-DAG: @_cry(%arg0: !quantum.reg, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CRY"} + # SKIP-CHECK-DAG: @_ry_to_rz_rx(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RY"} + # SKIP-CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} + # print(circuit_17.mlir) + + qp.decomposition.disable_graph() + + +skip_if_pauli_rot_issue(test_decomposition_rule_name_ctrl)() + + +# TODO: Reenable this once the underlying non-determinism issue is resolved +def test_qft_decomposition(): + """Test the decomposition of the QFT""" + + qp.decomposition.enable_graph() + + @qp.qjit(autograph=True, target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RX", "RY", "CNOT", "GlobalPhase"}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=4)) + # SKIP-CHECK: %0 = transform.apply_registered_pass "decompose-lowering" + # SKIP-CHECK: @circuit_18(%arg0: tensor<3xf64>) -> tensor attributes {decompose_gatesets + def circuit_18(): + # %6 = scf.for %arg1 = %c0 to %c4 step %c1 iter_args(%arg2 = %0) -> (!quantum.reg) { + # %23 = scf.for %arg3 = %c0 to %22 step %c1 iter_args(%arg4 = %21) -> (!quantum.reg) { + # %7 = scf.for %arg1 = %c0 to %c2 step %c1 iter_args(%arg2 = %6) -> (!quantum.reg) { + qp.QFT(wires=[0, 1, 2, 3]) + return qp.expval(qp.Z(0)) + + # SKIP-CHECK-DAG: @ag___cphase_to_rz_cnot(%arg0: !quantum.reg, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "ControlledPhaseShift"} + # SKIP-CHECK-DAG: @ag___rz_to_ry_rx(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} + # SKIP-CHECK-DAG: @ag___rot_to_rz_ry_rz(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} + # SKIP-CHECK-DAG: @ag___swap_to_cnot(%arg0: !quantum.reg, %arg1: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "SWAP"} + # SKIP-CHECK-DAG: @ag___hadamard_to_rz_ry(%arg0: !quantum.reg, %arg1: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Hadamard"} + # print(circuit_18.mlir) + + qp.decomposition.disable_graph() + + +skip_if_pauli_rot_issue(test_qft_decomposition)() + + +def test_decompose_lowering_with_other_passes(): + """Test the decompose lowering pass with other passes in a pass pipeline.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @qp.transforms.merge_rotations + @qp.transforms.cancel_inverses + @partial( + qp.transforms.decompose, + gate_set={"RZ", "RY", "CNOT", "GlobalPhase"}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK: module attributes {transform.with_named_sequence} { + # CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.op<"builtin.module">) { + # CHECK-NEXT: [[ONE:%.+]] = transform.apply_registered_pass "decompose-lowering" to %arg0 : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> + # CHECK-NEXT: [[TWO:%.+]] = transform.apply_registered_pass "cancel-inverses" to [[ONE]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> + # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to [[TWO]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> + # CHECK-NEXT: transform.yield + # CHECK-NEXT: } + def circuit_19(): + + # CHECK: [[QREG:%.+]] = qref.alloc( 1) : !qref.reg<1> + # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit + # CHECK: qref.custom "PauliX"() [[QUBIT]] : !qref.bit + # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit + # CHECK: qref.custom "PauliX"() [[QUBIT]] : !qref.bit + # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit + # CHECK: qref.custom "RX"({{%.+}}) [[QUBIT]] : !qref.bit + # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit + # CHECK: qref.custom "RX"({{%.+}}) [[QUBIT]] : !qref.bit + qp.PauliX(0) + qp.PauliX(0) + qp.RX(0.1, wires=0) + qp.RX(-0.1, wires=0) + return qp.expval(qp.PauliX(0)) + + # CHECK-DAG: @_paulix_to_rx(%arg0: !qref.reg<1>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PauliX"} + # CHECK-DAG: @_rx_to_rz_ry(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RX"} + print(circuit_19.mlir) + + qp.decomposition.disable_graph() + + +skip_if_pauli_rot_issue(test_decompose_lowering_with_other_passes)() + + +def test_decompose_lowering_multirz(): + """Test the decompose lowering pass with MultiRZ in the gate set.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"CNOT", "RZ"}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=3)) + # CHECK: %0 = transform.apply_registered_pass "decompose-lowering" + def circuit_20(x: float): + # CHECK: [[QREG:%.+]] = qref.alloc( 3) : !qref.reg<3> + # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit + # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor + # CHECK: qref.multirz([[angle]]) [[q0]] : !qref.bit + # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit + # CHECK: [[q1:%.+]] = qref.get [[QREG]][ 1] : !qref.reg<3> -> !qref.bit + # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor + # CHECK: qref.multirz([[angle]]) [[q0]], [[q1]] : !qref.bit, !qref.bit + # CHECK: [[q1:%.+]] = qref.get [[QREG]][ 1] : !qref.reg<3> -> !qref.bit + # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit + # CHECK: [[q2:%.+]] = qref.get [[QREG]][ 2] : !qref.reg<3> -> !qref.bit + # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor + # CHECK: qref.multirz([[angle]]) [[q1]], [[q0]], [[q2]] : !qref.bit, !qref.bit, !qref.bit + qp.MultiRZ(x, wires=[0]) + qp.MultiRZ(x, wires=[0, 1]) + qp.MultiRZ(x, wires=[1, 0, 2]) + return qp.expval(qp.PauliX(0)) + + # CHECK-DAG: @_multi_rz_decomposition_wires_1(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "MultiRZ"} + # CHECK-DAG: @_multi_rz_decomposition_wires_2(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "MultiRZ"} + # CHECK-DAG: @_multi_rz_decomposition_wires_3(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<3xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 3 : i64, target_gate = "MultiRZ"} + # CHECK-DAG: scf.for %arg3 = %c0 to %c2 step %c1 + # CHECK-DAG: scf.for %arg3 = %c1 to %c3 step %c1 + print(circuit_20.mlir) + + qp.decomposition.disable_graph() + + +test_decompose_lowering_multirz() + + +def test_decompose_lowering_with_ordered_passes(): + """Test the decompose lowering pass with other passes in a specific order in a pass pipeline.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RZ", "RY", "CNOT", "GlobalPhase"}, + ) + @qp.transforms.merge_rotations + @qp.transforms.cancel_inverses + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK: module attributes {transform.with_named_sequence} { + # CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.op<"builtin.module">) { + # CHECK-NEXT: [[FIRST:%.+]] = transform.apply_registered_pass "cancel-inverses" to %arg0 : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> + # CHECK-NEXT: [[SECOND:%.+]] = transform.apply_registered_pass "merge-rotations" to [[FIRST]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> + # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "decompose-lowering" to [[SECOND]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> + # CHECK-NEXT: transform.yield + # CHECK-NEXT: } + def circuit_21(x: float): + # CHECK: [[QREG:%.+]] = qref.alloc( 1) : !qref.reg<1> + # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit + # CHECK: qref.custom "PauliX"() [[q0]] : !qref.bit + # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit + # CHECK: qref.custom "PauliX"() [[q0]] : !qref.bit + # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit + # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor + # CHECK: qref.custom "RX"([[angle]]) [[q0]] : !qref.bit + # CHECK: [[negated:%.+]] = stablehlo.negate %arg0 : tensor + # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit + # CHECK: [[neg_angle:%.+]] = tensor.extract [[negated]][] : tensor + # CHECK: qref.custom "RX"([[neg_angle]]) [[q0]] : !qref.bit + qp.PauliX(0) + qp.PauliX(0) + qp.RX(x, wires=0) + qp.RX(-x, wires=0) + return qp.expval(qp.PauliX(0)) + + # CHECK-DAG: @_paulix_to_rx(%arg0: !qref.reg<1>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PauliX"} + # CHECK-DAG: @_rx_to_rz_ry(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RX"} + # CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} + print(circuit_21.mlir) + + qp.decomposition.disable_graph() + + +skip_if_pauli_rot_issue(test_decompose_lowering_with_ordered_passes)() + + +def test_decompose_lowering_with_gphase(): + """Test the decompose lowering pass with GlobalPhase.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RX", "RY", "GlobalPhase"}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=3)) + # CHECK: %0 = transform.apply_registered_pass "decompose-lowering" + def circuit_22(): + # CHECK: [[QREG:%.+]] = qref.alloc( 3) : !qref.reg<3> + # CHECK: qref.gphase({{%.+}}) + # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit + # CHECK: qref.custom "PhaseShift"({{%.+}}) [[q0]] : !qref.bit + # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit + # CHECK: qref.custom "PhaseShift"({{%.+}}) [[q0]] : !qref.bit + + qp.GlobalPhase(0.5) + qp.ctrl(qp.GlobalPhase, control=0)(0.3) + qp.ctrl(qp.GlobalPhase, control=0)(phi=0.3, wires=[1, 2]) + return qp.expval(qp.PauliX(0)) + + # CHECK-DAG: @_phaseshift_to_rz_gp(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PhaseShift"} + # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} + print(circuit_22.mlir) + + qp.decomposition.disable_graph() + + +skip_if_pauli_rot_issue(test_decompose_lowering_with_gphase)() + + +def test_decompose_lowering_alt_decomps(): + """Test the decompose lowering pass with alternative decompositions.""" + + qp.decomposition.enable_graph() + + @qp.register_resources({qp.RY: 1}) + def custom_rot_cheap(params, wires: WiresLike): + qp.RY(params[1], wires=wires) + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RY", "RZ"}, + alt_decomps={qp.Rot: [custom_rot_cheap]}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=3), shots=1000) + def circuit_23(x: float, y: float): + qp.Rot(x, y, x + y, wires=1) + return qp.expval(qp.PauliZ(0)) + + # CHECK-DAG: @custom_rot_cheap(%arg0: !qref.reg<3>, %arg1: tensor<3xf64>, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} + print(circuit_23.mlir) + + qp.decomposition.disable_graph() + + +test_decompose_lowering_alt_decomps() + + +def test_decompose_lowering_with_tensorlike(): + """Test the decompose lowering pass with fixed decompositions + using TensorLike parameters.""" + + qp.decomposition.enable_graph() + + @qp.register_resources({qp.RZ: 2, qp.RY: 1}) + def custom_rot(params: TensorLike, wires: WiresLike): + qp.RZ(params[0], wires=wires) + qp.RY(params[1], wires=wires) + qp.RZ(params[2], wires=wires) + + @qp.register_resources({qp.RZ: 1, qp.CNOT: 4}) + def custom_multirz(params: TensorLike, wires: WiresLike): + qp.CNOT(wires=(wires[2], wires[1])) + qp.CNOT(wires=(wires[1], wires[0])) + qp.RZ(params[0], wires=wires[0]) + qp.CNOT(wires=(wires[1], wires[0])) + qp.CNOT(wires=(wires[2], wires[1])) + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RY", "RX", qp.CNOT}, + fixed_decomps={qp.Rot: custom_rot, qp.MultiRZ: custom_multirz}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=3), shots=1000) + def circuit_24(x: float, y: float): + qp.Rot(x, y, x + y, wires=1) + qp.MultiRZ(x + y, wires=[0, 1, 2]) + return qp.expval(qp.PauliZ(0)) + + # CHECK-DAG: @custom_multirz_wires_3(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<3xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 3 : i64, target_gate = "MultiRZ"} + # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} + # CHECK-DAG: @custom_rot(%arg0: !qref.reg<3>, %arg1: tensor<3xf64>, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} + print(circuit_24.mlir) + + qp.decomposition.disable_graph() + + +skip_if_pauli_rot_issue(test_decompose_lowering_with_tensorlike)() + + +def test_decompose_lowering_fallback(): + """Test the decompose lowering pass when the graph is failed.""" + + qp.decomposition.enable_graph() + + @qp.qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set={qp.RX, qp.RZ}) + @qp.qnode(qp.device("lightning.qubit", wires=2)) + # CHECK-LABEL: @circuit_25() + def circuit_25(): + # CHECK: [[pi_over_2:%.+]] = arith.constant 1.5707963267948966 : f64 + # CHECK: [[QREG:%.+]] = qref.alloc( 2) : !qref.reg<2> + # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<2> -> !qref.bit + # CHECK: qref.custom "RZ"([[pi_over_2]]) [[QUBIT]] : !qref.bit + # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<2> -> !qref.bit + # CHECK: qref.custom "RX"([[pi_over_2]]) [[QUBIT]] : !qref.bit + # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<2> -> !qref.bit + # CHECK: qref.custom "RZ"([[pi_over_2]]) [[QUBIT]] : !qref.bit + qp.Hadamard(0) + return qp.state() + + print(circuit_25.mlir) + + qp.decomposition.disable_graph() + + +test_decompose_lowering_fallback() + + +def test_decompose_lowering_params_ordering(): + """Test the order of params and wires in the captured decomposition rule.""" + + qp.decomposition.enable_graph() + + @qjit(target="mlir", capture=True) + @partial(qp.transforms.decompose, gate_set=[qp.RX, qp.RY, qp.RZ]) + @qp.qnode(qp.device("lightning.qubit", wires=2)) + # CHECK-LABEL: @circuit_26(%arg0: tensor, %arg1: tensor, %arg2: tensor) + def circuit_26(x: float, y: float, z: float): + qp.Rot(x, y, z, wires=0) + return qp.expval(qp.PauliZ(0)) + + # CHECK-LABEL: @_rot_to_rz_ry_rz(%arg0: !qref.reg<2>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} + # CHECK: [[EXTRACTED_1:%.+]] = tensor.extract %arg1[] : tensor + # CHECK-NEXT: qref.custom "RZ"([[EXTRACTED_1]]) {{%.+}} : !qref.bit + # CHECK: [[EXTRACTED_2:%.+]] = tensor.extract %arg2[] : tensor + # CHECK-NEXT: qref.custom "RY"([[EXTRACTED_2]]) {{%.+}} : !qref.bit + # CHECK: [[EXTRACTED_3:%.+]] = tensor.extract %arg3[] : tensor + # CHECK-NEXT: qref.custom "RZ"([[EXTRACTED_3]]) {{%.+}} : !qref.bit + # CHECK: return + print(circuit_26.mlir) + + qp.decomposition.disable_graph() + + +test_decompose_lowering_params_ordering() + + +def test_decomposition_rule_with_allocation(): + """Test decomposition rule with dynamic qubit allocation""" + + @decomposition_rule(is_qreg=True) + def Hadamard0_with_alloc(wire: WiresLike): + with qp.allocate(1) as q: + qp.X(q[0]) + qp.CNOT(wires=[q[0], wire]) + + @qp.qjit(capture=True) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + # CHECK: module @circuit_27 + def circuit_27(): + Hadamard0_with_alloc(int) + return qp.probs() + + # CHECK-LABEL: @Hadamard0_with_alloc(%arg0: !qref.reg<1>, %arg1: tensor) + # CHECK: [[dynalloc_qreg:%.+]] = qref.alloc( 1) + # CHECK: [[dynalloc_bit0:%.+]] = qref.get [[dynalloc_qreg]][ 0] + # CHECK: qref.custom "PauliX"() [[dynalloc_bit0]] + # CHECK: [[detensor:%.+]] = tensor.extract %arg1[] + # CHECK: [[glob_bit:%.+]] = qref.get %arg0[[[detensor]]] + # CHECK: qref.custom "CNOT"() [[dynalloc_bit0]], [[glob_bit]] + # CHECK: qref.dealloc [[dynalloc_qreg]] + # CHECK: return + + print(circuit_27.mlir) + + +test_decomposition_rule_with_allocation() + + +def test_decompose_autograph_multi_blocks(): + """Test the decompose lowering pass with autograph in the program and rule.""" + + qp.decomposition.enable_graph() + + def _multi_rz_decomposition_resources(num_wires): + """Resources required for MultiRZ decomposition.""" + return {qp.RZ: 1, qp.CNOT: 2 * (num_wires - 1)} + + @qp.register_resources(_multi_rz_decomposition_resources) + @qp.capture.run_autograph + def _multi_rz_decomposition(theta: TensorLike, wires: WiresLike, **__): + """Decomposition of MultiRZ using CNOTs and RZs.""" + for i in range(len(wires) - 1): + qp.CNOT(wires=(wires[i], wires[i + 1])) + qp.RZ(theta, wires=wires[0]) + for i in range(len(wires) - 1, 0, -1): + qp.CNOT(wires=(wires[i], wires[i - 1])) + + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={"RZ", "CNOT"}, + fixed_decomps={qp.MultiRZ: _multi_rz_decomposition}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=5)) + def circuit_29(n: int): + + # CHECK: scf.for %arg1 = {{%.+}} to {{%.+}} step {{%.+}} { + @qp.for_loop(n) + def f(i): # pylint: disable=unused-argument + qp.MultiRZ(0.5, wires=[0, 1, 2, 3, 4]) + + f() # pylint: disable=no-value-for-parameter + + return qp.expval(qp.Z(0)) + + # CHECK-LABEL: @ag___multi_rz_decomposition_wires_5(%arg0: !qref.reg<5>, %arg1: tensor<1xf64>, %arg2: tensor<5xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 5 : i64, target_gate = "MultiRZ"} + # CHECK: scf.for %arg3 = {{%.+}} to {{%.+}} step {{%.+}} { + # CHECK: scf.for %arg3 = {{%.+}} to {{%.+}} step {{%.+}} { + print(circuit_29.mlir) + + qp.decomposition.disable_graph() + + +test_decompose_autograph_multi_blocks() + + +def test_decompose_work_wires_context_manager(): + """ + Test that decomposition with work wires is correctly applied when allocating with the context + manager. + """ + + @decomposition_rule(is_qreg=True, op_type="PauliZ") + def my_decomp(wires): + with qp.allocate(2, restored=False) as work_wires: + qp.X(wires[0]) + qp.X(wires[1]) + qp.H(work_wires[0]) + qp.H(work_wires[1]) + + @qp.qjit(capture=True) + @qp.transform(pass_name="decompose-lowering") + @qp.qnode(qp.device("lightning.qubit", wires=3)) + def my_circuit(): + my_decomp(jax.core.ShapedArray((2,), int)) + qp.Z(0) + return qp.probs() + + # check that decomp arrives properly + # CHECK-LABEL: @my_decomp({{.*}}) attributes {{{.*}} target_gate = "PauliZ"} + print(my_circuit.mlir) + + # check that decomp is applied properly + # CHECK-NOT: PauliZ + # CHECK-NOT: my_decomp + + # two allocates, one for main register and one for decomp register + # CHECK: allocate + # CHECK: allocate + # CHECK: PauliX + # CHECK: PauliX + # CHECK: Hadamard + # CHECK: Hadamard + # CHECK: release + # CHECK: release + print(my_circuit.mlir_opt) + + +test_decompose_work_wires_context_manager() + + +def test_decompose_work_wires_alloc_dealloc(): + """ + Test that decomposition with work wires is correctly applied when allocating/deallocating + explicitly. + """ + + @decomposition_rule(is_qreg=True, op_type="RY") + def my_decomp(angle, wires): + work_wires = qp.allocate(2) + qp.CNOT((work_wires[0], wires[0])) + qp.RX(-np.pi / 2, wires[0]) + qp.RZ(angle, wires[0]) + qp.RX(np.pi / 2, wires[0]) + qp.CNOT((work_wires[1], wires[1])) + qp.deallocate(work_wires) + + @qp.qjit(capture=True) + @qp.transform(pass_name="decompose-lowering") + @qp.qnode(qp.device("lightning.qubit", wires=3)) + def my_circuit(angle: float): + my_decomp(float, jax.core.ShapedArray((2,), int)) + qp.RY(angle, 0) + return qp.probs() + + # check that decomp arrives properly + # CHECK-LABEL: @my_decomp({{.*}}) attributes {{{.*}} target_gate = "RY"} + print(my_circuit.mlir) + + # check that the decomposition applies properly + # CHECK-NOT: my_decomp + # CHECK-NOT: RY + + # two allocates, one for main register and one for decomp register + # CHECK: allocate + # CHECK: allocate + # CHECK: CNOT + # CHECK: RX + # CHECK: RZ + # CHECK: RX + # CHECK: CNOT + # CHECK: release + # CHECK: release + print(my_circuit.mlir_opt) + + +test_decompose_work_wires_alloc_dealloc() + + +def test_decompose_work_wires_control_flow(): + """Test that decomposition with work wires + control flow is correctly applied.""" + + @decomposition_rule(is_qreg=True, op_type="CRX") + def my_decomp(angle, wires, **_): + def true_func(): + qp.CNOT(wires) + + with qp.allocate(2, state="any", restored=True) as w: + for _ in range(2): + qp.H(w[0]) + qp.X(w[1]) + + def false_func(): + with qp.allocate(1, state="any", restored=False) as w: + qp.H(w) + + m = qp.measure(wires[0]) + + qp.cond(m, qp.CNOT)(wires) + + qp.cond(angle > 1.2, true_func, false_func)() + + @qp.qjit(capture=True) + @qp.transform(pass_name="decompose-lowering") + @qp.qnode(qp.device("lightning.qubit", wires=4)) + def circuit(): + my_decomp(float, jax.core.ShapedArray((2,), int)) + qp.CRX(1.7, wires=[0, 1]) + qp.CRX(-7.2, wires=[0, 1]) + return qp.state() + + # target_gate attribute is correctly applied + # CHECK: my_decomp([[args:.*]]) attributes {[[other_attributes:.*]] target_gate = "CRX"} + print(circuit.mlir) + + # test that the decomposition is applied correctly + # CHECK-NOT: CRX + # CHECK-NOT: my_decomp + + # allocate for main register, subsequent allocates+releases for decomp registers + # CHECK: allocate + + # first CRX: true branch + # CHECK: CNOT + # CHECK: allocate + # CHECK: Hadamard + # CHECK: PauliX + # CHECK: Hadamard + # CHECK: PauliX + # CHECK: release + + # second CRX: false branch + # CHECK: allocate + # CHECK: Hadamard + # CHECK: Measure + # CHECK: cond + # CHECK: CNOT + # CHECK: release + + # release main register + # CHECK: release + + print(circuit.mlir_opt) + + +test_decompose_work_wires_control_flow() + + +def test_decompose_work_wires_with_decompose_transform(): + """Test that work wires are correctly lowered and decomposed by the decompose transform.""" + + qp.decomposition.enable_graph() + + @qp.register_resources({qp.X: 1, qp.Z: 1}) + def my_decomp(wire): + with qp.allocate(1) as work_wire: + qp.X(work_wire) + qp.Z(wire) + qp.X(work_wire) + + @qjit(capture=True) + @partial( + qp.transforms.decompose, + gate_set={ + "X", + "Z", + }, + fixed_decomps={ + qp.Y: my_decomp, + }, + ) + @qp.qnode(qp.device("lightning.qubit", wires=2)) + def my_circuit(): + qp.Y(0) + return qp.probs() + + # CHECK-NOT: Y + # CHECK-NOT: my_decomp + + # two allocates, one for main register and one for decomp register + # CHECK: allocate + # CHECK: allocate + # CHECK: X + # CHECK: Z + # CHECK: X + # CHECK: release + # CHECK: release + print(my_circuit.mlir_opt) + + qp.decomposition.disable_graph() + + +test_decompose_work_wires_with_decompose_transform() + + +def test_num_work_wires(): + """Test that num_work_wires can be passed and is correctly used in solving the graph.""" + + qp.decomposition.enable_graph() + + @qp.register_resources( + {qp.CNOT: 3, qp.H: 1, qp.X: 1, qp.ops.op_math.Conditional: 2}, + work_wires={"borrowed": 2, "garbage": 1}, + ) + def my_decomp(angle, wires, **_): + def true_func(): + qp.CNOT(wires) + + with qp.allocate(2, state="any", restored=True) as w: + qp.H(w[0]) + qp.H(w[0]) + qp.X(w[1]) + qp.X(w[1]) + + return + + def false_func(): + with qp.allocate(1, state="any", restored=False) as w: + qp.H(w) + + m = qp.measure(wires[0]) + + qp.cond(m, qp.CNOT)(wires) + + return + + qp.cond(angle > 1.2, true_func, false_func)() + + @qp.qjit(capture=True) + @partial( + qp.transforms.decompose, + gate_set={qp.CNOT, qp.H, qp.X, "Conditional", "MidMeasure"}, + fixed_decomps={qp.CRX: my_decomp}, + num_work_wires=3, + ) + @qp.qnode(qp.device("lightning.qubit", wires=5)) + def circuit(): + qp.CRX(1.7, wires=[0, 1]) + qp.CRX(-7.2, wires=[0, 1]) + return qp.state() + + # CHECK-NOT: CRX + # CHECK-NOT: my_decomp + + # CHECK: allocate + # CHECK: allocate + # CHECK: CNOT + # CHECK: Hadamard + # CHECK: Hadamard + # CHECK: PauliX + # CHECK: PauliX + # CHECK: Hadamard + # CHECK: Measure + # CHECK: CNOT + # CHECK: release + # CHECK: release + print(circuit.mlir_opt) + + qp.decomposition.disable_graph() + + +test_num_work_wires() + + +def test_default_decomps(): + """Test that default decompositions are correctly applied with qjit.""" + qp.decomposition.enable_graph() + + # Toffoli's decomposition to this gateset includes a wire allocation + @qp.qjit(target="mlir", capture=True) + @partial( + qp.transforms.decompose, + gate_set={qp.ops.ChangeOpBasis}, + num_work_wires=1, + ) + @qp.qnode(qp.device("lightning.qubit", wires=4)) + def circuit(): + qp.Toffoli(wires=[0, 1, 2]) + return qp.state() + + # CHECK-NOT: toffoli_elbow + # CHECK-NOT: Toffoli + + # two allocates/releases, for default register + work wires + # CHECK: allocate + # CHECK: allocate + # CHECK: TemporaryAND + # CHECK: release + # CHECK: release + print(circuit.mlir_opt) + + qp.decomposition.disable_graph() + + +test_default_decomps() + + +def test_graph_decomp_registered(): + """Test that the `graph_decomposition` pass is registered correctly.""" + + @qjit(target="mlir", capture=True) + # CHECK: transform.apply_registered_pass "graph-decomposition" + @graph_decomposition(gate_set={qp.RX}) + @qp.qnode(qp.device("lightning.qubit", wires=2)) + def catalyst_circuit(): + return + + print(catalyst_circuit.mlir) + + my_transform = qp.transform(pass_name="graph-decomposition") + + @qjit(target="mlir", capture=True) + # CHECK: transform.apply_registered_pass "graph-decomposition" + @my_transform(gate_set=["RX"]) + @qp.qnode(qp.device("lightning.qubit", wires=2)) + def pennylane_circuit(): + return + + print(pennylane_circuit.mlir) + + +test_graph_decomp_registered() + + +def test_cpp_decomp_args(): + """Test that the `graph_decomposition` pass lowers arguments to mlir correctly.""" + + def x_to_rx(wire): + qp.RX(np.pi, wire) + + def y_to_ry(wire): + qp.RY(np.pi, wire) + + def h_to_rx_ry(wire): + qp.RX(np.pi / 2, wire) + qp.RY(np.pi / 2, wire) + + @qjit(target="mlir") + # CHECK: "graph-decomposition" with options = { + # CHECK-DAG: "gate-set" = {Hadamard = 1.000000e+00 : f64, RX = 1.000000e+00 : f64, RY = 1.000000e+00 : f64} + # CHECK-DAG: "fixed-decomps" = {PauliX = "x_to_rx", PauliY = "y_to_ry"} + # CHECK-DAG: "alt-decomps" = {Hadamard = ["h_to_rx_ry"]} + # CHECK-DAG: "bytecode-rules" = "{{.*}}decomposition_rules_{{.*}}.mlirbc" + # CHECK: } to {{%.+}} : (!transform.op<"builtin.module">) + @graph_decomposition( + gate_set={qp.RX, qp.H, qp.RY}, + fixed_decomps={qp.X: x_to_rx, qp.Y: y_to_ry}, + alt_decomps={qp.H: [h_to_rx_ry]}, + _builtin_rule_path="/decomp_rules.mlirbc", + ) + @qp.qnode(qp.device("lightning.qubit", wires=2)) + def circuit(): + return + + print(circuit.mlir) + + +test_cpp_decomp_args() + + +def test_cpp_decomp_empty_args(): + """ + Test that the `graph_decomposition` pass correctly handled arg lowering when no values are + supplied. + """ + + @qjit(target="mlir", capture=True) + # CHECK: transform.apply_registered_pass "graph-decomposition" + # CHECK-NOT: fixed-decomps + # CHECK-NOT: alt-decomps + # CHECK: "bytecode-rules" = "{{.*}}/decomposition_rules{{.*}}.mlirbc" + @graph_decomposition(gate_set={qp.RX}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + def circuit(): + return + + print(circuit.mlir) + + @qjit(target="mlir", capture=True) + # CHECK: transform.apply_registered_pass "graph-decomposition" + # CHECK-NOT: fixed-decomps + # CHECK-NOT: alt-decomps + # CHECK: "bytecode-rules" = "{{.*}}/decomposition_rules{{.*}}.mlirbc" + @graph_decomposition(gate_set={qp.RX}, fixed_decomps={}, alt_decomps={}) + @qp.qnode(qp.device("lightning.qubit", wires=1)) + def circuit2(): + return + + print(circuit2.mlir) + + +test_cpp_decomp_empty_args() + + +def test_cpp_decomp_string_op_names(): + """Test that cpp decomp args work with string op names.""" + + def y_to_xz(wires): + qp.RX(np.pi, wires) + qp.RZ(np.pi, wires) + + @qjit(target="mlir", capture=True) + # CHECK: transform.apply_registered_pass "graph-decomposition" with options = { + # CHECK-DAG: "fixed-decomps" = {PauliX = "{{.*}}", PauliZ = "{{.*}}"} + # CHECK-DAG: "alt-decomps" = {PauliY = ["{{.*}}", "y_to_xz"]} + # CHECK: } to {{%.+}} : (!transform.op<"builtin.module">) + @graph_decomposition( + gate_set={"RX", "RY", "RZ"}, + fixed_decomps={ + "X": lambda wires: qp.RX(np.pi, wires), + "PauliZ": lambda wires: qp.RZ(np.pi, wires), + }, + alt_decomps={ + "PauliY": [ + lambda wires: qp.RY(np.pi, wires), + y_to_xz, + ] + }, + ) + @qp.qnode(qp.device("lightning.qubit", wires=2)) + def circuit(): + return + + print(circuit.mlir) + + +test_cpp_decomp_string_op_names() + + +def test_cpp_decomp_builtin_rules(): + """Test that cpp decomp applies builtin rules.""" + + @qjit(target="mlir", capture=True) + @graph_decomposition( + gate_set={qp.RX, qp.RY, qp.RZ, qp.GlobalPhase}, + ) + @qp.qnode(qp.device("lightning.qubit", wires=2)) + def circuit(): + # CHECK-NOT: PauliX + # CHECK-NOT: PauliY + # CHECK-NOT: PauliZ + # CHECK-DAG: RX + # CHECK-DAG: RY + # CHECK-DAG: RZ + qp.X(0) + qp.Y(1) + qp.Z(0) + return qp.probs() + + print(circuit.mlir_opt) + + +test_cpp_decomp_builtin_rules() + + +def test_cpp_decomp_user_rules(): + """Test that cpp decomp applies user rules.""" + + @decomposition_rule(is_qreg=True, op_type="PauliY") + def y_to_rx(wire): + qp.RX(np.pi, wire) + + @decomposition_rule(is_qreg=True, op_type="PauliZ") + def z_to_rx(wire): + qp.RX(np.pi, wire) + + @qp.qjit(target="mlir", capture=True) + @graph_decomposition( + gate_set={qp.RX}, fixed_decomps={qp.Y: y_to_rx}, alt_decomps={qp.Z: [z_to_rx]} + ) + @qp.qnode(qp.device("null.qubit", wires=1)) + def circuit(): + y_to_rx(jax.core.ShapedArray((1,), int)) + z_to_rx(jax.core.ShapedArray((1,), int)) + # CHECK-NOT: PauliY + # CHECK-NOT: PauliZ + # CHECK: RX + # CHECK: RX + # CHECK: return + qp.Y(0) + qp.Z(0) + return qp.probs() + + print(circuit.mlir_opt) + + +test_cpp_decomp_user_rules() + + +def test_cpp_decomp_user_rule_cleanup(): + """Test that user rules do not pollute the IR after the quantum compilation stage.""" + + @decomposition_rule(is_qreg=True, op_type="PauliX") + def x_to_h(wire): + return qp.H(wire) + + @qjit(capture=True) + @graph_decomposition(gate_set={qp.H}, fixed_decomps={qp.X: x_to_h}) + @qp.qnode(qp.device("null.qubit", wires=1)) + def circuit(): + # CHECK-NOT: PauliX + # CHECK-NOT: x_to_h + x_to_h(jax.core.ShapedArray((1,), int)) + qp.X(0) + + print(circuit.mlir_opt) + + +test_cpp_decomp_user_rule_cleanup() + + +def test_paulirot_python_decomposition(): + """Test that paulirots are decomposed by the mlir graph.""" + + @qjit(capture=True) + @graph_decomposition(gate_set={qp.H, qp.MultiRZ, qp.GlobalPhase, qp.RX}) + @qp.qnode(qp.device("null.qubit", wires=4)) + def circuit(): + qp.PauliRot(0.9, "XZXY", [0, 1, 2, 3]) + return qp.probs() + + print(circuit.mlir_opt) + + # CHECK-NOT: quantum.paulirot + # CHECK: Hadamard + # CHECK: Hadamard + # CHECK: RX + # CHECK: MultiRZ + # CHECK: Hadamard + # CHECK: Hadamard + # CHECK: RX + + +test_paulirot_python_decomposition() diff --git a/frontend/test/lit/test_decomposition.py b/frontend/test/lit/test_decomposition.py index e8f9448e7c..fb8fe117eb 100644 --- a/frontend/test/lit/test_decomposition.py +++ b/frontend/test/lit/test_decomposition.py @@ -1,22 +1,4 @@ -# Copyright 2022-2025 Xanadu Quantum Technologies Inc. -import os -import pathlib -import platform -from copy import deepcopy -from functools import partial - -import jax -import numpy as np -import pennylane as qp -from pennylane.devices.capabilities import OperatorProperties -from pennylane.typing import TensorLike -from pennylane.wires import WiresLike - -from catalyst import measure, qjit -from catalyst.compiler import get_lib_path -from catalyst.device import get_device_capabilities -from catalyst.jax_primitives import decomposition_rule -from catalyst.passes import graph_decomposition +# 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. @@ -29,1835 +11,3 @@ # 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. - -# RUN: %PYTHON %s | FileCheck %s -# pylint: disable=line-too-long -# pylint: disable=too-many-lines - - -# Helper to skip tests that fail due to PauliRot type annotation issue -# TODO: Remove this once PennyLane fixes the PauliRot decomposition type annotations -def skip_if_pauli_rot_issue(test_func): - """Wrapper to skip tests that fail due to PauliRot type annotation issues.""" - - def wrapper(): - try: - test_func() - except (ValueError, IndexError) as e: - error_msg = str(e) - if ( - "Unsupported type annotation None for parameter pauli_word" in error_msg - or "Unsupported type annotation for parameter pauli_word" in error_msg - or "index is out of bounds for axis" in error_msg - ): - print(f"# SKIPPED {test_func.__name__}: PauliRot type annotation issue") - else: - raise - - return wrapper - - -TEST_PATH = os.path.dirname(__file__) -CONFIG_CUSTOM_DEVICE = pathlib.Path(f"{TEST_PATH}/../custom_device/custom_device.toml") - - -def get_custom_device_without(num_wires, discards=frozenset(), force_matrix=frozenset()): - """Generate a custom device without gates in discards.""" - - class CustomDevice(qp.devices.Device): - """Custom Gate Set Device""" - - name = "Custom Device" - config_filepath = CONFIG_CUSTOM_DEVICE - - _to_matrix_ops = {} - - def __init__(self, wires=None): - super().__init__(wires=wires) - self.qjit_capabilities = deepcopy(get_device_capabilities(self)) - for gate in discards: - self.qjit_capabilities.operations.pop(gate, None) - for gate in force_matrix: - self.qjit_capabilities.operations.pop(gate, None) - self._to_matrix_ops[gate] = OperatorProperties(False, False, False) - - def apply(self, operations, **kwargs): - """Unused""" - raise RuntimeError("Only C/C++ interface is defined") - - @staticmethod - def get_c_interface(): - """Returns a tuple consisting of the device name, and - the location to the shared object with the C/C++ device implementation. - """ - system_extension = ".dylib" if platform.system() == "Darwin" else ".so" - lib_path = ( - get_lib_path("runtime", "RUNTIME_LIB_DIR") + "/librtd_null_qubit" + system_extension - ) - return "NullQubit", lib_path - - def execute(self, circuits, execution_config): - """Execution.""" - return circuits, execution_config - - return CustomDevice(wires=num_wires) - - -def test_decompose_multicontrolledx(): - """Test decomposition of MultiControlledX as an aliased gate.""" - dev = get_custom_device_without(5, discards={"MultiControlledX"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_multicontrolled_x1 - def decompose_multicontrolled_x1(theta: float): - qp.RX(theta, wires=[0]) - # CHECK-NOT: name = "MultiControlledX" - # CHECK: quantum.custom "PauliX"() {{%[a-zA-Z0-9_]+}} ctrls({{%[a-zA-Z0-9_]+}}, {{%[a-zA-Z0-9_]+}}, {{%[a-zA-Z0-9_]+}}) - # CHECK-NOT: name = "MultiControlledX" - qp.MultiControlledX(wires=[0, 1, 2, 3]) - return qp.state() - - print(decompose_multicontrolled_x1.mlir) - - -test_decompose_multicontrolledx() - - -def test_decompose_rot(): - """Test decomposition of Rot gate.""" - dev = get_custom_device_without(1, discards={"Rot", "C(Rot)"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_rot - def decompose_rot(phi: float, theta: float, omega: float): - # CHECK-NOT: name = "Rot" - # CHECK: [[phi:%.+]] = tensor.extract %arg0 - # CHECK-NOT: name = "Rot" - # CHECK: {{%.+}} = quantum.custom "RZ"([[phi]]) - # CHECK-NOT: name = "Rot" - # CHECK: [[theta:%.+]] = tensor.extract %arg1 - # CHECK-NOT: name = "Rot" - # CHECK: {{%.+}} = quantum.custom "RY"([[theta]]) - # CHECK-NOT: name = "Rot" - # CHECK: [[omega:%.+]] = tensor.extract %arg2 - # CHECK-NOT: name = "Rot" - # CHECK: {{%.+}} = quantum.custom "RZ"([[omega]]) - # CHECK-NOT: name = "Rot" - qp.Rot(phi, theta, omega, wires=0) - return measure(wires=0) - - print(decompose_rot.mlir) - - -test_decompose_rot() - - -def test_decompose_s(): - """Test decomposition of S gate.""" - dev = get_custom_device_without(1, discards={"S", "C(S)"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_s - def decompose_s(): - # CHECK-NOT: name="S" - # CHECK: [[pi_div_2:%.+]] = arith.constant 1.57079{{.+}} : f64 - # CHECK-NOT: name = "S" - # CHECK: {{%.+}} = quantum.custom "PhaseShift"([[pi_div_2]]) - # CHECK-NOT: name = "S" - qp.S(wires=0) - return measure(wires=0) - - print(decompose_s.mlir) - - -test_decompose_s() - - -def test_decompose_qubitunitary(): - """Test decomposition of QubitUnitary""" - dev = get_custom_device_without(1, discards={"QubitUnitary"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_qubit_unitary - def decompose_qubit_unitary(U: jax.core.ShapedArray([2, 2], float)): - # CHECK-NOT: name = "QubitUnitary" - # CHECK: quantum.custom "RZ" - # CHECK: quantum.custom "RY" - # CHECK: quantum.custom "RZ" - # CHECK-NOT: name = "QubitUnitary" - qp.QubitUnitary(U, wires=0) - return measure(wires=0) - - print(decompose_qubit_unitary.mlir) - - -test_decompose_qubitunitary() - - -def test_decompose_singleexcitation(): - """ - Test that single excitation is not decomposed. - """ - dev = get_custom_device_without(2) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_singleexcitation - def decompose_singleexcitation(theta: float): - # CHECK: quantum.custom "SingleExcitation" - - qp.SingleExcitation(theta, wires=[0, 1]) - return measure(wires=0) - - print(decompose_singleexcitation.mlir) - - -test_decompose_singleexcitation() - - -def test_decompose_doubleexcitation(): - """ - Test that Double excitation is not decomposed. - """ - dev = get_custom_device_without(4) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_doubleexcitation - def decompose_doubleexcitation(theta: float): - # CHECK: quantum.custom "DoubleExcitation" - - qp.DoubleExcitation(theta, wires=[0, 1, 2, 3]) - return measure(wires=0) - - print(decompose_doubleexcitation.mlir) - - -test_decompose_doubleexcitation() - - -def test_decompose_singleexcitationplus(): - """ - Test decomposition of single excitation plus. - See - https://github.com/PennyLaneAI/pennylane/blob/main/pennylane/ops/qubit/qchem_ops.py - for the decomposition of qp.SingleExcitationPlus - """ - dev = get_custom_device_without(2, discards={"SingleExcitationPlus", "C(SingleExcitationPlus)"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_singleexcitationplus - def decompose_singleexcitationplus(theta: float): - # CHECK-NOT: "SingleExcitationPlus" - # CHECK: quantum.custom "Hadamard" - # CHECK: quantum.custom "CNOT" - # CHECK: quantum.custom "RY" - # CHECK: quantum.custom "RY" - # CHECK: quantum.custom "CY" - # CHECK: quantum.custom "S" - # CHECK: quantum.custom "Hadamard" - # CHECK: quantum.custom "RZ" - # CHECK: quantum.custom "CNOT" - # CHECK: quantum.gphase - - qp.SingleExcitationPlus(theta, wires=[0, 1]) - return measure(wires=0) - - print(decompose_singleexcitationplus.mlir) - - -test_decompose_singleexcitationplus() - - -def test_decompose_to_matrix(): - """Test decomposition of QubitUnitary""" - dev = get_custom_device_without(1, force_matrix={"PauliY"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_to_matrix - def decompose_to_matrix(): - # CHECK: quantum.custom "PauliX" - qp.PauliX(wires=0) - # CHECK: quantum.unitary - qp.PauliY(wires=0) - # CHECK: quantum.custom "PauliZ" - qp.PauliZ(wires=0) - return measure(wires=0) - - print(decompose_to_matrix.mlir) - - -test_decompose_to_matrix() - - -def test_decomposition_rule_lowering(): - """Test that decomposition rules are lowered to private functions.""" - - @decomposition_rule(is_qreg=True) - def my_decomp(): - return - - @qp.qjit(capture=True) - @qp.qnode(qp.device("null.qubit", wires=1)) - def circuit(): - # CHECK-LABEL: func.func private @my_decomp - my_decomp() - return - - print(circuit.mlir) - - -test_decomposition_rule_lowering() - - -def test_decomposition_rule_wire_param(): - """Test decomposition rule with passing a parameter that is a wire/integer""" - - @decomposition_rule(is_qreg=False) - def Hadamard0(wire: WiresLike): - qp.Hadamard(wire) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @circuit - def circuit(_: float): - # CHECK: @circuit([[ARG0:%.+]] - # CHECK: [[QREG:%.+]] = qref.alloc - Hadamard0(int) - return qp.probs() - - # CHECK: @Hadamard0([[QBIT:%.+]]: !qref.bit) - # CHECK-NEXT: qref.custom "Hadamard"() [[QBIT]] : !qref.bit - # CHECK-NEXT: return - - print(circuit.mlir) - - -test_decomposition_rule_wire_param() - - -def test_decomposition_rule_gate_param_param(): - """Test decomposition rule with passing a regular parameter""" - - @decomposition_rule(is_qreg=False, num_params=1) - def RX_on_wire_0(param: TensorLike, w0: WiresLike): - qp.RX(param, wires=w0) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module @circuit_2 - def circuit_2(_: float): - RX_on_wire_0(float, int) - return qp.probs() - - # CHECK: @RX_on_wire_0([[PARAM_TENSOR:%.+]]: tensor, [[QUBIT:%.+]]: !qref.bit) - # CHECK-NEXT: [[PARAM:%.+]] = tensor.extract [[PARAM_TENSOR]][] : tensor - # CHECK-NEXT: qref.custom "RX"([[PARAM]]) [[QUBIT]] : !qref.bit - # CHECK-NEXT: return - print(circuit_2.mlir) - - -test_decomposition_rule_gate_param_param() - - -def test_multiple_decomposition_rules(): - """Test with multiple decomposition rules""" - - @decomposition_rule - def identity(): ... - - @decomposition_rule(is_qreg=True) - def all_wires_rx(param: TensorLike, w0: WiresLike, w1: WiresLike, w2: WiresLike): - qp.RX(param, wires=w0) - qp.RX(param, wires=w1) - qp.RX(param, wires=w2) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit_3(_: float): - # CHECK: [[QREG:%.+]] = qref.alloc - # CHECK-NEXT: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK-NEXT: qref.custom "Hadamard"() [[QUBIT]] : !qref.bit - # CHECK-NEXT: qref.compbasis(qreg [[QREG]] : !qref.reg<1>) : !quantum.obs - identity() - all_wires_rx(float, int, int, int) - qp.Hadamard(0) - return qp.probs() - - # CHECK-LABEL: @identity - # CHECK-LABEL: @all_wires_rx - - print(circuit_3.mlir) - - -test_multiple_decomposition_rules() - - -def test_decomposition_rule_shaped_wires(): - """Test decomposition rule with passing a shaped array of wires""" - - @decomposition_rule(is_qreg=True) - def shaped_wires_rule(param: TensorLike, wires: WiresLike): - qp.RX(param, wires=wires[0]) - qp.RX(param, wires=wires[1]) - qp.RX(param, wires=wires[2]) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit_4(_: float): - # CHECK: module @circuit_4 - shaped_wires_rule(float, jax.core.ShapedArray((3,), int)) - qp.Hadamard(0) - return qp.probs() - - # CHECK: @shaped_wires_rule([[QREG:%.+]]: !qref.reg<1>, [[PARAM_TENSOR:%.+]]: tensor, [[QUBITS:%.+]]: tensor<3xi64>) - # CHECK-NEXT: [[IDX_0:%.+]] = stablehlo.slice [[QUBITS]] [0:1] : (tensor<3xi64>) -> tensor<1xi64> - # CHECK-NEXT: [[RIDX_0:%.+]] = stablehlo.reshape [[IDX_0]] : (tensor<1xi64>) -> tensor - # CHECK-NEXT: [[EXTRACTED:%.+]] = tensor.extract [[RIDX_0]][] : tensor - # CHECK-NEXT: [[QUBIT:%.+]] = qref.get [[QREG]][[[EXTRACTED]]] : !qref.reg<1>, i64 -> !qref.bit - # CHECK-NEXT: [[EXTRACTED_0:%.+]] = tensor.extract [[PARAM_TENSOR]][] : tensor - # CHECK-NEXT: qref.custom "RX"([[EXTRACTED_0]]) [[QUBIT]] : !qref.bit - - print(circuit_4.mlir) - - -test_decomposition_rule_shaped_wires() - - -def test_decomposition_rule_expanded_wires(): - """Test decomposition rule with passing expanding wires as a Python list""" - - def shaped_wires_rule(param: TensorLike, wires: WiresLike): - qp.RX(param, wires=wires[0]) - qp.RX(param, wires=wires[1]) - qp.RX(param, wires=wires[2]) - - @decomposition_rule(is_qreg=False, num_params=1) - def expanded_wires_rule(param: TensorLike, w1, w2, w3): - shaped_wires_rule(param, [w1, w2, w3]) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit_5(_: float): - # CHECK: module @circuit_5 - expanded_wires_rule(float, int, int, int) - qp.Hadamard(0) - return qp.probs() - - # CHECK-LABEL: @expanded_wires_rule(%arg0: tensor, %arg1: !qref.bit, %arg2: !qref.bit, %arg3: !qref.bit) - - print(circuit_5.mlir) - - -test_decomposition_rule_expanded_wires() - - -def test_decomposition_rule_with_cond(): - """Test decomposition rule with a conditional path""" - - @decomposition_rule(is_qreg=True) - def cond_RX(param: TensorLike, w0: WiresLike): - - def true_path(): - qp.RX(param, wires=w0) - - def false_path(): ... - - qp.cond(param != 0.0, true_path, false_path)() - - @qp.qjit(autograph=False, capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit_6(): - # CHECK: module @circuit_6 - cond_RX(float, jax.core.ShapedArray((1,), int)) - return qp.probs() - - # CHECK: @cond_RX([[QREG:%.+]]: !qref.reg<1>, [[PARAM_TENSOR:%.+]]: tensor, [[QUBITS:%.+]]: tensor<1xi64>) - # CHECK-NEXT: [[ZERO:%.+]] = stablehlo.constant dense<0.000000e+00> : tensor - # CHECK-NEXT: [[COND_TENSOR:%.+]] = stablehlo.compare NE, [[PARAM_TENSOR]], [[ZERO]], FLOAT : (tensor, tensor) -> tensor - # CHECK-NEXT: [[COND:%.+]] = tensor.extract [[COND_TENSOR]][] : tensor - # CHECK-NEXT: scf.if [[COND]] - # CHECK-DAG: [[QUBIT:%.+]] = qref.get [[QREG]][%extracted_0] : !qref.reg<1>, i64 -> !qref.bit - # CHECK-DAG: [[PARAM:%.+]] = tensor.extract [[PARAM_TENSOR]][] : tensor - # CHECK: qref.custom "RX"([[PARAM]]) [[QUBIT]] : !qref.bit - # CHECK: return - - print(circuit_6.mlir) - - -test_decomposition_rule_with_cond() - - -def test_decomposition_rule_caller(): - """Test decomposition rules with a caller""" - - @decomposition_rule(is_qreg=True) - def rule_op1_decomp(_: TensorLike, wires: WiresLike): - qp.Hadamard(wires=wires[0]) - qp.Hadamard(wires=[1]) - - @decomposition_rule(is_qreg=True) - def rule_op2_decomp(param: TensorLike, wires: WiresLike): - qp.RX(param, wires=wires[0]) - - def decomps_caller(param: TensorLike, wires: WiresLike): - rule_op1_decomp(param, wires) - rule_op2_decomp(param, wires) - - @qp.qjit(autograph=False, capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module @circuit_7 - def circuit_7(): - # CHECK: [[QREG:%.+]] = qref.alloc - # CHECK: qref.compbasis(qreg [[QREG]] : !qref.reg<1>) : !quantum.obs - decomps_caller(float, jax.core.ShapedArray((2,), int)) - return qp.probs() - - # CHECK-LABEL: @rule_op1_decomp(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<2xi64>) - # CHECK-LABEL: @rule_op2_decomp(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<2xi64>) - print(circuit_7.mlir) - - -test_decomposition_rule_caller() - - -def test_decompose_gateset_without_graph(): - """Test the decompose transform to a target gate set without the graph decomposition.""" - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RX", "RZ"}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @circuit_8() -> tensor attributes {diff_method = "adjoint", llvm.linkage = #llvm.linkage, quantum.node} - def circuit_8(): - return qp.expval(qp.Z(0)) - - print(circuit_8.mlir) - - -test_decompose_gateset_without_graph() - - -def test_decompose_gateset_with_graph(): - """Test the decompose transform to a target gate set with the graph decomposition.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RX"}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @simple_circuit_9() -> tensor attributes {decompose_gatesets - def simple_circuit_9(): - return qp.expval(qp.Z(0)) - - print(simple_circuit_9.mlir) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RX", "RZ"}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_9() -> tensor attributes {decompose_gatesets - def circuit_9(): - return qp.expval(qp.Z(0)) - - print(circuit_9.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_gateset_with_graph() - - -def test_decompose_gateset_operator_with_graph(): - """Test the decompose transform to a target gate set with the graph decomposition.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.RX}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @simple_circuit_10() -> tensor attributes {decompose_gatesets - def simple_circuit_10(): - return qp.expval(qp.Z(0)) - - print(simple_circuit_10.mlir) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.RX, qp.RZ, "PauliZ", qp.PauliX, qp.Hadamard}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @circuit_10() -> tensor attributes {decompose_gatesets - def circuit_10(): - return qp.expval(qp.Z(0)) - - print(circuit_10.mlir) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.RX, qp.RZ, qp.PauliZ, qp.PauliX, qp.Hadamard}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_11() -> tensor attributes {decompose_gatesets - def circuit_11(): - return qp.expval(qp.Z(0)) - - print(circuit_11.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_gateset_operator_with_graph() - - -def test_decompose_gateset_with_rotxzx(): - """Test the decompose transform with a custom operator with the graph decomposition.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RotXZX"}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @simple_circuit_12() -> tensor attributes {decompose_gatesets - def simple_circuit_12(): - return qp.expval(qp.Z(0)) - - print(simple_circuit_12.mlir) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.ftqc.RotXZX}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_12() -> tensor attributes {decompose_gatesets - def circuit_12(): - return qp.expval(qp.Z(0)) - - print(circuit_12.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_gateset_with_rotxzx() - - -def test_decomposition_rule_name(): - """Test the name of the decomposition rule is not updated with circuit instantiation.""" - - qp.decomposition.enable_graph() - - @decomposition_rule - def _ry_to_rz_rx(phi, wires: WiresLike, **__): - """Decomposition of RY gate using RZ and RX gates.""" - qp.RZ(-np.pi / 2, wires=wires) - qp.RX(phi, wires=wires) - qp.RZ(np.pi / 2, wires=wires) - - @decomposition_rule - def _rot_to_rz_ry_rz(phi, theta, omega, wires: WiresLike, **__): - """Decomposition of Rot gate using RZ and RY gates.""" - qp.RZ(phi, wires=wires) - qp.RY(theta, wires=wires) - qp.RZ(omega, wires=wires) - - @decomposition_rule - def _u2_phaseshift_rot_decomposition(phi, delta, wires, **__): - """Decomposition of U2 gate using Rot and PhaseShift gates.""" - pi_half = qp.math.ones_like(delta) * (np.pi / 2) - qp.Rot(delta, pi_half, -delta, wires=wires) - qp.PhaseShift(delta, wires=wires) - qp.PhaseShift(phi, wires=wires) - - @decomposition_rule - def _xzx_decompose(phi, theta, omega, wires, **__): - """Decomposition of Rot gate using RX and RZ gates in XZX format.""" - qp.RX(phi, wires=wires) - qp.RZ(theta, wires=wires) - qp.RX(omega, wires=wires) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RX", "RZ", "PhaseShift"}) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_13() -> tensor attributes {decompose_gatesets - def circuit_13(): - _ry_to_rz_rx(float, int) - _rot_to_rz_ry_rz(float, float, float, int) - _u2_phaseshift_rot_decomposition(float, float, int) - _xzx_decompose(float, float, float, int) - return qp.expval(qp.Z(0)) - - # CHECK-LABEL: @_ry_to_rz_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor) - # CHECK-LABEL: @_rot_to_rz_ry_rz(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor) - # CHECK-LABEL: @_u2_phaseshift_rot_decomposition(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor) - # CHECK-LABEL: @_xzx_decompose(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor) - print(circuit_13.mlir) - - qp.decomposition.disable_graph() - - -test_decomposition_rule_name() - - -def test_decomposition_rule_name_update(): - """Test the name of the decomposition rule is updated in the MLIR output.""" - - qp.decomposition.enable_graph() - - @qp.register_resources({qp.RZ: 2, qp.RX: 1}) - def rz_rx(phi, wires: WiresLike, **__): - """Decomposition of RY gate using RZ and RX gates.""" - qp.RZ(-np.pi / 2, wires=wires) - qp.RX(phi, wires=wires) - qp.RZ(np.pi / 2, wires=wires) - - @qp.register_resources({qp.RZ: 2, qp.RY: 1}) - def rz_ry_rz(phi, theta, omega, wires: WiresLike, **__): - """Decomposition of Rot gate using RZ and RY gates.""" - qp.RZ(phi, wires=wires) - qp.RY(theta, wires=wires) - qp.RZ(omega, wires=wires) - - @qp.register_resources({qp.RY: 1, qp.GlobalPhase: 1}) - def ry_gp(wires: WiresLike, **__): - """Decomposition of PauliY gate using RY and GlobalPhase gates.""" - qp.RY(np.pi, wires=wires) - qp.GlobalPhase(-np.pi / 2, wires=wires) - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RX", "RZ", "GlobalPhase"}, - fixed_decomps={ - qp.RY: rz_rx, - qp.Rot: rz_ry_rz, - qp.PauliY: ry_gp, - }, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_14() -> tensor attributes {decompose_gatesets - def circuit_14(): - qp.RY(0.5, wires=0) - qp.Rot(0.1, 0.2, 0.3, wires=1) - qp.PauliY(wires=2) - return qp.expval(qp.Z(0)) - - # CHECK-DAG: @rz_ry_rz(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) - # CHECK-DAG: @rz_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) - # CHECK-DAG: @ry_gp(%arg0: !qref.reg<3>, %arg1: tensor<1xi64>) - print(circuit_14.mlir) - - qp.decomposition.disable_graph() - - -test_decomposition_rule_name_update() - - -def test_decomposition_inside_subroutine(): - """Test that operators inside subroutines can be decomposed.""" - - qp.decomposition.enable_graph() - - @qp.templates.Subroutine - def f(x, wires): - qp.IsingXX(x, wires) - - @qp.qjit(capture=True, target="mlir") - @qp.decompose(gate_set=qp.gate_sets.ROTATIONS_PLUS_CNOT) - @qp.qnode(qp.device("lightning.qubit", wires=5)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - def subroutine_circuit(): - # CHECK-DAG: [[FIRST_CONST:%.+]] = stablehlo.constant dense<5.000000e-01> : tensor - # CHECK-DAG: [[SECOND_CONST:%.+]] = stablehlo.constant dense<1.200000e+00> : tensor - - # CHECK: [[QREG:%.+]] = qref.alloc - # CHECK: call @f([[QREG]], [[FIRST_CONST]], {{%.+}}) : (!qref.reg<5>, tensor, tensor<2xi64>) - # CHECK: call @f([[QREG]], [[SECOND_CONST]], {{%.+}}) : (!qref.reg<5>, tensor, tensor<2xi64>) - - f(0.5, (0, 1)) - f(1.2, (2, 3)) - return qp.probs(wires=0) - - # CHECK-DAG: @_isingxx_to_cnot_rx_cnot(%arg0: !qref.reg<5>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) - print(subroutine_circuit.mlir) - qp.decomposition.disable_graph() - - -test_decomposition_inside_subroutine() - - -def test_decomposition_rule_name_update_multi_qubits(): - """Test the name of the decomposition rule with multi-qubit gates.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RY", "RX", "CNOT", "Hadamard", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_15() -> tensor attributes {decompose_gatesets - def circuit_15(): - qp.SingleExcitation(0.5, wires=[0, 1]) - qp.SingleExcitationPlus(0.5, wires=[0, 1]) - qp.SingleExcitationMinus(0.5, wires=[0, 1]) - qp.DoubleExcitation(0.5, wires=[0, 1, 2, 3]) - return qp.expval(qp.Z(0)) - - # CHECK-DAG: @_cry(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CRY"} - # CHECK-DAG: @_s_phaseshift(%arg0: !qref.reg<4>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "S"} - # CHECK-DAG: @_phaseshift_to_rz_gp(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PhaseShift"} - # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - # CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # CHECK-DAG: @_doublexcit(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<4xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 4 : i64, target_gate = "DoubleExcitation"} - # CHECK-DAG: @_single_excitation_decomp(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "SingleExcitation"} - print(circuit_15.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decomposition_rule_name_update_multi_qubits)() - - -def test_decomposition_rule_name_adjoint(): - """Test decomposition rule with qp.adjoint.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RY", "RX", "CZ", "GlobalPhase", "Adjoint(SingleExcitation)"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - def circuit_16(x: float): - # CHECK: qref.adjoint { - # CHECK: qref.adjoint { - # CHECK: qref.adjoint { - # CHECK: qref.adjoint { - qp.adjoint(qp.CNOT)(wires=[0, 1]) - qp.adjoint(qp.Hadamard)(wires=2) - qp.adjoint(qp.RZ)(0.5, wires=3) - qp.adjoint(qp.SingleExcitation)(0.1, wires=[0, 1]) - qp.adjoint(qp.SingleExcitation(x, wires=[0, 1])) - return qp.expval(qp.Z(0)) - - # CHECK-DAG: @_single_excitation_decomp(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "SingleExcitation"} - # CHECK-DAG: @_hadamard_to_rz_ry(%arg0: !qref.reg<4>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Hadamard"} - # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - # CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # CHECK-DAG: @_cnot_to_cz_h(%arg0: !qref.reg<4>, %arg1: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CNOT"} - print(circuit_16.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decomposition_rule_name_adjoint)() - - -# TODO: Reenable this once the underlying non-determinism issue is resolved -def test_decomposition_rule_name_ctrl(): - """Test decomposition rule with qp.ctrl.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RX", "RZ", "H", "CZ"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - # SKIP-CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # SKIP-CHECK{LITERAL}: @circuit_17() -> tensor attributes {decompose_gatesets - def circuit_17(): - # SKIP-CHECK: %out_qubits:2 = quantum.custom "CRY"(%cst) %1, %2 : !quantum.bit, !quantum.bit - # SKIP-CHECK-NEXT: %out_qubits_0:2 = quantum.custom "CNOT"() %out_qubits#0, %out_qubits#1 : !quantum.bit, !quantum.bit - qp.ctrl(qp.RY, control=0)(0.5, 1) - qp.ctrl(qp.PauliX, control=0)(1) - return qp.expval(qp.Z(0)) - - # SKIP-CHECK-DAG: @_cnot_to_cz_h(%arg0: !quantum.reg, %arg1: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CNOT"} - # SKIP-CHECK-DAG: @_cry(%arg0: !quantum.reg, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CRY"} - # SKIP-CHECK-DAG: @_ry_to_rz_rx(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RY"} - # SKIP-CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # print(circuit_17.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decomposition_rule_name_ctrl)() - - -# TODO: Reenable this once the underlying non-determinism issue is resolved -def test_qft_decomposition(): - """Test the decomposition of the QFT""" - - qp.decomposition.enable_graph() - - @qp.qjit(autograph=True, target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RX", "RY", "CNOT", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - # SKIP-CHECK: %0 = transform.apply_registered_pass "decompose-lowering" - # SKIP-CHECK: @circuit_18(%arg0: tensor<3xf64>) -> tensor attributes {decompose_gatesets - def circuit_18(): - # %6 = scf.for %arg1 = %c0 to %c4 step %c1 iter_args(%arg2 = %0) -> (!quantum.reg) { - # %23 = scf.for %arg3 = %c0 to %22 step %c1 iter_args(%arg4 = %21) -> (!quantum.reg) { - # %7 = scf.for %arg1 = %c0 to %c2 step %c1 iter_args(%arg2 = %6) -> (!quantum.reg) { - qp.QFT(wires=[0, 1, 2, 3]) - return qp.expval(qp.Z(0)) - - # SKIP-CHECK-DAG: @ag___cphase_to_rz_cnot(%arg0: !quantum.reg, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "ControlledPhaseShift"} - # SKIP-CHECK-DAG: @ag___rz_to_ry_rx(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - # SKIP-CHECK-DAG: @ag___rot_to_rz_ry_rz(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # SKIP-CHECK-DAG: @ag___swap_to_cnot(%arg0: !quantum.reg, %arg1: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "SWAP"} - # SKIP-CHECK-DAG: @ag___hadamard_to_rz_ry(%arg0: !quantum.reg, %arg1: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Hadamard"} - # print(circuit_18.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_qft_decomposition)() - - -def test_decompose_lowering_with_other_passes(): - """Test the decompose lowering pass with other passes in a pass pipeline.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @qp.transforms.merge_rotations - @qp.transforms.cancel_inverses - @partial( - qp.transforms.decompose, - gate_set={"RZ", "RY", "CNOT", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module attributes {transform.with_named_sequence} { - # CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.op<"builtin.module">) { - # CHECK-NEXT: [[ONE:%.+]] = transform.apply_registered_pass "decompose-lowering" to %arg0 : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: [[TWO:%.+]] = transform.apply_registered_pass "cancel-inverses" to [[ONE]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to [[TWO]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: transform.yield - # CHECK-NEXT: } - def circuit_19(): - - # CHECK: [[QREG:%.+]] = qref.alloc( 1) : !qref.reg<1> - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "PauliX"() [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "PauliX"() [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "RX"({{%.+}}) [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "RX"({{%.+}}) [[QUBIT]] : !qref.bit - qp.PauliX(0) - qp.PauliX(0) - qp.RX(0.1, wires=0) - qp.RX(-0.1, wires=0) - return qp.expval(qp.PauliX(0)) - - # CHECK-DAG: @_paulix_to_rx(%arg0: !qref.reg<1>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PauliX"} - # CHECK-DAG: @_rx_to_rz_ry(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RX"} - print(circuit_19.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decompose_lowering_with_other_passes)() - - -def test_decompose_lowering_multirz(): - """Test the decompose lowering pass with MultiRZ in the gate set.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"CNOT", "RZ"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - # CHECK: %0 = transform.apply_registered_pass "decompose-lowering" - def circuit_20(x: float): - # CHECK: [[QREG:%.+]] = qref.alloc( 3) : !qref.reg<3> - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor - # CHECK: qref.multirz([[angle]]) [[q0]] : !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: [[q1:%.+]] = qref.get [[QREG]][ 1] : !qref.reg<3> -> !qref.bit - # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor - # CHECK: qref.multirz([[angle]]) [[q0]], [[q1]] : !qref.bit, !qref.bit - # CHECK: [[q1:%.+]] = qref.get [[QREG]][ 1] : !qref.reg<3> -> !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: [[q2:%.+]] = qref.get [[QREG]][ 2] : !qref.reg<3> -> !qref.bit - # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor - # CHECK: qref.multirz([[angle]]) [[q1]], [[q0]], [[q2]] : !qref.bit, !qref.bit, !qref.bit - qp.MultiRZ(x, wires=[0]) - qp.MultiRZ(x, wires=[0, 1]) - qp.MultiRZ(x, wires=[1, 0, 2]) - return qp.expval(qp.PauliX(0)) - - # CHECK-DAG: @_multi_rz_decomposition_wires_1(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "MultiRZ"} - # CHECK-DAG: @_multi_rz_decomposition_wires_2(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "MultiRZ"} - # CHECK-DAG: @_multi_rz_decomposition_wires_3(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<3xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 3 : i64, target_gate = "MultiRZ"} - # CHECK-DAG: scf.for %arg3 = %c0 to %c2 step %c1 - # CHECK-DAG: scf.for %arg3 = %c1 to %c3 step %c1 - print(circuit_20.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_lowering_multirz() - - -def test_decompose_lowering_with_ordered_passes(): - """Test the decompose lowering pass with other passes in a specific order in a pass pipeline.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RZ", "RY", "CNOT", "GlobalPhase"}, - ) - @qp.transforms.merge_rotations - @qp.transforms.cancel_inverses - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module attributes {transform.with_named_sequence} { - # CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.op<"builtin.module">) { - # CHECK-NEXT: [[FIRST:%.+]] = transform.apply_registered_pass "cancel-inverses" to %arg0 : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: [[SECOND:%.+]] = transform.apply_registered_pass "merge-rotations" to [[FIRST]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "decompose-lowering" to [[SECOND]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: transform.yield - # CHECK-NEXT: } - def circuit_21(x: float): - # CHECK: [[QREG:%.+]] = qref.alloc( 1) : !qref.reg<1> - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "PauliX"() [[q0]] : !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "PauliX"() [[q0]] : !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor - # CHECK: qref.custom "RX"([[angle]]) [[q0]] : !qref.bit - # CHECK: [[negated:%.+]] = stablehlo.negate %arg0 : tensor - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: [[neg_angle:%.+]] = tensor.extract [[negated]][] : tensor - # CHECK: qref.custom "RX"([[neg_angle]]) [[q0]] : !qref.bit - qp.PauliX(0) - qp.PauliX(0) - qp.RX(x, wires=0) - qp.RX(-x, wires=0) - return qp.expval(qp.PauliX(0)) - - # CHECK-DAG: @_paulix_to_rx(%arg0: !qref.reg<1>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PauliX"} - # CHECK-DAG: @_rx_to_rz_ry(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RX"} - # CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - print(circuit_21.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decompose_lowering_with_ordered_passes)() - - -def test_decompose_lowering_with_gphase(): - """Test the decompose lowering pass with GlobalPhase.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RX", "RY", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - # CHECK: %0 = transform.apply_registered_pass "decompose-lowering" - def circuit_22(): - # CHECK: [[QREG:%.+]] = qref.alloc( 3) : !qref.reg<3> - # CHECK: qref.gphase({{%.+}}) - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: qref.custom "PhaseShift"({{%.+}}) [[q0]] : !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: qref.custom "PhaseShift"({{%.+}}) [[q0]] : !qref.bit - - qp.GlobalPhase(0.5) - qp.ctrl(qp.GlobalPhase, control=0)(0.3) - qp.ctrl(qp.GlobalPhase, control=0)(phi=0.3, wires=[1, 2]) - return qp.expval(qp.PauliX(0)) - - # CHECK-DAG: @_phaseshift_to_rz_gp(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PhaseShift"} - # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - print(circuit_22.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decompose_lowering_with_gphase)() - - -def test_decompose_lowering_alt_decomps(): - """Test the decompose lowering pass with alternative decompositions.""" - - qp.decomposition.enable_graph() - - @qp.register_resources({qp.RY: 1}) - def custom_rot_cheap(params, wires: WiresLike): - qp.RY(params[1], wires=wires) - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RY", "RZ"}, - alt_decomps={qp.Rot: [custom_rot_cheap]}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3), shots=1000) - def circuit_23(x: float, y: float): - qp.Rot(x, y, x + y, wires=1) - return qp.expval(qp.PauliZ(0)) - - # CHECK-DAG: @custom_rot_cheap(%arg0: !qref.reg<3>, %arg1: tensor<3xf64>, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - print(circuit_23.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_lowering_alt_decomps() - - -def test_decompose_lowering_with_tensorlike(): - """Test the decompose lowering pass with fixed decompositions - using TensorLike parameters.""" - - qp.decomposition.enable_graph() - - @qp.register_resources({qp.RZ: 2, qp.RY: 1}) - def custom_rot(params: TensorLike, wires: WiresLike): - qp.RZ(params[0], wires=wires) - qp.RY(params[1], wires=wires) - qp.RZ(params[2], wires=wires) - - @qp.register_resources({qp.RZ: 1, qp.CNOT: 4}) - def custom_multirz(params: TensorLike, wires: WiresLike): - qp.CNOT(wires=(wires[2], wires[1])) - qp.CNOT(wires=(wires[1], wires[0])) - qp.RZ(params[0], wires=wires[0]) - qp.CNOT(wires=(wires[1], wires[0])) - qp.CNOT(wires=(wires[2], wires[1])) - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RY", "RX", qp.CNOT}, - fixed_decomps={qp.Rot: custom_rot, qp.MultiRZ: custom_multirz}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3), shots=1000) - def circuit_24(x: float, y: float): - qp.Rot(x, y, x + y, wires=1) - qp.MultiRZ(x + y, wires=[0, 1, 2]) - return qp.expval(qp.PauliZ(0)) - - # CHECK-DAG: @custom_multirz_wires_3(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<3xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 3 : i64, target_gate = "MultiRZ"} - # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - # CHECK-DAG: @custom_rot(%arg0: !qref.reg<3>, %arg1: tensor<3xf64>, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - print(circuit_24.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decompose_lowering_with_tensorlike)() - - -def test_decompose_lowering_fallback(): - """Test the decompose lowering pass when the graph is failed.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.RX, qp.RZ}) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - # CHECK-LABEL: @circuit_25() - def circuit_25(): - # CHECK: [[pi_over_2:%.+]] = arith.constant 1.5707963267948966 : f64 - # CHECK: [[QREG:%.+]] = qref.alloc( 2) : !qref.reg<2> - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<2> -> !qref.bit - # CHECK: qref.custom "RZ"([[pi_over_2]]) [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<2> -> !qref.bit - # CHECK: qref.custom "RX"([[pi_over_2]]) [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<2> -> !qref.bit - # CHECK: qref.custom "RZ"([[pi_over_2]]) [[QUBIT]] : !qref.bit - qp.Hadamard(0) - return qp.state() - - print(circuit_25.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_lowering_fallback() - - -def test_decompose_lowering_params_ordering(): - """Test the order of params and wires in the captured decomposition rule.""" - - qp.decomposition.enable_graph() - - @qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set=[qp.RX, qp.RY, qp.RZ]) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - # CHECK-LABEL: @circuit_26(%arg0: tensor, %arg1: tensor, %arg2: tensor) - def circuit_26(x: float, y: float, z: float): - qp.Rot(x, y, z, wires=0) - return qp.expval(qp.PauliZ(0)) - - # CHECK-LABEL: @_rot_to_rz_ry_rz(%arg0: !qref.reg<2>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # CHECK: [[EXTRACTED_1:%.+]] = tensor.extract %arg1[] : tensor - # CHECK-NEXT: qref.custom "RZ"([[EXTRACTED_1]]) {{%.+}} : !qref.bit - # CHECK: [[EXTRACTED_2:%.+]] = tensor.extract %arg2[] : tensor - # CHECK-NEXT: qref.custom "RY"([[EXTRACTED_2]]) {{%.+}} : !qref.bit - # CHECK: [[EXTRACTED_3:%.+]] = tensor.extract %arg3[] : tensor - # CHECK-NEXT: qref.custom "RZ"([[EXTRACTED_3]]) {{%.+}} : !qref.bit - # CHECK: return - print(circuit_26.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_lowering_params_ordering() - - -def test_decomposition_rule_with_allocation(): - """Test decomposition rule with dynamic qubit allocation""" - - @decomposition_rule(is_qreg=True) - def Hadamard0_with_alloc(wire: WiresLike): - with qp.allocate(1) as q: - qp.X(q[0]) - qp.CNOT(wires=[q[0], wire]) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module @circuit_27 - def circuit_27(): - Hadamard0_with_alloc(int) - return qp.probs() - - # CHECK-LABEL: @Hadamard0_with_alloc(%arg0: !qref.reg<1>, %arg1: tensor) - # CHECK: [[dynalloc_qreg:%.+]] = qref.alloc( 1) - # CHECK: [[dynalloc_bit0:%.+]] = qref.get [[dynalloc_qreg]][ 0] - # CHECK: qref.custom "PauliX"() [[dynalloc_bit0]] - # CHECK: [[detensor:%.+]] = tensor.extract %arg1[] - # CHECK: [[glob_bit:%.+]] = qref.get %arg0[[[detensor]]] - # CHECK: qref.custom "CNOT"() [[dynalloc_bit0]], [[glob_bit]] - # CHECK: qref.dealloc [[dynalloc_qreg]] - # CHECK: return - - print(circuit_27.mlir) - - -test_decomposition_rule_with_allocation() - - -def test_decompose_autograph_multi_blocks(): - """Test the decompose lowering pass with autograph in the program and rule.""" - - qp.decomposition.enable_graph() - - def _multi_rz_decomposition_resources(num_wires): - """Resources required for MultiRZ decomposition.""" - return {qp.RZ: 1, qp.CNOT: 2 * (num_wires - 1)} - - @qp.register_resources(_multi_rz_decomposition_resources) - @qp.capture.run_autograph - def _multi_rz_decomposition(theta: TensorLike, wires: WiresLike, **__): - """Decomposition of MultiRZ using CNOTs and RZs.""" - for i in range(len(wires) - 1): - qp.CNOT(wires=(wires[i], wires[i + 1])) - qp.RZ(theta, wires=wires[0]) - for i in range(len(wires) - 1, 0, -1): - qp.CNOT(wires=(wires[i], wires[i - 1])) - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RZ", "CNOT"}, - fixed_decomps={qp.MultiRZ: _multi_rz_decomposition}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=5)) - def circuit_29(n: int): - - # CHECK: scf.for %arg1 = {{%.+}} to {{%.+}} step {{%.+}} { - @qp.for_loop(n) - def f(i): # pylint: disable=unused-argument - qp.MultiRZ(0.5, wires=[0, 1, 2, 3, 4]) - - f() # pylint: disable=no-value-for-parameter - - return qp.expval(qp.Z(0)) - - # CHECK-LABEL: @ag___multi_rz_decomposition_wires_5(%arg0: !qref.reg<5>, %arg1: tensor<1xf64>, %arg2: tensor<5xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 5 : i64, target_gate = "MultiRZ"} - # CHECK: scf.for %arg3 = {{%.+}} to {{%.+}} step {{%.+}} { - # CHECK: scf.for %arg3 = {{%.+}} to {{%.+}} step {{%.+}} { - print(circuit_29.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_autograph_multi_blocks() - - -def test_decompose_work_wires_context_manager(): - """ - Test that decomposition with work wires is correctly applied when allocating with the context - manager. - """ - - @decomposition_rule(is_qreg=True, op_type="PauliZ") - def my_decomp(wires): - with qp.allocate(2, restored=False) as work_wires: - qp.X(wires[0]) - qp.X(wires[1]) - qp.H(work_wires[0]) - qp.H(work_wires[1]) - - @qp.qjit(capture=True) - @qp.transform(pass_name="decompose-lowering") - @qp.qnode(qp.device("lightning.qubit", wires=3)) - def my_circuit(): - my_decomp(jax.core.ShapedArray((2,), int)) - qp.Z(0) - return qp.probs() - - # check that decomp arrives properly - # CHECK-LABEL: @my_decomp({{.*}}) attributes {{{.*}} target_gate = "PauliZ"} - print(my_circuit.mlir) - - # check that decomp is applied properly - # CHECK-NOT: PauliZ - # CHECK-NOT: my_decomp - - # two allocates, one for main register and one for decomp register - # CHECK: allocate - # CHECK: allocate - # CHECK: PauliX - # CHECK: PauliX - # CHECK: Hadamard - # CHECK: Hadamard - # CHECK: release - # CHECK: release - print(my_circuit.mlir_opt) - - -test_decompose_work_wires_context_manager() - - -def test_decompose_work_wires_alloc_dealloc(): - """ - Test that decomposition with work wires is correctly applied when allocating/deallocating - explicitly. - """ - - @decomposition_rule(is_qreg=True, op_type="RY") - def my_decomp(angle, wires): - work_wires = qp.allocate(2) - qp.CNOT((work_wires[0], wires[0])) - qp.RX(-np.pi / 2, wires[0]) - qp.RZ(angle, wires[0]) - qp.RX(np.pi / 2, wires[0]) - qp.CNOT((work_wires[1], wires[1])) - qp.deallocate(work_wires) - - @qp.qjit(capture=True) - @qp.transform(pass_name="decompose-lowering") - @qp.qnode(qp.device("lightning.qubit", wires=3)) - def my_circuit(angle: float): - my_decomp(float, jax.core.ShapedArray((2,), int)) - qp.RY(angle, 0) - return qp.probs() - - # check that decomp arrives properly - # CHECK-LABEL: @my_decomp({{.*}}) attributes {{{.*}} target_gate = "RY"} - print(my_circuit.mlir) - - # check that the decomposition applies properly - # CHECK-NOT: my_decomp - # CHECK-NOT: RY - - # two allocates, one for main register and one for decomp register - # CHECK: allocate - # CHECK: allocate - # CHECK: CNOT - # CHECK: RX - # CHECK: RZ - # CHECK: RX - # CHECK: CNOT - # CHECK: release - # CHECK: release - print(my_circuit.mlir_opt) - - -test_decompose_work_wires_alloc_dealloc() - - -def test_decompose_work_wires_control_flow(): - """Test that decomposition with work wires + control flow is correctly applied.""" - - @decomposition_rule(is_qreg=True, op_type="CRX") - def my_decomp(angle, wires, **_): - def true_func(): - qp.CNOT(wires) - - with qp.allocate(2, state="any", restored=True) as w: - for _ in range(2): - qp.H(w[0]) - qp.X(w[1]) - - def false_func(): - with qp.allocate(1, state="any", restored=False) as w: - qp.H(w) - - m = qp.measure(wires[0]) - - qp.cond(m, qp.CNOT)(wires) - - qp.cond(angle > 1.2, true_func, false_func)() - - @qp.qjit(capture=True) - @qp.transform(pass_name="decompose-lowering") - @qp.qnode(qp.device("lightning.qubit", wires=4)) - def circuit(): - my_decomp(float, jax.core.ShapedArray((2,), int)) - qp.CRX(1.7, wires=[0, 1]) - qp.CRX(-7.2, wires=[0, 1]) - return qp.state() - - # target_gate attribute is correctly applied - # CHECK: my_decomp([[args:.*]]) attributes {[[other_attributes:.*]] target_gate = "CRX"} - print(circuit.mlir) - - # test that the decomposition is applied correctly - # CHECK-NOT: CRX - # CHECK-NOT: my_decomp - - # allocate for main register, subsequent allocates+releases for decomp registers - # CHECK: allocate - - # first CRX: true branch - # CHECK: CNOT - # CHECK: allocate - # CHECK: Hadamard - # CHECK: PauliX - # CHECK: Hadamard - # CHECK: PauliX - # CHECK: release - - # second CRX: false branch - # CHECK: allocate - # CHECK: Hadamard - # CHECK: Measure - # CHECK: cond - # CHECK: CNOT - # CHECK: release - - # release main register - # CHECK: release - - print(circuit.mlir_opt) - - -test_decompose_work_wires_control_flow() - - -def test_decompose_work_wires_with_decompose_transform(): - """Test that work wires are correctly lowered and decomposed by the decompose transform.""" - - qp.decomposition.enable_graph() - - @qp.register_resources({qp.X: 1, qp.Z: 1}) - def my_decomp(wire): - with qp.allocate(1) as work_wire: - qp.X(work_wire) - qp.Z(wire) - qp.X(work_wire) - - @qjit(capture=True) - @partial( - qp.transforms.decompose, - gate_set={ - "X", - "Z", - }, - fixed_decomps={ - qp.Y: my_decomp, - }, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def my_circuit(): - qp.Y(0) - return qp.probs() - - # CHECK-NOT: Y - # CHECK-NOT: my_decomp - - # two allocates, one for main register and one for decomp register - # CHECK: allocate - # CHECK: allocate - # CHECK: X - # CHECK: Z - # CHECK: X - # CHECK: release - # CHECK: release - print(my_circuit.mlir_opt) - - qp.decomposition.disable_graph() - - -test_decompose_work_wires_with_decompose_transform() - - -def test_num_work_wires(): - """Test that num_work_wires can be passed and is correctly used in solving the graph.""" - - qp.decomposition.enable_graph() - - @qp.register_resources( - {qp.CNOT: 3, qp.H: 1, qp.X: 1, qp.ops.op_math.Conditional: 2}, - work_wires={"borrowed": 2, "garbage": 1}, - ) - def my_decomp(angle, wires, **_): - def true_func(): - qp.CNOT(wires) - - with qp.allocate(2, state="any", restored=True) as w: - qp.H(w[0]) - qp.H(w[0]) - qp.X(w[1]) - qp.X(w[1]) - - return - - def false_func(): - with qp.allocate(1, state="any", restored=False) as w: - qp.H(w) - - m = qp.measure(wires[0]) - - qp.cond(m, qp.CNOT)(wires) - - return - - qp.cond(angle > 1.2, true_func, false_func)() - - @qp.qjit(capture=True) - @partial( - qp.transforms.decompose, - gate_set={qp.CNOT, qp.H, qp.X, "Conditional", "MidMeasure"}, - fixed_decomps={qp.CRX: my_decomp}, - num_work_wires=3, - ) - @qp.qnode(qp.device("lightning.qubit", wires=5)) - def circuit(): - qp.CRX(1.7, wires=[0, 1]) - qp.CRX(-7.2, wires=[0, 1]) - return qp.state() - - # CHECK-NOT: CRX - # CHECK-NOT: my_decomp - - # CHECK: allocate - # CHECK: allocate - # CHECK: CNOT - # CHECK: Hadamard - # CHECK: Hadamard - # CHECK: PauliX - # CHECK: PauliX - # CHECK: Hadamard - # CHECK: Measure - # CHECK: CNOT - # CHECK: release - # CHECK: release - print(circuit.mlir_opt) - - qp.decomposition.disable_graph() - - -test_num_work_wires() - - -def test_default_decomps(): - """Test that default decompositions are correctly applied with qjit.""" - qp.decomposition.enable_graph() - - # Toffoli's decomposition to this gateset includes a wire allocation - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={qp.ops.ChangeOpBasis}, - num_work_wires=1, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - def circuit(): - qp.Toffoli(wires=[0, 1, 2]) - return qp.state() - - # CHECK-NOT: toffoli_elbow - # CHECK-NOT: Toffoli - - # two allocates/releases, for default register + work wires - # CHECK: allocate - # CHECK: allocate - # CHECK: TemporaryAND - # CHECK: release - # CHECK: release - print(circuit.mlir_opt) - - qp.decomposition.disable_graph() - - -test_default_decomps() - - -def test_graph_decomp_registered(): - """Test that the `graph_decomposition` pass is registered correctly.""" - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" - @graph_decomposition(gate_set={qp.RX}) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def catalyst_circuit(): - return - - print(catalyst_circuit.mlir) - - my_transform = qp.transform(pass_name="graph-decomposition") - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" - @my_transform(gate_set=["RX"]) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def pennylane_circuit(): - return - - print(pennylane_circuit.mlir) - - -test_graph_decomp_registered() - - -def test_cpp_decomp_args(): - """Test that the `graph_decomposition` pass lowers arguments to mlir correctly.""" - - def x_to_rx(wire): - qp.RX(np.pi, wire) - - def y_to_ry(wire): - qp.RY(np.pi, wire) - - def h_to_rx_ry(wire): - qp.RX(np.pi / 2, wire) - qp.RY(np.pi / 2, wire) - - @qjit(target="mlir") - # CHECK: "graph-decomposition" with options = { - # CHECK-DAG: "gate-set" = {Hadamard = 1.000000e+00 : f64, RX = 1.000000e+00 : f64, RY = 1.000000e+00 : f64} - # CHECK-DAG: "fixed-decomps" = {PauliX = "x_to_rx", PauliY = "y_to_ry"} - # CHECK-DAG: "alt-decomps" = {Hadamard = ["h_to_rx_ry"]} - # CHECK-DAG: "bytecode-rules" = "{{.*}}decomposition_rules_{{.*}}.mlirbc" - # CHECK: } to {{%.+}} : (!transform.op<"builtin.module">) - @graph_decomposition( - gate_set={qp.RX, qp.H, qp.RY}, - fixed_decomps={qp.X: x_to_rx, qp.Y: y_to_ry}, - alt_decomps={qp.H: [h_to_rx_ry]}, - _builtin_rule_path="/decomp_rules.mlirbc", - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - return - - print(circuit.mlir) - - -test_cpp_decomp_args() - - -def test_cpp_decomp_empty_args(): - """ - Test that the `graph_decomposition` pass correctly handled arg lowering when no values are - supplied. - """ - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" - # CHECK-NOT: fixed-decomps - # CHECK-NOT: alt-decomps - # CHECK: "bytecode-rules" = "{{.*}}/decomposition_rules{{.*}}.mlirbc" - @graph_decomposition(gate_set={qp.RX}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit(): - return - - print(circuit.mlir) - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" - # CHECK-NOT: fixed-decomps - # CHECK-NOT: alt-decomps - # CHECK: "bytecode-rules" = "{{.*}}/decomposition_rules{{.*}}.mlirbc" - @graph_decomposition(gate_set={qp.RX}, fixed_decomps={}, alt_decomps={}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit2(): - return - - print(circuit2.mlir) - - -test_cpp_decomp_empty_args() - - -def test_cpp_decomp_string_op_names(): - """Test that cpp decomp args work with string op names.""" - - def y_to_xz(wires): - qp.RX(np.pi, wires) - qp.RZ(np.pi, wires) - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" with options = { - # CHECK-DAG: "fixed-decomps" = {PauliX = "{{.*}}", PauliZ = "{{.*}}"} - # CHECK-DAG: "alt-decomps" = {PauliY = ["{{.*}}", "y_to_xz"]} - # CHECK: } to {{%.+}} : (!transform.op<"builtin.module">) - @graph_decomposition( - gate_set={"RX", "RY", "RZ"}, - fixed_decomps={ - "X": lambda wires: qp.RX(np.pi, wires), - "PauliZ": lambda wires: qp.RZ(np.pi, wires), - }, - alt_decomps={ - "PauliY": [ - lambda wires: qp.RY(np.pi, wires), - y_to_xz, - ] - }, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - return - - print(circuit.mlir) - - -test_cpp_decomp_string_op_names() - - -def test_cpp_decomp_builtin_rules(): - """Test that cpp decomp applies builtin rules.""" - - @qjit(target="mlir", capture=True) - @graph_decomposition( - gate_set={qp.RX, qp.RY, qp.RZ, qp.GlobalPhase}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - # CHECK-NOT: PauliX - # CHECK-NOT: PauliY - # CHECK-NOT: PauliZ - # CHECK-DAG: RX - # CHECK-DAG: RY - # CHECK-DAG: RZ - qp.X(0) - qp.Y(1) - qp.Z(0) - return qp.probs() - - print(circuit.mlir_opt) - - -test_cpp_decomp_builtin_rules() - - -def test_cpp_decomp_user_rules(): - """Test that cpp decomp applies user rules.""" - - @decomposition_rule(is_qreg=True, op_type="PauliY") - def y_to_rx(wire): - qp.RX(np.pi, wire) - - @decomposition_rule(is_qreg=True, op_type="PauliZ") - def z_to_rx(wire): - qp.RX(np.pi, wire) - - @qp.qjit(target="mlir", capture=True) - @graph_decomposition( - gate_set={qp.RX}, fixed_decomps={qp.Y: y_to_rx}, alt_decomps={qp.Z: [z_to_rx]} - ) - @qp.qnode(qp.device("null.qubit", wires=1)) - def circuit(): - y_to_rx(jax.core.ShapedArray((1,), int)) - z_to_rx(jax.core.ShapedArray((1,), int)) - # CHECK-NOT: PauliY - # CHECK-NOT: PauliZ - # CHECK: RX - # CHECK: RX - # CHECK: return - qp.Y(0) - qp.Z(0) - return qp.probs() - - print(circuit.mlir_opt) - - -test_cpp_decomp_user_rules() - - -def test_cpp_decomp_user_rule_cleanup(): - """Test that user rules do not pollute the IR after the quantum compilation stage.""" - - @decomposition_rule(is_qreg=True, op_type="PauliX") - def x_to_h(wire): - return qp.H(wire) - - @qjit(capture=True) - @graph_decomposition(gate_set={qp.H}, fixed_decomps={qp.X: x_to_h}) - @qp.qnode(qp.device("null.qubit", wires=1)) - def circuit(): - # CHECK-NOT: PauliX - # CHECK-NOT: x_to_h - x_to_h(jax.core.ShapedArray((1,), int)) - qp.X(0) - - print(circuit.mlir_opt) - - -test_cpp_decomp_user_rule_cleanup() - - -def test_paulirot_python_decomposition(): - """Test that paulirots are decomposed by the mlir graph.""" - - @qjit(capture=True) - @graph_decomposition(gate_set={qp.H, qp.MultiRZ, qp.GlobalPhase, qp.RX}) - @qp.qnode(qp.device("null.qubit", wires=4)) - def circuit(): - qp.PauliRot(0.9, "XZXY", [0, 1, 2, 3]) - return qp.probs() - - print(circuit.mlir_opt) - - # CHECK-NOT: quantum.paulirot - # CHECK: Hadamard - # CHECK: Hadamard - # CHECK: RX - # CHECK: MultiRZ - # CHECK: Hadamard - # CHECK: Hadamard - # CHECK: RX - - -test_paulirot_python_decomposition() From ff30d1d3524f7ab9712c4ae43ba7fe2d5aa9dcf2 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 23 Jul 2026 14:47:32 -0400 Subject: [PATCH 18/36] . --- frontend/test/lit/_old_test_decomposition.py | 1863 ------------------ 1 file changed, 1863 deletions(-) delete mode 100644 frontend/test/lit/_old_test_decomposition.py diff --git a/frontend/test/lit/_old_test_decomposition.py b/frontend/test/lit/_old_test_decomposition.py deleted file mode 100644 index e8f9448e7c..0000000000 --- a/frontend/test/lit/_old_test_decomposition.py +++ /dev/null @@ -1,1863 +0,0 @@ -# Copyright 2022-2025 Xanadu Quantum Technologies Inc. -import os -import pathlib -import platform -from copy import deepcopy -from functools import partial - -import jax -import numpy as np -import pennylane as qp -from pennylane.devices.capabilities import OperatorProperties -from pennylane.typing import TensorLike -from pennylane.wires import WiresLike - -from catalyst import measure, qjit -from catalyst.compiler import get_lib_path -from catalyst.device import get_device_capabilities -from catalyst.jax_primitives import decomposition_rule -from catalyst.passes import graph_decomposition - -# 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. - -# RUN: %PYTHON %s | FileCheck %s -# pylint: disable=line-too-long -# pylint: disable=too-many-lines - - -# Helper to skip tests that fail due to PauliRot type annotation issue -# TODO: Remove this once PennyLane fixes the PauliRot decomposition type annotations -def skip_if_pauli_rot_issue(test_func): - """Wrapper to skip tests that fail due to PauliRot type annotation issues.""" - - def wrapper(): - try: - test_func() - except (ValueError, IndexError) as e: - error_msg = str(e) - if ( - "Unsupported type annotation None for parameter pauli_word" in error_msg - or "Unsupported type annotation for parameter pauli_word" in error_msg - or "index is out of bounds for axis" in error_msg - ): - print(f"# SKIPPED {test_func.__name__}: PauliRot type annotation issue") - else: - raise - - return wrapper - - -TEST_PATH = os.path.dirname(__file__) -CONFIG_CUSTOM_DEVICE = pathlib.Path(f"{TEST_PATH}/../custom_device/custom_device.toml") - - -def get_custom_device_without(num_wires, discards=frozenset(), force_matrix=frozenset()): - """Generate a custom device without gates in discards.""" - - class CustomDevice(qp.devices.Device): - """Custom Gate Set Device""" - - name = "Custom Device" - config_filepath = CONFIG_CUSTOM_DEVICE - - _to_matrix_ops = {} - - def __init__(self, wires=None): - super().__init__(wires=wires) - self.qjit_capabilities = deepcopy(get_device_capabilities(self)) - for gate in discards: - self.qjit_capabilities.operations.pop(gate, None) - for gate in force_matrix: - self.qjit_capabilities.operations.pop(gate, None) - self._to_matrix_ops[gate] = OperatorProperties(False, False, False) - - def apply(self, operations, **kwargs): - """Unused""" - raise RuntimeError("Only C/C++ interface is defined") - - @staticmethod - def get_c_interface(): - """Returns a tuple consisting of the device name, and - the location to the shared object with the C/C++ device implementation. - """ - system_extension = ".dylib" if platform.system() == "Darwin" else ".so" - lib_path = ( - get_lib_path("runtime", "RUNTIME_LIB_DIR") + "/librtd_null_qubit" + system_extension - ) - return "NullQubit", lib_path - - def execute(self, circuits, execution_config): - """Execution.""" - return circuits, execution_config - - return CustomDevice(wires=num_wires) - - -def test_decompose_multicontrolledx(): - """Test decomposition of MultiControlledX as an aliased gate.""" - dev = get_custom_device_without(5, discards={"MultiControlledX"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_multicontrolled_x1 - def decompose_multicontrolled_x1(theta: float): - qp.RX(theta, wires=[0]) - # CHECK-NOT: name = "MultiControlledX" - # CHECK: quantum.custom "PauliX"() {{%[a-zA-Z0-9_]+}} ctrls({{%[a-zA-Z0-9_]+}}, {{%[a-zA-Z0-9_]+}}, {{%[a-zA-Z0-9_]+}}) - # CHECK-NOT: name = "MultiControlledX" - qp.MultiControlledX(wires=[0, 1, 2, 3]) - return qp.state() - - print(decompose_multicontrolled_x1.mlir) - - -test_decompose_multicontrolledx() - - -def test_decompose_rot(): - """Test decomposition of Rot gate.""" - dev = get_custom_device_without(1, discards={"Rot", "C(Rot)"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_rot - def decompose_rot(phi: float, theta: float, omega: float): - # CHECK-NOT: name = "Rot" - # CHECK: [[phi:%.+]] = tensor.extract %arg0 - # CHECK-NOT: name = "Rot" - # CHECK: {{%.+}} = quantum.custom "RZ"([[phi]]) - # CHECK-NOT: name = "Rot" - # CHECK: [[theta:%.+]] = tensor.extract %arg1 - # CHECK-NOT: name = "Rot" - # CHECK: {{%.+}} = quantum.custom "RY"([[theta]]) - # CHECK-NOT: name = "Rot" - # CHECK: [[omega:%.+]] = tensor.extract %arg2 - # CHECK-NOT: name = "Rot" - # CHECK: {{%.+}} = quantum.custom "RZ"([[omega]]) - # CHECK-NOT: name = "Rot" - qp.Rot(phi, theta, omega, wires=0) - return measure(wires=0) - - print(decompose_rot.mlir) - - -test_decompose_rot() - - -def test_decompose_s(): - """Test decomposition of S gate.""" - dev = get_custom_device_without(1, discards={"S", "C(S)"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_s - def decompose_s(): - # CHECK-NOT: name="S" - # CHECK: [[pi_div_2:%.+]] = arith.constant 1.57079{{.+}} : f64 - # CHECK-NOT: name = "S" - # CHECK: {{%.+}} = quantum.custom "PhaseShift"([[pi_div_2]]) - # CHECK-NOT: name = "S" - qp.S(wires=0) - return measure(wires=0) - - print(decompose_s.mlir) - - -test_decompose_s() - - -def test_decompose_qubitunitary(): - """Test decomposition of QubitUnitary""" - dev = get_custom_device_without(1, discards={"QubitUnitary"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_qubit_unitary - def decompose_qubit_unitary(U: jax.core.ShapedArray([2, 2], float)): - # CHECK-NOT: name = "QubitUnitary" - # CHECK: quantum.custom "RZ" - # CHECK: quantum.custom "RY" - # CHECK: quantum.custom "RZ" - # CHECK-NOT: name = "QubitUnitary" - qp.QubitUnitary(U, wires=0) - return measure(wires=0) - - print(decompose_qubit_unitary.mlir) - - -test_decompose_qubitunitary() - - -def test_decompose_singleexcitation(): - """ - Test that single excitation is not decomposed. - """ - dev = get_custom_device_without(2) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_singleexcitation - def decompose_singleexcitation(theta: float): - # CHECK: quantum.custom "SingleExcitation" - - qp.SingleExcitation(theta, wires=[0, 1]) - return measure(wires=0) - - print(decompose_singleexcitation.mlir) - - -test_decompose_singleexcitation() - - -def test_decompose_doubleexcitation(): - """ - Test that Double excitation is not decomposed. - """ - dev = get_custom_device_without(4) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_doubleexcitation - def decompose_doubleexcitation(theta: float): - # CHECK: quantum.custom "DoubleExcitation" - - qp.DoubleExcitation(theta, wires=[0, 1, 2, 3]) - return measure(wires=0) - - print(decompose_doubleexcitation.mlir) - - -test_decompose_doubleexcitation() - - -def test_decompose_singleexcitationplus(): - """ - Test decomposition of single excitation plus. - See - https://github.com/PennyLaneAI/pennylane/blob/main/pennylane/ops/qubit/qchem_ops.py - for the decomposition of qp.SingleExcitationPlus - """ - dev = get_custom_device_without(2, discards={"SingleExcitationPlus", "C(SingleExcitationPlus)"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_singleexcitationplus - def decompose_singleexcitationplus(theta: float): - # CHECK-NOT: "SingleExcitationPlus" - # CHECK: quantum.custom "Hadamard" - # CHECK: quantum.custom "CNOT" - # CHECK: quantum.custom "RY" - # CHECK: quantum.custom "RY" - # CHECK: quantum.custom "CY" - # CHECK: quantum.custom "S" - # CHECK: quantum.custom "Hadamard" - # CHECK: quantum.custom "RZ" - # CHECK: quantum.custom "CNOT" - # CHECK: quantum.gphase - - qp.SingleExcitationPlus(theta, wires=[0, 1]) - return measure(wires=0) - - print(decompose_singleexcitationplus.mlir) - - -test_decompose_singleexcitationplus() - - -def test_decompose_to_matrix(): - """Test decomposition of QubitUnitary""" - dev = get_custom_device_without(1, force_matrix={"PauliY"}) - - @qjit(target="mlir") - @qp.qnode(dev) - # CHECK-LABEL: @jit_decompose_to_matrix - def decompose_to_matrix(): - # CHECK: quantum.custom "PauliX" - qp.PauliX(wires=0) - # CHECK: quantum.unitary - qp.PauliY(wires=0) - # CHECK: quantum.custom "PauliZ" - qp.PauliZ(wires=0) - return measure(wires=0) - - print(decompose_to_matrix.mlir) - - -test_decompose_to_matrix() - - -def test_decomposition_rule_lowering(): - """Test that decomposition rules are lowered to private functions.""" - - @decomposition_rule(is_qreg=True) - def my_decomp(): - return - - @qp.qjit(capture=True) - @qp.qnode(qp.device("null.qubit", wires=1)) - def circuit(): - # CHECK-LABEL: func.func private @my_decomp - my_decomp() - return - - print(circuit.mlir) - - -test_decomposition_rule_lowering() - - -def test_decomposition_rule_wire_param(): - """Test decomposition rule with passing a parameter that is a wire/integer""" - - @decomposition_rule(is_qreg=False) - def Hadamard0(wire: WiresLike): - qp.Hadamard(wire) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @circuit - def circuit(_: float): - # CHECK: @circuit([[ARG0:%.+]] - # CHECK: [[QREG:%.+]] = qref.alloc - Hadamard0(int) - return qp.probs() - - # CHECK: @Hadamard0([[QBIT:%.+]]: !qref.bit) - # CHECK-NEXT: qref.custom "Hadamard"() [[QBIT]] : !qref.bit - # CHECK-NEXT: return - - print(circuit.mlir) - - -test_decomposition_rule_wire_param() - - -def test_decomposition_rule_gate_param_param(): - """Test decomposition rule with passing a regular parameter""" - - @decomposition_rule(is_qreg=False, num_params=1) - def RX_on_wire_0(param: TensorLike, w0: WiresLike): - qp.RX(param, wires=w0) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module @circuit_2 - def circuit_2(_: float): - RX_on_wire_0(float, int) - return qp.probs() - - # CHECK: @RX_on_wire_0([[PARAM_TENSOR:%.+]]: tensor, [[QUBIT:%.+]]: !qref.bit) - # CHECK-NEXT: [[PARAM:%.+]] = tensor.extract [[PARAM_TENSOR]][] : tensor - # CHECK-NEXT: qref.custom "RX"([[PARAM]]) [[QUBIT]] : !qref.bit - # CHECK-NEXT: return - print(circuit_2.mlir) - - -test_decomposition_rule_gate_param_param() - - -def test_multiple_decomposition_rules(): - """Test with multiple decomposition rules""" - - @decomposition_rule - def identity(): ... - - @decomposition_rule(is_qreg=True) - def all_wires_rx(param: TensorLike, w0: WiresLike, w1: WiresLike, w2: WiresLike): - qp.RX(param, wires=w0) - qp.RX(param, wires=w1) - qp.RX(param, wires=w2) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit_3(_: float): - # CHECK: [[QREG:%.+]] = qref.alloc - # CHECK-NEXT: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK-NEXT: qref.custom "Hadamard"() [[QUBIT]] : !qref.bit - # CHECK-NEXT: qref.compbasis(qreg [[QREG]] : !qref.reg<1>) : !quantum.obs - identity() - all_wires_rx(float, int, int, int) - qp.Hadamard(0) - return qp.probs() - - # CHECK-LABEL: @identity - # CHECK-LABEL: @all_wires_rx - - print(circuit_3.mlir) - - -test_multiple_decomposition_rules() - - -def test_decomposition_rule_shaped_wires(): - """Test decomposition rule with passing a shaped array of wires""" - - @decomposition_rule(is_qreg=True) - def shaped_wires_rule(param: TensorLike, wires: WiresLike): - qp.RX(param, wires=wires[0]) - qp.RX(param, wires=wires[1]) - qp.RX(param, wires=wires[2]) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit_4(_: float): - # CHECK: module @circuit_4 - shaped_wires_rule(float, jax.core.ShapedArray((3,), int)) - qp.Hadamard(0) - return qp.probs() - - # CHECK: @shaped_wires_rule([[QREG:%.+]]: !qref.reg<1>, [[PARAM_TENSOR:%.+]]: tensor, [[QUBITS:%.+]]: tensor<3xi64>) - # CHECK-NEXT: [[IDX_0:%.+]] = stablehlo.slice [[QUBITS]] [0:1] : (tensor<3xi64>) -> tensor<1xi64> - # CHECK-NEXT: [[RIDX_0:%.+]] = stablehlo.reshape [[IDX_0]] : (tensor<1xi64>) -> tensor - # CHECK-NEXT: [[EXTRACTED:%.+]] = tensor.extract [[RIDX_0]][] : tensor - # CHECK-NEXT: [[QUBIT:%.+]] = qref.get [[QREG]][[[EXTRACTED]]] : !qref.reg<1>, i64 -> !qref.bit - # CHECK-NEXT: [[EXTRACTED_0:%.+]] = tensor.extract [[PARAM_TENSOR]][] : tensor - # CHECK-NEXT: qref.custom "RX"([[EXTRACTED_0]]) [[QUBIT]] : !qref.bit - - print(circuit_4.mlir) - - -test_decomposition_rule_shaped_wires() - - -def test_decomposition_rule_expanded_wires(): - """Test decomposition rule with passing expanding wires as a Python list""" - - def shaped_wires_rule(param: TensorLike, wires: WiresLike): - qp.RX(param, wires=wires[0]) - qp.RX(param, wires=wires[1]) - qp.RX(param, wires=wires[2]) - - @decomposition_rule(is_qreg=False, num_params=1) - def expanded_wires_rule(param: TensorLike, w1, w2, w3): - shaped_wires_rule(param, [w1, w2, w3]) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit_5(_: float): - # CHECK: module @circuit_5 - expanded_wires_rule(float, int, int, int) - qp.Hadamard(0) - return qp.probs() - - # CHECK-LABEL: @expanded_wires_rule(%arg0: tensor, %arg1: !qref.bit, %arg2: !qref.bit, %arg3: !qref.bit) - - print(circuit_5.mlir) - - -test_decomposition_rule_expanded_wires() - - -def test_decomposition_rule_with_cond(): - """Test decomposition rule with a conditional path""" - - @decomposition_rule(is_qreg=True) - def cond_RX(param: TensorLike, w0: WiresLike): - - def true_path(): - qp.RX(param, wires=w0) - - def false_path(): ... - - qp.cond(param != 0.0, true_path, false_path)() - - @qp.qjit(autograph=False, capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit_6(): - # CHECK: module @circuit_6 - cond_RX(float, jax.core.ShapedArray((1,), int)) - return qp.probs() - - # CHECK: @cond_RX([[QREG:%.+]]: !qref.reg<1>, [[PARAM_TENSOR:%.+]]: tensor, [[QUBITS:%.+]]: tensor<1xi64>) - # CHECK-NEXT: [[ZERO:%.+]] = stablehlo.constant dense<0.000000e+00> : tensor - # CHECK-NEXT: [[COND_TENSOR:%.+]] = stablehlo.compare NE, [[PARAM_TENSOR]], [[ZERO]], FLOAT : (tensor, tensor) -> tensor - # CHECK-NEXT: [[COND:%.+]] = tensor.extract [[COND_TENSOR]][] : tensor - # CHECK-NEXT: scf.if [[COND]] - # CHECK-DAG: [[QUBIT:%.+]] = qref.get [[QREG]][%extracted_0] : !qref.reg<1>, i64 -> !qref.bit - # CHECK-DAG: [[PARAM:%.+]] = tensor.extract [[PARAM_TENSOR]][] : tensor - # CHECK: qref.custom "RX"([[PARAM]]) [[QUBIT]] : !qref.bit - # CHECK: return - - print(circuit_6.mlir) - - -test_decomposition_rule_with_cond() - - -def test_decomposition_rule_caller(): - """Test decomposition rules with a caller""" - - @decomposition_rule(is_qreg=True) - def rule_op1_decomp(_: TensorLike, wires: WiresLike): - qp.Hadamard(wires=wires[0]) - qp.Hadamard(wires=[1]) - - @decomposition_rule(is_qreg=True) - def rule_op2_decomp(param: TensorLike, wires: WiresLike): - qp.RX(param, wires=wires[0]) - - def decomps_caller(param: TensorLike, wires: WiresLike): - rule_op1_decomp(param, wires) - rule_op2_decomp(param, wires) - - @qp.qjit(autograph=False, capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module @circuit_7 - def circuit_7(): - # CHECK: [[QREG:%.+]] = qref.alloc - # CHECK: qref.compbasis(qreg [[QREG]] : !qref.reg<1>) : !quantum.obs - decomps_caller(float, jax.core.ShapedArray((2,), int)) - return qp.probs() - - # CHECK-LABEL: @rule_op1_decomp(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<2xi64>) - # CHECK-LABEL: @rule_op2_decomp(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<2xi64>) - print(circuit_7.mlir) - - -test_decomposition_rule_caller() - - -def test_decompose_gateset_without_graph(): - """Test the decompose transform to a target gate set without the graph decomposition.""" - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RX", "RZ"}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @circuit_8() -> tensor attributes {diff_method = "adjoint", llvm.linkage = #llvm.linkage, quantum.node} - def circuit_8(): - return qp.expval(qp.Z(0)) - - print(circuit_8.mlir) - - -test_decompose_gateset_without_graph() - - -def test_decompose_gateset_with_graph(): - """Test the decompose transform to a target gate set with the graph decomposition.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RX"}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @simple_circuit_9() -> tensor attributes {decompose_gatesets - def simple_circuit_9(): - return qp.expval(qp.Z(0)) - - print(simple_circuit_9.mlir) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RX", "RZ"}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_9() -> tensor attributes {decompose_gatesets - def circuit_9(): - return qp.expval(qp.Z(0)) - - print(circuit_9.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_gateset_with_graph() - - -def test_decompose_gateset_operator_with_graph(): - """Test the decompose transform to a target gate set with the graph decomposition.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.RX}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @simple_circuit_10() -> tensor attributes {decompose_gatesets - def simple_circuit_10(): - return qp.expval(qp.Z(0)) - - print(simple_circuit_10.mlir) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.RX, qp.RZ, "PauliZ", qp.PauliX, qp.Hadamard}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @circuit_10() -> tensor attributes {decompose_gatesets - def circuit_10(): - return qp.expval(qp.Z(0)) - - print(circuit_10.mlir) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.RX, qp.RZ, qp.PauliZ, qp.PauliX, qp.Hadamard}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_11() -> tensor attributes {decompose_gatesets - def circuit_11(): - return qp.expval(qp.Z(0)) - - print(circuit_11.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_gateset_operator_with_graph() - - -def test_decompose_gateset_with_rotxzx(): - """Test the decompose transform with a custom operator with the graph decomposition.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RotXZX"}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-LABEL: @simple_circuit_12() -> tensor attributes {decompose_gatesets - def simple_circuit_12(): - return qp.expval(qp.Z(0)) - - print(simple_circuit_12.mlir) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.ftqc.RotXZX}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_12() -> tensor attributes {decompose_gatesets - def circuit_12(): - return qp.expval(qp.Z(0)) - - print(circuit_12.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_gateset_with_rotxzx() - - -def test_decomposition_rule_name(): - """Test the name of the decomposition rule is not updated with circuit instantiation.""" - - qp.decomposition.enable_graph() - - @decomposition_rule - def _ry_to_rz_rx(phi, wires: WiresLike, **__): - """Decomposition of RY gate using RZ and RX gates.""" - qp.RZ(-np.pi / 2, wires=wires) - qp.RX(phi, wires=wires) - qp.RZ(np.pi / 2, wires=wires) - - @decomposition_rule - def _rot_to_rz_ry_rz(phi, theta, omega, wires: WiresLike, **__): - """Decomposition of Rot gate using RZ and RY gates.""" - qp.RZ(phi, wires=wires) - qp.RY(theta, wires=wires) - qp.RZ(omega, wires=wires) - - @decomposition_rule - def _u2_phaseshift_rot_decomposition(phi, delta, wires, **__): - """Decomposition of U2 gate using Rot and PhaseShift gates.""" - pi_half = qp.math.ones_like(delta) * (np.pi / 2) - qp.Rot(delta, pi_half, -delta, wires=wires) - qp.PhaseShift(delta, wires=wires) - qp.PhaseShift(phi, wires=wires) - - @decomposition_rule - def _xzx_decompose(phi, theta, omega, wires, **__): - """Decomposition of Rot gate using RX and RZ gates in XZX format.""" - qp.RX(phi, wires=wires) - qp.RZ(theta, wires=wires) - qp.RX(omega, wires=wires) - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={"RX", "RZ", "PhaseShift"}) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_13() -> tensor attributes {decompose_gatesets - def circuit_13(): - _ry_to_rz_rx(float, int) - _rot_to_rz_ry_rz(float, float, float, int) - _u2_phaseshift_rot_decomposition(float, float, int) - _xzx_decompose(float, float, float, int) - return qp.expval(qp.Z(0)) - - # CHECK-LABEL: @_ry_to_rz_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor) - # CHECK-LABEL: @_rot_to_rz_ry_rz(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor) - # CHECK-LABEL: @_u2_phaseshift_rot_decomposition(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor) - # CHECK-LABEL: @_xzx_decompose(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor) - print(circuit_13.mlir) - - qp.decomposition.disable_graph() - - -test_decomposition_rule_name() - - -def test_decomposition_rule_name_update(): - """Test the name of the decomposition rule is updated in the MLIR output.""" - - qp.decomposition.enable_graph() - - @qp.register_resources({qp.RZ: 2, qp.RX: 1}) - def rz_rx(phi, wires: WiresLike, **__): - """Decomposition of RY gate using RZ and RX gates.""" - qp.RZ(-np.pi / 2, wires=wires) - qp.RX(phi, wires=wires) - qp.RZ(np.pi / 2, wires=wires) - - @qp.register_resources({qp.RZ: 2, qp.RY: 1}) - def rz_ry_rz(phi, theta, omega, wires: WiresLike, **__): - """Decomposition of Rot gate using RZ and RY gates.""" - qp.RZ(phi, wires=wires) - qp.RY(theta, wires=wires) - qp.RZ(omega, wires=wires) - - @qp.register_resources({qp.RY: 1, qp.GlobalPhase: 1}) - def ry_gp(wires: WiresLike, **__): - """Decomposition of PauliY gate using RY and GlobalPhase gates.""" - qp.RY(np.pi, wires=wires) - qp.GlobalPhase(-np.pi / 2, wires=wires) - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RX", "RZ", "GlobalPhase"}, - fixed_decomps={ - qp.RY: rz_rx, - qp.Rot: rz_ry_rz, - qp.PauliY: ry_gp, - }, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_14() -> tensor attributes {decompose_gatesets - def circuit_14(): - qp.RY(0.5, wires=0) - qp.Rot(0.1, 0.2, 0.3, wires=1) - qp.PauliY(wires=2) - return qp.expval(qp.Z(0)) - - # CHECK-DAG: @rz_ry_rz(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) - # CHECK-DAG: @rz_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) - # CHECK-DAG: @ry_gp(%arg0: !qref.reg<3>, %arg1: tensor<1xi64>) - print(circuit_14.mlir) - - qp.decomposition.disable_graph() - - -test_decomposition_rule_name_update() - - -def test_decomposition_inside_subroutine(): - """Test that operators inside subroutines can be decomposed.""" - - qp.decomposition.enable_graph() - - @qp.templates.Subroutine - def f(x, wires): - qp.IsingXX(x, wires) - - @qp.qjit(capture=True, target="mlir") - @qp.decompose(gate_set=qp.gate_sets.ROTATIONS_PLUS_CNOT) - @qp.qnode(qp.device("lightning.qubit", wires=5)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - def subroutine_circuit(): - # CHECK-DAG: [[FIRST_CONST:%.+]] = stablehlo.constant dense<5.000000e-01> : tensor - # CHECK-DAG: [[SECOND_CONST:%.+]] = stablehlo.constant dense<1.200000e+00> : tensor - - # CHECK: [[QREG:%.+]] = qref.alloc - # CHECK: call @f([[QREG]], [[FIRST_CONST]], {{%.+}}) : (!qref.reg<5>, tensor, tensor<2xi64>) - # CHECK: call @f([[QREG]], [[SECOND_CONST]], {{%.+}}) : (!qref.reg<5>, tensor, tensor<2xi64>) - - f(0.5, (0, 1)) - f(1.2, (2, 3)) - return qp.probs(wires=0) - - # CHECK-DAG: @_isingxx_to_cnot_rx_cnot(%arg0: !qref.reg<5>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) - print(subroutine_circuit.mlir) - qp.decomposition.disable_graph() - - -test_decomposition_inside_subroutine() - - -def test_decomposition_rule_name_update_multi_qubits(): - """Test the name of the decomposition rule with multi-qubit gates.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RY", "RX", "CNOT", "Hadamard", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # CHECK-LABEL: @circuit_15() -> tensor attributes {decompose_gatesets - def circuit_15(): - qp.SingleExcitation(0.5, wires=[0, 1]) - qp.SingleExcitationPlus(0.5, wires=[0, 1]) - qp.SingleExcitationMinus(0.5, wires=[0, 1]) - qp.DoubleExcitation(0.5, wires=[0, 1, 2, 3]) - return qp.expval(qp.Z(0)) - - # CHECK-DAG: @_cry(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CRY"} - # CHECK-DAG: @_s_phaseshift(%arg0: !qref.reg<4>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "S"} - # CHECK-DAG: @_phaseshift_to_rz_gp(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PhaseShift"} - # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - # CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # CHECK-DAG: @_doublexcit(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<4xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 4 : i64, target_gate = "DoubleExcitation"} - # CHECK-DAG: @_single_excitation_decomp(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "SingleExcitation"} - print(circuit_15.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decomposition_rule_name_update_multi_qubits)() - - -def test_decomposition_rule_name_adjoint(): - """Test decomposition rule with qp.adjoint.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RY", "RX", "CZ", "GlobalPhase", "Adjoint(SingleExcitation)"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - # CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - def circuit_16(x: float): - # CHECK: qref.adjoint { - # CHECK: qref.adjoint { - # CHECK: qref.adjoint { - # CHECK: qref.adjoint { - qp.adjoint(qp.CNOT)(wires=[0, 1]) - qp.adjoint(qp.Hadamard)(wires=2) - qp.adjoint(qp.RZ)(0.5, wires=3) - qp.adjoint(qp.SingleExcitation)(0.1, wires=[0, 1]) - qp.adjoint(qp.SingleExcitation(x, wires=[0, 1])) - return qp.expval(qp.Z(0)) - - # CHECK-DAG: @_single_excitation_decomp(%arg0: !qref.reg<4>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "SingleExcitation"} - # CHECK-DAG: @_hadamard_to_rz_ry(%arg0: !qref.reg<4>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Hadamard"} - # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - # CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !qref.reg<4>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # CHECK-DAG: @_cnot_to_cz_h(%arg0: !qref.reg<4>, %arg1: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CNOT"} - print(circuit_16.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decomposition_rule_name_adjoint)() - - -# TODO: Reenable this once the underlying non-determinism issue is resolved -def test_decomposition_rule_name_ctrl(): - """Test decomposition rule with qp.ctrl.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RX", "RZ", "H", "CZ"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - # SKIP-CHECK-DAG: %0 = transform.apply_registered_pass "decompose-lowering" - # SKIP-CHECK{LITERAL}: @circuit_17() -> tensor attributes {decompose_gatesets - def circuit_17(): - # SKIP-CHECK: %out_qubits:2 = quantum.custom "CRY"(%cst) %1, %2 : !quantum.bit, !quantum.bit - # SKIP-CHECK-NEXT: %out_qubits_0:2 = quantum.custom "CNOT"() %out_qubits#0, %out_qubits#1 : !quantum.bit, !quantum.bit - qp.ctrl(qp.RY, control=0)(0.5, 1) - qp.ctrl(qp.PauliX, control=0)(1) - return qp.expval(qp.Z(0)) - - # SKIP-CHECK-DAG: @_cnot_to_cz_h(%arg0: !quantum.reg, %arg1: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CNOT"} - # SKIP-CHECK-DAG: @_cry(%arg0: !quantum.reg, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "CRY"} - # SKIP-CHECK-DAG: @_ry_to_rz_rx(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RY"} - # SKIP-CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # print(circuit_17.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decomposition_rule_name_ctrl)() - - -# TODO: Reenable this once the underlying non-determinism issue is resolved -def test_qft_decomposition(): - """Test the decomposition of the QFT""" - - qp.decomposition.enable_graph() - - @qp.qjit(autograph=True, target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RX", "RY", "CNOT", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - # SKIP-CHECK: %0 = transform.apply_registered_pass "decompose-lowering" - # SKIP-CHECK: @circuit_18(%arg0: tensor<3xf64>) -> tensor attributes {decompose_gatesets - def circuit_18(): - # %6 = scf.for %arg1 = %c0 to %c4 step %c1 iter_args(%arg2 = %0) -> (!quantum.reg) { - # %23 = scf.for %arg3 = %c0 to %22 step %c1 iter_args(%arg4 = %21) -> (!quantum.reg) { - # %7 = scf.for %arg1 = %c0 to %c2 step %c1 iter_args(%arg2 = %6) -> (!quantum.reg) { - qp.QFT(wires=[0, 1, 2, 3]) - return qp.expval(qp.Z(0)) - - # SKIP-CHECK-DAG: @ag___cphase_to_rz_cnot(%arg0: !quantum.reg, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "ControlledPhaseShift"} - # SKIP-CHECK-DAG: @ag___rz_to_ry_rx(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - # SKIP-CHECK-DAG: @ag___rot_to_rz_ry_rz(%arg0: !quantum.reg, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # SKIP-CHECK-DAG: @ag___swap_to_cnot(%arg0: !quantum.reg, %arg1: tensor<2xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "SWAP"} - # SKIP-CHECK-DAG: @ag___hadamard_to_rz_ry(%arg0: !quantum.reg, %arg1: tensor<1xi64>) -> !quantum.reg attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Hadamard"} - # print(circuit_18.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_qft_decomposition)() - - -def test_decompose_lowering_with_other_passes(): - """Test the decompose lowering pass with other passes in a pass pipeline.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @qp.transforms.merge_rotations - @qp.transforms.cancel_inverses - @partial( - qp.transforms.decompose, - gate_set={"RZ", "RY", "CNOT", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module attributes {transform.with_named_sequence} { - # CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.op<"builtin.module">) { - # CHECK-NEXT: [[ONE:%.+]] = transform.apply_registered_pass "decompose-lowering" to %arg0 : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: [[TWO:%.+]] = transform.apply_registered_pass "cancel-inverses" to [[ONE]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to [[TWO]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: transform.yield - # CHECK-NEXT: } - def circuit_19(): - - # CHECK: [[QREG:%.+]] = qref.alloc( 1) : !qref.reg<1> - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "PauliX"() [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "PauliX"() [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "RX"({{%.+}}) [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "RX"({{%.+}}) [[QUBIT]] : !qref.bit - qp.PauliX(0) - qp.PauliX(0) - qp.RX(0.1, wires=0) - qp.RX(-0.1, wires=0) - return qp.expval(qp.PauliX(0)) - - # CHECK-DAG: @_paulix_to_rx(%arg0: !qref.reg<1>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PauliX"} - # CHECK-DAG: @_rx_to_rz_ry(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RX"} - print(circuit_19.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decompose_lowering_with_other_passes)() - - -def test_decompose_lowering_multirz(): - """Test the decompose lowering pass with MultiRZ in the gate set.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"CNOT", "RZ"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - # CHECK: %0 = transform.apply_registered_pass "decompose-lowering" - def circuit_20(x: float): - # CHECK: [[QREG:%.+]] = qref.alloc( 3) : !qref.reg<3> - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor - # CHECK: qref.multirz([[angle]]) [[q0]] : !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: [[q1:%.+]] = qref.get [[QREG]][ 1] : !qref.reg<3> -> !qref.bit - # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor - # CHECK: qref.multirz([[angle]]) [[q0]], [[q1]] : !qref.bit, !qref.bit - # CHECK: [[q1:%.+]] = qref.get [[QREG]][ 1] : !qref.reg<3> -> !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: [[q2:%.+]] = qref.get [[QREG]][ 2] : !qref.reg<3> -> !qref.bit - # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor - # CHECK: qref.multirz([[angle]]) [[q1]], [[q0]], [[q2]] : !qref.bit, !qref.bit, !qref.bit - qp.MultiRZ(x, wires=[0]) - qp.MultiRZ(x, wires=[0, 1]) - qp.MultiRZ(x, wires=[1, 0, 2]) - return qp.expval(qp.PauliX(0)) - - # CHECK-DAG: @_multi_rz_decomposition_wires_1(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "MultiRZ"} - # CHECK-DAG: @_multi_rz_decomposition_wires_2(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<2xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 2 : i64, target_gate = "MultiRZ"} - # CHECK-DAG: @_multi_rz_decomposition_wires_3(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<3xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 3 : i64, target_gate = "MultiRZ"} - # CHECK-DAG: scf.for %arg3 = %c0 to %c2 step %c1 - # CHECK-DAG: scf.for %arg3 = %c1 to %c3 step %c1 - print(circuit_20.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_lowering_multirz() - - -def test_decompose_lowering_with_ordered_passes(): - """Test the decompose lowering pass with other passes in a specific order in a pass pipeline.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RZ", "RY", "CNOT", "GlobalPhase"}, - ) - @qp.transforms.merge_rotations - @qp.transforms.cancel_inverses - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module attributes {transform.with_named_sequence} { - # CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.op<"builtin.module">) { - # CHECK-NEXT: [[FIRST:%.+]] = transform.apply_registered_pass "cancel-inverses" to %arg0 : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: [[SECOND:%.+]] = transform.apply_registered_pass "merge-rotations" to [[FIRST]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "decompose-lowering" to [[SECOND]] : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> - # CHECK-NEXT: transform.yield - # CHECK-NEXT: } - def circuit_21(x: float): - # CHECK: [[QREG:%.+]] = qref.alloc( 1) : !qref.reg<1> - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "PauliX"() [[q0]] : !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: qref.custom "PauliX"() [[q0]] : !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: [[angle:%.+]] = tensor.extract %arg0[] : tensor - # CHECK: qref.custom "RX"([[angle]]) [[q0]] : !qref.bit - # CHECK: [[negated:%.+]] = stablehlo.negate %arg0 : tensor - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<1> -> !qref.bit - # CHECK: [[neg_angle:%.+]] = tensor.extract [[negated]][] : tensor - # CHECK: qref.custom "RX"([[neg_angle]]) [[q0]] : !qref.bit - qp.PauliX(0) - qp.PauliX(0) - qp.RX(x, wires=0) - qp.RX(-x, wires=0) - return qp.expval(qp.PauliX(0)) - - # CHECK-DAG: @_paulix_to_rx(%arg0: !qref.reg<1>, %arg1: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PauliX"} - # CHECK-DAG: @_rx_to_rz_ry(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RX"} - # CHECK-DAG: @_rot_to_rz_ry_rz(%arg0: !qref.reg<1>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - print(circuit_21.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decompose_lowering_with_ordered_passes)() - - -def test_decompose_lowering_with_gphase(): - """Test the decompose lowering pass with GlobalPhase.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RX", "RY", "GlobalPhase"}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3)) - # CHECK: %0 = transform.apply_registered_pass "decompose-lowering" - def circuit_22(): - # CHECK: [[QREG:%.+]] = qref.alloc( 3) : !qref.reg<3> - # CHECK: qref.gphase({{%.+}}) - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: qref.custom "PhaseShift"({{%.+}}) [[q0]] : !qref.bit - # CHECK: [[q0:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<3> -> !qref.bit - # CHECK: qref.custom "PhaseShift"({{%.+}}) [[q0]] : !qref.bit - - qp.GlobalPhase(0.5) - qp.ctrl(qp.GlobalPhase, control=0)(0.3) - qp.ctrl(qp.GlobalPhase, control=0)(phi=0.3, wires=[1, 2]) - return qp.expval(qp.PauliX(0)) - - # CHECK-DAG: @_phaseshift_to_rz_gp(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "PhaseShift"} - # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - print(circuit_22.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decompose_lowering_with_gphase)() - - -def test_decompose_lowering_alt_decomps(): - """Test the decompose lowering pass with alternative decompositions.""" - - qp.decomposition.enable_graph() - - @qp.register_resources({qp.RY: 1}) - def custom_rot_cheap(params, wires: WiresLike): - qp.RY(params[1], wires=wires) - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RY", "RZ"}, - alt_decomps={qp.Rot: [custom_rot_cheap]}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3), shots=1000) - def circuit_23(x: float, y: float): - qp.Rot(x, y, x + y, wires=1) - return qp.expval(qp.PauliZ(0)) - - # CHECK-DAG: @custom_rot_cheap(%arg0: !qref.reg<3>, %arg1: tensor<3xf64>, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - print(circuit_23.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_lowering_alt_decomps() - - -def test_decompose_lowering_with_tensorlike(): - """Test the decompose lowering pass with fixed decompositions - using TensorLike parameters.""" - - qp.decomposition.enable_graph() - - @qp.register_resources({qp.RZ: 2, qp.RY: 1}) - def custom_rot(params: TensorLike, wires: WiresLike): - qp.RZ(params[0], wires=wires) - qp.RY(params[1], wires=wires) - qp.RZ(params[2], wires=wires) - - @qp.register_resources({qp.RZ: 1, qp.CNOT: 4}) - def custom_multirz(params: TensorLike, wires: WiresLike): - qp.CNOT(wires=(wires[2], wires[1])) - qp.CNOT(wires=(wires[1], wires[0])) - qp.RZ(params[0], wires=wires[0]) - qp.CNOT(wires=(wires[1], wires[0])) - qp.CNOT(wires=(wires[2], wires[1])) - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RY", "RX", qp.CNOT}, - fixed_decomps={qp.Rot: custom_rot, qp.MultiRZ: custom_multirz}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=3), shots=1000) - def circuit_24(x: float, y: float): - qp.Rot(x, y, x + y, wires=1) - qp.MultiRZ(x + y, wires=[0, 1, 2]) - return qp.expval(qp.PauliZ(0)) - - # CHECK-DAG: @custom_multirz_wires_3(%arg0: !qref.reg<3>, %arg1: tensor<1xf64>, %arg2: tensor<3xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 3 : i64, target_gate = "MultiRZ"} - # CHECK-DAG: @_rz_to_ry_rx(%arg0: !qref.reg<3>, %arg1: tensor, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "RZ"} - # CHECK-DAG: @custom_rot(%arg0: !qref.reg<3>, %arg1: tensor<3xf64>, %arg2: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - print(circuit_24.mlir) - - qp.decomposition.disable_graph() - - -skip_if_pauli_rot_issue(test_decompose_lowering_with_tensorlike)() - - -def test_decompose_lowering_fallback(): - """Test the decompose lowering pass when the graph is failed.""" - - qp.decomposition.enable_graph() - - @qp.qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set={qp.RX, qp.RZ}) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - # CHECK-LABEL: @circuit_25() - def circuit_25(): - # CHECK: [[pi_over_2:%.+]] = arith.constant 1.5707963267948966 : f64 - # CHECK: [[QREG:%.+]] = qref.alloc( 2) : !qref.reg<2> - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<2> -> !qref.bit - # CHECK: qref.custom "RZ"([[pi_over_2]]) [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<2> -> !qref.bit - # CHECK: qref.custom "RX"([[pi_over_2]]) [[QUBIT]] : !qref.bit - # CHECK: [[QUBIT:%.+]] = qref.get [[QREG]][ 0] : !qref.reg<2> -> !qref.bit - # CHECK: qref.custom "RZ"([[pi_over_2]]) [[QUBIT]] : !qref.bit - qp.Hadamard(0) - return qp.state() - - print(circuit_25.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_lowering_fallback() - - -def test_decompose_lowering_params_ordering(): - """Test the order of params and wires in the captured decomposition rule.""" - - qp.decomposition.enable_graph() - - @qjit(target="mlir", capture=True) - @partial(qp.transforms.decompose, gate_set=[qp.RX, qp.RY, qp.RZ]) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - # CHECK-LABEL: @circuit_26(%arg0: tensor, %arg1: tensor, %arg2: tensor) - def circuit_26(x: float, y: float, z: float): - qp.Rot(x, y, z, wires=0) - return qp.expval(qp.PauliZ(0)) - - # CHECK-LABEL: @_rot_to_rz_ry_rz(%arg0: !qref.reg<2>, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor<1xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 1 : i64, target_gate = "Rot"} - # CHECK: [[EXTRACTED_1:%.+]] = tensor.extract %arg1[] : tensor - # CHECK-NEXT: qref.custom "RZ"([[EXTRACTED_1]]) {{%.+}} : !qref.bit - # CHECK: [[EXTRACTED_2:%.+]] = tensor.extract %arg2[] : tensor - # CHECK-NEXT: qref.custom "RY"([[EXTRACTED_2]]) {{%.+}} : !qref.bit - # CHECK: [[EXTRACTED_3:%.+]] = tensor.extract %arg3[] : tensor - # CHECK-NEXT: qref.custom "RZ"([[EXTRACTED_3]]) {{%.+}} : !qref.bit - # CHECK: return - print(circuit_26.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_lowering_params_ordering() - - -def test_decomposition_rule_with_allocation(): - """Test decomposition rule with dynamic qubit allocation""" - - @decomposition_rule(is_qreg=True) - def Hadamard0_with_alloc(wire: WiresLike): - with qp.allocate(1) as q: - qp.X(q[0]) - qp.CNOT(wires=[q[0], wire]) - - @qp.qjit(capture=True) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - # CHECK: module @circuit_27 - def circuit_27(): - Hadamard0_with_alloc(int) - return qp.probs() - - # CHECK-LABEL: @Hadamard0_with_alloc(%arg0: !qref.reg<1>, %arg1: tensor) - # CHECK: [[dynalloc_qreg:%.+]] = qref.alloc( 1) - # CHECK: [[dynalloc_bit0:%.+]] = qref.get [[dynalloc_qreg]][ 0] - # CHECK: qref.custom "PauliX"() [[dynalloc_bit0]] - # CHECK: [[detensor:%.+]] = tensor.extract %arg1[] - # CHECK: [[glob_bit:%.+]] = qref.get %arg0[[[detensor]]] - # CHECK: qref.custom "CNOT"() [[dynalloc_bit0]], [[glob_bit]] - # CHECK: qref.dealloc [[dynalloc_qreg]] - # CHECK: return - - print(circuit_27.mlir) - - -test_decomposition_rule_with_allocation() - - -def test_decompose_autograph_multi_blocks(): - """Test the decompose lowering pass with autograph in the program and rule.""" - - qp.decomposition.enable_graph() - - def _multi_rz_decomposition_resources(num_wires): - """Resources required for MultiRZ decomposition.""" - return {qp.RZ: 1, qp.CNOT: 2 * (num_wires - 1)} - - @qp.register_resources(_multi_rz_decomposition_resources) - @qp.capture.run_autograph - def _multi_rz_decomposition(theta: TensorLike, wires: WiresLike, **__): - """Decomposition of MultiRZ using CNOTs and RZs.""" - for i in range(len(wires) - 1): - qp.CNOT(wires=(wires[i], wires[i + 1])) - qp.RZ(theta, wires=wires[0]) - for i in range(len(wires) - 1, 0, -1): - qp.CNOT(wires=(wires[i], wires[i - 1])) - - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={"RZ", "CNOT"}, - fixed_decomps={qp.MultiRZ: _multi_rz_decomposition}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=5)) - def circuit_29(n: int): - - # CHECK: scf.for %arg1 = {{%.+}} to {{%.+}} step {{%.+}} { - @qp.for_loop(n) - def f(i): # pylint: disable=unused-argument - qp.MultiRZ(0.5, wires=[0, 1, 2, 3, 4]) - - f() # pylint: disable=no-value-for-parameter - - return qp.expval(qp.Z(0)) - - # CHECK-LABEL: @ag___multi_rz_decomposition_wires_5(%arg0: !qref.reg<5>, %arg1: tensor<1xf64>, %arg2: tensor<5xi64>) attributes {llvm.linkage = #llvm.linkage, num_wires = 5 : i64, target_gate = "MultiRZ"} - # CHECK: scf.for %arg3 = {{%.+}} to {{%.+}} step {{%.+}} { - # CHECK: scf.for %arg3 = {{%.+}} to {{%.+}} step {{%.+}} { - print(circuit_29.mlir) - - qp.decomposition.disable_graph() - - -test_decompose_autograph_multi_blocks() - - -def test_decompose_work_wires_context_manager(): - """ - Test that decomposition with work wires is correctly applied when allocating with the context - manager. - """ - - @decomposition_rule(is_qreg=True, op_type="PauliZ") - def my_decomp(wires): - with qp.allocate(2, restored=False) as work_wires: - qp.X(wires[0]) - qp.X(wires[1]) - qp.H(work_wires[0]) - qp.H(work_wires[1]) - - @qp.qjit(capture=True) - @qp.transform(pass_name="decompose-lowering") - @qp.qnode(qp.device("lightning.qubit", wires=3)) - def my_circuit(): - my_decomp(jax.core.ShapedArray((2,), int)) - qp.Z(0) - return qp.probs() - - # check that decomp arrives properly - # CHECK-LABEL: @my_decomp({{.*}}) attributes {{{.*}} target_gate = "PauliZ"} - print(my_circuit.mlir) - - # check that decomp is applied properly - # CHECK-NOT: PauliZ - # CHECK-NOT: my_decomp - - # two allocates, one for main register and one for decomp register - # CHECK: allocate - # CHECK: allocate - # CHECK: PauliX - # CHECK: PauliX - # CHECK: Hadamard - # CHECK: Hadamard - # CHECK: release - # CHECK: release - print(my_circuit.mlir_opt) - - -test_decompose_work_wires_context_manager() - - -def test_decompose_work_wires_alloc_dealloc(): - """ - Test that decomposition with work wires is correctly applied when allocating/deallocating - explicitly. - """ - - @decomposition_rule(is_qreg=True, op_type="RY") - def my_decomp(angle, wires): - work_wires = qp.allocate(2) - qp.CNOT((work_wires[0], wires[0])) - qp.RX(-np.pi / 2, wires[0]) - qp.RZ(angle, wires[0]) - qp.RX(np.pi / 2, wires[0]) - qp.CNOT((work_wires[1], wires[1])) - qp.deallocate(work_wires) - - @qp.qjit(capture=True) - @qp.transform(pass_name="decompose-lowering") - @qp.qnode(qp.device("lightning.qubit", wires=3)) - def my_circuit(angle: float): - my_decomp(float, jax.core.ShapedArray((2,), int)) - qp.RY(angle, 0) - return qp.probs() - - # check that decomp arrives properly - # CHECK-LABEL: @my_decomp({{.*}}) attributes {{{.*}} target_gate = "RY"} - print(my_circuit.mlir) - - # check that the decomposition applies properly - # CHECK-NOT: my_decomp - # CHECK-NOT: RY - - # two allocates, one for main register and one for decomp register - # CHECK: allocate - # CHECK: allocate - # CHECK: CNOT - # CHECK: RX - # CHECK: RZ - # CHECK: RX - # CHECK: CNOT - # CHECK: release - # CHECK: release - print(my_circuit.mlir_opt) - - -test_decompose_work_wires_alloc_dealloc() - - -def test_decompose_work_wires_control_flow(): - """Test that decomposition with work wires + control flow is correctly applied.""" - - @decomposition_rule(is_qreg=True, op_type="CRX") - def my_decomp(angle, wires, **_): - def true_func(): - qp.CNOT(wires) - - with qp.allocate(2, state="any", restored=True) as w: - for _ in range(2): - qp.H(w[0]) - qp.X(w[1]) - - def false_func(): - with qp.allocate(1, state="any", restored=False) as w: - qp.H(w) - - m = qp.measure(wires[0]) - - qp.cond(m, qp.CNOT)(wires) - - qp.cond(angle > 1.2, true_func, false_func)() - - @qp.qjit(capture=True) - @qp.transform(pass_name="decompose-lowering") - @qp.qnode(qp.device("lightning.qubit", wires=4)) - def circuit(): - my_decomp(float, jax.core.ShapedArray((2,), int)) - qp.CRX(1.7, wires=[0, 1]) - qp.CRX(-7.2, wires=[0, 1]) - return qp.state() - - # target_gate attribute is correctly applied - # CHECK: my_decomp([[args:.*]]) attributes {[[other_attributes:.*]] target_gate = "CRX"} - print(circuit.mlir) - - # test that the decomposition is applied correctly - # CHECK-NOT: CRX - # CHECK-NOT: my_decomp - - # allocate for main register, subsequent allocates+releases for decomp registers - # CHECK: allocate - - # first CRX: true branch - # CHECK: CNOT - # CHECK: allocate - # CHECK: Hadamard - # CHECK: PauliX - # CHECK: Hadamard - # CHECK: PauliX - # CHECK: release - - # second CRX: false branch - # CHECK: allocate - # CHECK: Hadamard - # CHECK: Measure - # CHECK: cond - # CHECK: CNOT - # CHECK: release - - # release main register - # CHECK: release - - print(circuit.mlir_opt) - - -test_decompose_work_wires_control_flow() - - -def test_decompose_work_wires_with_decompose_transform(): - """Test that work wires are correctly lowered and decomposed by the decompose transform.""" - - qp.decomposition.enable_graph() - - @qp.register_resources({qp.X: 1, qp.Z: 1}) - def my_decomp(wire): - with qp.allocate(1) as work_wire: - qp.X(work_wire) - qp.Z(wire) - qp.X(work_wire) - - @qjit(capture=True) - @partial( - qp.transforms.decompose, - gate_set={ - "X", - "Z", - }, - fixed_decomps={ - qp.Y: my_decomp, - }, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def my_circuit(): - qp.Y(0) - return qp.probs() - - # CHECK-NOT: Y - # CHECK-NOT: my_decomp - - # two allocates, one for main register and one for decomp register - # CHECK: allocate - # CHECK: allocate - # CHECK: X - # CHECK: Z - # CHECK: X - # CHECK: release - # CHECK: release - print(my_circuit.mlir_opt) - - qp.decomposition.disable_graph() - - -test_decompose_work_wires_with_decompose_transform() - - -def test_num_work_wires(): - """Test that num_work_wires can be passed and is correctly used in solving the graph.""" - - qp.decomposition.enable_graph() - - @qp.register_resources( - {qp.CNOT: 3, qp.H: 1, qp.X: 1, qp.ops.op_math.Conditional: 2}, - work_wires={"borrowed": 2, "garbage": 1}, - ) - def my_decomp(angle, wires, **_): - def true_func(): - qp.CNOT(wires) - - with qp.allocate(2, state="any", restored=True) as w: - qp.H(w[0]) - qp.H(w[0]) - qp.X(w[1]) - qp.X(w[1]) - - return - - def false_func(): - with qp.allocate(1, state="any", restored=False) as w: - qp.H(w) - - m = qp.measure(wires[0]) - - qp.cond(m, qp.CNOT)(wires) - - return - - qp.cond(angle > 1.2, true_func, false_func)() - - @qp.qjit(capture=True) - @partial( - qp.transforms.decompose, - gate_set={qp.CNOT, qp.H, qp.X, "Conditional", "MidMeasure"}, - fixed_decomps={qp.CRX: my_decomp}, - num_work_wires=3, - ) - @qp.qnode(qp.device("lightning.qubit", wires=5)) - def circuit(): - qp.CRX(1.7, wires=[0, 1]) - qp.CRX(-7.2, wires=[0, 1]) - return qp.state() - - # CHECK-NOT: CRX - # CHECK-NOT: my_decomp - - # CHECK: allocate - # CHECK: allocate - # CHECK: CNOT - # CHECK: Hadamard - # CHECK: Hadamard - # CHECK: PauliX - # CHECK: PauliX - # CHECK: Hadamard - # CHECK: Measure - # CHECK: CNOT - # CHECK: release - # CHECK: release - print(circuit.mlir_opt) - - qp.decomposition.disable_graph() - - -test_num_work_wires() - - -def test_default_decomps(): - """Test that default decompositions are correctly applied with qjit.""" - qp.decomposition.enable_graph() - - # Toffoli's decomposition to this gateset includes a wire allocation - @qp.qjit(target="mlir", capture=True) - @partial( - qp.transforms.decompose, - gate_set={qp.ops.ChangeOpBasis}, - num_work_wires=1, - ) - @qp.qnode(qp.device("lightning.qubit", wires=4)) - def circuit(): - qp.Toffoli(wires=[0, 1, 2]) - return qp.state() - - # CHECK-NOT: toffoli_elbow - # CHECK-NOT: Toffoli - - # two allocates/releases, for default register + work wires - # CHECK: allocate - # CHECK: allocate - # CHECK: TemporaryAND - # CHECK: release - # CHECK: release - print(circuit.mlir_opt) - - qp.decomposition.disable_graph() - - -test_default_decomps() - - -def test_graph_decomp_registered(): - """Test that the `graph_decomposition` pass is registered correctly.""" - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" - @graph_decomposition(gate_set={qp.RX}) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def catalyst_circuit(): - return - - print(catalyst_circuit.mlir) - - my_transform = qp.transform(pass_name="graph-decomposition") - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" - @my_transform(gate_set=["RX"]) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def pennylane_circuit(): - return - - print(pennylane_circuit.mlir) - - -test_graph_decomp_registered() - - -def test_cpp_decomp_args(): - """Test that the `graph_decomposition` pass lowers arguments to mlir correctly.""" - - def x_to_rx(wire): - qp.RX(np.pi, wire) - - def y_to_ry(wire): - qp.RY(np.pi, wire) - - def h_to_rx_ry(wire): - qp.RX(np.pi / 2, wire) - qp.RY(np.pi / 2, wire) - - @qjit(target="mlir") - # CHECK: "graph-decomposition" with options = { - # CHECK-DAG: "gate-set" = {Hadamard = 1.000000e+00 : f64, RX = 1.000000e+00 : f64, RY = 1.000000e+00 : f64} - # CHECK-DAG: "fixed-decomps" = {PauliX = "x_to_rx", PauliY = "y_to_ry"} - # CHECK-DAG: "alt-decomps" = {Hadamard = ["h_to_rx_ry"]} - # CHECK-DAG: "bytecode-rules" = "{{.*}}decomposition_rules_{{.*}}.mlirbc" - # CHECK: } to {{%.+}} : (!transform.op<"builtin.module">) - @graph_decomposition( - gate_set={qp.RX, qp.H, qp.RY}, - fixed_decomps={qp.X: x_to_rx, qp.Y: y_to_ry}, - alt_decomps={qp.H: [h_to_rx_ry]}, - _builtin_rule_path="/decomp_rules.mlirbc", - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - return - - print(circuit.mlir) - - -test_cpp_decomp_args() - - -def test_cpp_decomp_empty_args(): - """ - Test that the `graph_decomposition` pass correctly handled arg lowering when no values are - supplied. - """ - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" - # CHECK-NOT: fixed-decomps - # CHECK-NOT: alt-decomps - # CHECK: "bytecode-rules" = "{{.*}}/decomposition_rules{{.*}}.mlirbc" - @graph_decomposition(gate_set={qp.RX}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit(): - return - - print(circuit.mlir) - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" - # CHECK-NOT: fixed-decomps - # CHECK-NOT: alt-decomps - # CHECK: "bytecode-rules" = "{{.*}}/decomposition_rules{{.*}}.mlirbc" - @graph_decomposition(gate_set={qp.RX}, fixed_decomps={}, alt_decomps={}) - @qp.qnode(qp.device("lightning.qubit", wires=1)) - def circuit2(): - return - - print(circuit2.mlir) - - -test_cpp_decomp_empty_args() - - -def test_cpp_decomp_string_op_names(): - """Test that cpp decomp args work with string op names.""" - - def y_to_xz(wires): - qp.RX(np.pi, wires) - qp.RZ(np.pi, wires) - - @qjit(target="mlir", capture=True) - # CHECK: transform.apply_registered_pass "graph-decomposition" with options = { - # CHECK-DAG: "fixed-decomps" = {PauliX = "{{.*}}", PauliZ = "{{.*}}"} - # CHECK-DAG: "alt-decomps" = {PauliY = ["{{.*}}", "y_to_xz"]} - # CHECK: } to {{%.+}} : (!transform.op<"builtin.module">) - @graph_decomposition( - gate_set={"RX", "RY", "RZ"}, - fixed_decomps={ - "X": lambda wires: qp.RX(np.pi, wires), - "PauliZ": lambda wires: qp.RZ(np.pi, wires), - }, - alt_decomps={ - "PauliY": [ - lambda wires: qp.RY(np.pi, wires), - y_to_xz, - ] - }, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - return - - print(circuit.mlir) - - -test_cpp_decomp_string_op_names() - - -def test_cpp_decomp_builtin_rules(): - """Test that cpp decomp applies builtin rules.""" - - @qjit(target="mlir", capture=True) - @graph_decomposition( - gate_set={qp.RX, qp.RY, qp.RZ, qp.GlobalPhase}, - ) - @qp.qnode(qp.device("lightning.qubit", wires=2)) - def circuit(): - # CHECK-NOT: PauliX - # CHECK-NOT: PauliY - # CHECK-NOT: PauliZ - # CHECK-DAG: RX - # CHECK-DAG: RY - # CHECK-DAG: RZ - qp.X(0) - qp.Y(1) - qp.Z(0) - return qp.probs() - - print(circuit.mlir_opt) - - -test_cpp_decomp_builtin_rules() - - -def test_cpp_decomp_user_rules(): - """Test that cpp decomp applies user rules.""" - - @decomposition_rule(is_qreg=True, op_type="PauliY") - def y_to_rx(wire): - qp.RX(np.pi, wire) - - @decomposition_rule(is_qreg=True, op_type="PauliZ") - def z_to_rx(wire): - qp.RX(np.pi, wire) - - @qp.qjit(target="mlir", capture=True) - @graph_decomposition( - gate_set={qp.RX}, fixed_decomps={qp.Y: y_to_rx}, alt_decomps={qp.Z: [z_to_rx]} - ) - @qp.qnode(qp.device("null.qubit", wires=1)) - def circuit(): - y_to_rx(jax.core.ShapedArray((1,), int)) - z_to_rx(jax.core.ShapedArray((1,), int)) - # CHECK-NOT: PauliY - # CHECK-NOT: PauliZ - # CHECK: RX - # CHECK: RX - # CHECK: return - qp.Y(0) - qp.Z(0) - return qp.probs() - - print(circuit.mlir_opt) - - -test_cpp_decomp_user_rules() - - -def test_cpp_decomp_user_rule_cleanup(): - """Test that user rules do not pollute the IR after the quantum compilation stage.""" - - @decomposition_rule(is_qreg=True, op_type="PauliX") - def x_to_h(wire): - return qp.H(wire) - - @qjit(capture=True) - @graph_decomposition(gate_set={qp.H}, fixed_decomps={qp.X: x_to_h}) - @qp.qnode(qp.device("null.qubit", wires=1)) - def circuit(): - # CHECK-NOT: PauliX - # CHECK-NOT: x_to_h - x_to_h(jax.core.ShapedArray((1,), int)) - qp.X(0) - - print(circuit.mlir_opt) - - -test_cpp_decomp_user_rule_cleanup() - - -def test_paulirot_python_decomposition(): - """Test that paulirots are decomposed by the mlir graph.""" - - @qjit(capture=True) - @graph_decomposition(gate_set={qp.H, qp.MultiRZ, qp.GlobalPhase, qp.RX}) - @qp.qnode(qp.device("null.qubit", wires=4)) - def circuit(): - qp.PauliRot(0.9, "XZXY", [0, 1, 2, 3]) - return qp.probs() - - print(circuit.mlir_opt) - - # CHECK-NOT: quantum.paulirot - # CHECK: Hadamard - # CHECK: Hadamard - # CHECK: RX - # CHECK: MultiRZ - # CHECK: Hadamard - # CHECK: Hadamard - # CHECK: RX - - -test_paulirot_python_decomposition() From 59bb9f4c76845c5d0329c6abfd616cde6f58c0e6 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 23 Jul 2026 15:08:15 -0400 Subject: [PATCH 19/36] generic pytest layout --- frontend/test/pytest/test_decomposition.py | 69 ++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/frontend/test/pytest/test_decomposition.py b/frontend/test/pytest/test_decomposition.py index fb8fe117eb..91d207de8c 100644 --- a/frontend/test/pytest/test_decomposition.py +++ b/frontend/test/pytest/test_decomposition.py @@ -11,3 +11,72 @@ # 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. + +"""Unit tests for the python decompositions module.""" + +from pathlib import Path + +import pennylane as qp +import pytest + +from catalyst.compiler import _quantum_opt +from catalyst.decomposition.precompile_decomposition_rules import ( + get_abstract_args, + precompile_decomp_rules, +) +from catalyst.decomposition.python_decompositions import python_decomposition_wrapper +from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH + + +class TestGenericUtilities: + """Tests for common decomposition rule lowering utilities.""" + + def test_paulirot(self): + """Test that the QPD wrapper correctly returns the IR as a string.""" + result = python_decomposition_wrapper( + "PauliRot", "PauliRot[f64][3]{pauli_word:XZZ}", ["i32"], [3], {"pauli_word": "XZZ"} + ) + assert isinstance(result, str) + assert "_pauli_rot_decomposition" in result + assert 'target_gate = "PauliRot[f64][3]{pauli_word:XZZ}"' in result + assert "Hadamard" in result + assert "multirz" in result + + def test_multiple_rules(self): + """Test that the python decomposition wrapper supports multiple rules.""" + with qp.decomposition.local_decomps(): + + def test_resources(pauli_word): # pylint: disable=unused-argument + return {qp.X: 1} + + @qp.register_resources(test_resources) + def test_decomp(angle, wires, pauli_word): # pylint: disable=unused-argument + qp.RX(angle, wires[0]) + + qp.add_decomps(qp.PauliRot, test_decomp) + + result = python_decomposition_wrapper( + "PauliRot", "PauliRot[f64][3]{pauli_word:XYX}", ["f64"], [3], {"pauli_word": "XYX"} + ) + + assert "test_decomp" in result + assert "_pauli_rot_decomp" in result + assert 'target_gate = "PauliRot[f64][3]{pauli_word:XYX}"' in result + + +class TestPrecompiled: + """Tests for precompiled decomposition rules.""" + + +class TestTraceTime: + """Placeholder for future tests of trace-time decomposition rule lowering.""" + + +class TestOnDemand: + """ + Test the python wrapper functions used for on-demand, compile-time decomposition rule lowering. + """ + + +if __name__ == "__main__": + pytest.main(["-x", __file__]) From 7bc64f09e1ce64764931326e873aabb787f6e2c3 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 23 Jul 2026 16:02:19 -0400 Subject: [PATCH 20/36] rule.compute_resources need dynamic args and wires too, not just static data --- .../catalyst/decomposition/python_decompositions.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/decomposition/python_decompositions.py b/frontend/catalyst/decomposition/python_decompositions.py index 55b95dc5e9..cd8de269ed 100644 --- a/frontend/catalyst/decomposition/python_decompositions.py +++ b/frontend/catalyst/decomposition/python_decompositions.py @@ -135,7 +135,7 @@ def getID(self): return ID_string -def collect_resources_for_op(op_name, static_data): +def collect_resources_for_op(op_name, dummy_dynamic_args, dummy_wires, static_data): decomp_rules = list(qp.decomposition.list_decomps(op_name)) # map rules to resource resources, in a more generic format @@ -144,7 +144,7 @@ def collect_resources_for_op(op_name, static_data): for rule in decomp_rules: # The `compute_resources` function's signature is the same as the Operator2 signature # for the original op of the rule - resources = rule.compute_resources(**static_data) + resources = rule.compute_resources(*dummy_dynamic_args, *dummy_wires, **static_data) name_to_resources[rule.name] = resources.gate_counts name_to_resource_ids[rule.name] = { GraphOpID(op).getID(): count for op, count in resources.gate_counts.items() @@ -157,9 +157,12 @@ def python_decomposition(op_name, op_id, dynamic_shape, wire_lens, static_data) """Python decomposition rule lowering.""" # TODO update docstring device = qp.device("null.qubit", wires=sum(wire_lens)) - wires = tuple(jnp.array(range(length), dtype=int) for length in wire_lens) + dummy_wires = tuple(jnp.array(range(length), dtype=int) for length in wire_lens) + dummy_dynamic_args = get_dummy_values_for_container(dynamic_shape) - _, name_to_resource_ids, decomp_rules = collect_resources_for_op(op_name, static_data) + _, name_to_resource_ids, decomp_rules = collect_resources_for_op( + op_name, dummy_dynamic_args, dummy_wires, static_data + ) def rule_to_subroutine(rule): def decomp_rule(*args, **kwargs): @@ -179,7 +182,7 @@ def decomp_rule(*args, **kwargs): @qp.qnode(device=device) def circuit(): for subroutine in subroutines: - subroutine(*get_dummy_values_for_container(dynamic_shape), wires=wires) + subroutine(*dummy_dynamic_args, wires=dummy_wires) module = circuit.mlir_module From 54c3fca5036c75fb0af409bc8c24615cf4a60d69 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 23 Jul 2026 16:23:40 -0400 Subject: [PATCH 21/36] unify type utils --- .../decomposition/python_decompositions.py | 25 +++++------------ ...{type_stringify_utils.py => type_utils.py} | 27 +++++++++++-------- 2 files changed, 23 insertions(+), 29 deletions(-) rename frontend/catalyst/decomposition/{type_stringify_utils.py => type_utils.py} (69%) diff --git a/frontend/catalyst/decomposition/python_decompositions.py b/frontend/catalyst/decomposition/python_decompositions.py index cd8de269ed..298c124940 100644 --- a/frontend/catalyst/decomposition/python_decompositions.py +++ b/frontend/catalyst/decomposition/python_decompositions.py @@ -18,38 +18,27 @@ # pylint: disable=protected-access,bare-except -import warnings - import jax.numpy as jnp import pennylane as qp from jax._src.lib.mlir import ir from jaxlib.mlir.dialects.builtin import ModuleOp -from catalyst.decomposition.type_stringify_utils import mlir_stringify_type +from catalyst.decomposition.type_utils import ( + _MLIR_DTYPES_TO_PY_DTYPES, + _PY_DTYPES_TO_MLIR_DTYPES, + mlir_stringify_type, +) from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval -_MLIR_DTYPES = { - "i1": jnp.bool_, - "i8": jnp.int8, - "i16": jnp.int16, - "i32": jnp.int32, - "i64": jnp.int64, - "f16": jnp.float16, - "f32": jnp.float32, - "f64": jnp.float64, - "complex": jnp.complex64, - "complex": jnp.complex128, -} - def get_dummy_values_for_container(container): """Given a container of python types, replace the types with corresponding dummy values.""" dummy_args = [] for dtype in container: if isinstance(dtype, str): - if dtype in _MLIR_DTYPES: + if dtype in _MLIR_DTYPES_TO_PY_DTYPES: count = 1 - dtype = _MLIR_DTYPES[dtype] + dtype = _MLIR_DTYPES_TO_PY_DTYPES[dtype] elif dtype.startswith("tensor"): # tensor<{number}x{type}> dtype = dtype.removeprefix("tensor<") diff --git a/frontend/catalyst/decomposition/type_stringify_utils.py b/frontend/catalyst/decomposition/type_utils.py similarity index 69% rename from frontend/catalyst/decomposition/type_stringify_utils.py rename to frontend/catalyst/decomposition/type_utils.py index 7cd0ab570f..70d7f97ebc 100644 --- a/frontend/catalyst/decomposition/type_stringify_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -12,23 +12,28 @@ # See the License for the specific language governing permissions and # limitations under the License. -import jax.numpy as jnp +import numpy as np import pennylane as qp -from catalyst.utils.exceptions import CompileError +_MLIR_DTYPES_TO_PY_DTYPES = { + "i1": np.bool_, + "i8": np.int8, + "i16": np.int16, + "i32": np.int32, + "i64": np.int64, + "f16": np.float16, + "f32": np.float32, + "f64": np.float64, + "complex": np.complex64, + "complex": np.complex128, +} - -def _py_dtype_to_mlir_type_string(python_dtype: type): - match python_dtype: - case jnp.float64: - return "f64" - case _: - raise CompileError("Unknown data type") +_PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} def _stringify_shaped_type(shape: tuple, dim: int, element_type): if dim + 1 == len(shape): - inner_content = _py_dtype_to_mlir_type_string(element_type) + inner_content = _PY_DTYPES_TO_MLIR_DTYPES[element_type] else: inner_content = _stringify_shaped_type(shape, dim + 1, element_type) length = shape[dim] @@ -39,6 +44,6 @@ def mlir_stringify_type(dtype: qp.typing.AbstractArray): assert isinstance(dtype, qp.typing.AbstractArray) element_type = dtype.dtype.type if dtype.shape == (): - return _py_dtype_to_mlir_type_string(element_type) + return _PY_DTYPES_TO_MLIR_DTYPES[element_type] else: return _stringify_shaped_type(dtype.shape, 0, element_type) From c10bc23b3e4439f78142e055c10d492eb06c2158 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 23 Jul 2026 16:03:16 -0400 Subject: [PATCH 22/36] update docs for frontend --- .../decomposition/python_decompositions.py | 46 ++++++++++++++----- 1 file changed, 34 insertions(+), 12 deletions(-) diff --git a/frontend/catalyst/decomposition/python_decompositions.py b/frontend/catalyst/decomposition/python_decompositions.py index 298c124940..b0ac95edae 100644 --- a/frontend/catalyst/decomposition/python_decompositions.py +++ b/frontend/catalyst/decomposition/python_decompositions.py @@ -13,7 +13,7 @@ # limitations under the License. """ -This module provides infrastructure for compile-time lowering of decomposition rules via python. +This module provides infrastructure for lowering decomposition rules via python. """ # pylint: disable=protected-access,bare-except @@ -77,6 +77,7 @@ class GraphOpID: """ def __init__(self, op: qp.core.Operator2, uid=None): + """Create a new GraphOpId.""" assert isinstance( op, qp.core.Operator2 ), "Graph-based decomposition expects an Operator2 instance" @@ -88,31 +89,44 @@ def __init__(self, op: qp.core.Operator2, uid=None): self.static_data = self.parse_static_data() self.extra_data = uid - def parse_dynamic_shape(self): + def parse_dynamic_shape(self) -> list: + """Return the dynamic shape as a list of dtypes.""" return list(self.op.dynamic_args.values()) - def parse_wire_lens(self): + def parse_wire_lens(self) -> list[int]: + """Return the length of each of the wire args.""" return list(map(len, self.op.wire_args.values())) - def parse_static_data(self): + def parse_static_data(self) -> dict: + """Return a dictionary of names to static data values.""" return { static_argname: getattr(self.op, static_argname) for static_argname in self.op.compilable_argnames } - def get_operator_name(self): + def get_operator_name(self) -> str: + """Return the name of the operator.""" return self.operator_name - def get_dynamic_shape_id_format(self): + def get_dynamic_shape_id_format(self) -> str: + """Return the dynamic shape formatted for GraphOpId.""" return f"[{','.join(map(mlir_stringify_type, self.dynamic_shape))}]" - def get_wire_lens_id_format(self): + def get_wire_lens_id_format(self) -> str: + """Return the wire lengths formatted for GraphOpId.""" return f"[{','.join(map(str, self.wire_lens))}]" - def get_static_data_id_format(self): + def get_static_data_id_format(self) -> str: + """Return the static data formatted for GraphOpId.""" return f"{{{','.join(f'{k}:{v}' for k, v in self.static_data.items())}}}" - def getID(self): + def getID(self) -> str: + """ + Return the GraphOpId as a string. + + NOTE: do not modify this method without also modifying the corresponding DecomposableGate + interface in MLIR. + """ ID_string = ( self.get_operator_name() + self.get_dynamic_shape_id_format() @@ -125,6 +139,11 @@ def getID(self): def collect_resources_for_op(op_name, dummy_dynamic_args, dummy_wires, static_data): + """ + Return resource data for all decomposition rules associated to op_name. + + This includes a dictionary + """ decomp_rules = list(qp.decomposition.list_decomps(op_name)) # map rules to resource resources, in a more generic format @@ -143,8 +162,11 @@ def collect_resources_for_op(op_name, dummy_dynamic_args, dummy_wires, static_da def python_decomposition(op_name, op_id, dynamic_shape, wire_lens, static_data) -> ModuleOp: - """Python decomposition rule lowering.""" - # TODO update docstring + """ + Return a ModuleOp containing the decomposition rules for an operator instance. + + The decomposition rules will be decorated with appropriate resource and target_gate attributes. + """ device = qp.device("null.qubit", wires=sum(wire_lens)) dummy_wires = tuple(jnp.array(range(length), dtype=int) for length in wire_lens) dummy_dynamic_args = get_dummy_values_for_container(dynamic_shape) @@ -201,5 +223,5 @@ def update_funcop_attributes(op): def python_decomposition_wrapper(op_name, op_id, dynamic_shape, wire_lens, static_data) -> str: - """Generic decomposition wrapper.""" + """Return a string MLIR module containing the decomposition rules for an operator instance.""" return str(python_decomposition(op_name, op_id, dynamic_shape, wire_lens, static_data)) From 185628a347dbd6538254fca5e60d8047feb21869 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Thu, 23 Jul 2026 16:48:06 -0400 Subject: [PATCH 23/36] rename decomposition rule functions --- .../{python_decompositions.py => decomposition_rules.py} | 8 +++++--- .../decomposition/precompile_decomposition_rules.py | 4 ++-- frontend/test/lit/GraphDecomposition/TestMultiDecomp.mlir | 4 ++-- frontend/test/pytest/test_QPD.py | 6 +++--- frontend/test/pytest/test_decomposition.py | 6 +++--- .../QuantumPythonDecompositions/PythonFunction.cpp | 2 +- 6 files changed, 16 insertions(+), 14 deletions(-) rename frontend/catalyst/decomposition/{python_decompositions.py => decomposition_rules.py} (96%) diff --git a/frontend/catalyst/decomposition/python_decompositions.py b/frontend/catalyst/decomposition/decomposition_rules.py similarity index 96% rename from frontend/catalyst/decomposition/python_decompositions.py rename to frontend/catalyst/decomposition/decomposition_rules.py index b0ac95edae..2ff460d6fc 100644 --- a/frontend/catalyst/decomposition/python_decompositions.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -161,7 +161,7 @@ def collect_resources_for_op(op_name, dummy_dynamic_args, dummy_wires, static_da return name_to_resources, name_to_resource_ids, decomp_rules -def python_decomposition(op_name, op_id, dynamic_shape, wire_lens, static_data) -> ModuleOp: +def compile_decomposition_rules(op_name, op_id, dynamic_shape, wire_lens, static_data) -> ModuleOp: """ Return a ModuleOp containing the decomposition rules for an operator instance. @@ -222,6 +222,8 @@ def update_funcop_attributes(op): return module -def python_decomposition_wrapper(op_name, op_id, dynamic_shape, wire_lens, static_data) -> str: +def compile_decomposition_rules_wrapper( + op_name, op_id, dynamic_shape, wire_lens, static_data +) -> str: """Return a string MLIR module containing the decomposition rules for an operator instance.""" - return str(python_decomposition(op_name, op_id, dynamic_shape, wire_lens, static_data)) + return str(compile_decomposition_rules(op_name, op_id, dynamic_shape, wire_lens, static_data)) diff --git a/frontend/catalyst/decomposition/precompile_decomposition_rules.py b/frontend/catalyst/decomposition/precompile_decomposition_rules.py index 659c7db03b..27405e0834 100644 --- a/frontend/catalyst/decomposition/precompile_decomposition_rules.py +++ b/frontend/catalyst/decomposition/precompile_decomposition_rules.py @@ -21,7 +21,7 @@ from pennylane.operation import Operator, Operator2 from catalyst.compiler import _quantum_opt -from catalyst.decomposition.python_decompositions import GraphOpID, python_decomposition +from catalyst.decomposition.decomposition_rules import GraphOpID, compile_decomposition_rules from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH # TODO: Uncomment dynamic size wires ops once they are supported @@ -169,7 +169,7 @@ def precompile_decomp_rules(decomp_file_path: str = BYTECODE_FILE_PATH): # we cannot precompile if the rule takes static data continue - mlir_rules = python_decomposition( + mlir_rules = compile_decomposition_rules( op.__name__, GraphOpID(op).getID(), dynamic_data, wire_lens, {} ) diff --git a/frontend/test/lit/GraphDecomposition/TestMultiDecomp.mlir b/frontend/test/lit/GraphDecomposition/TestMultiDecomp.mlir index a6f04685de..5ba159fe1c 100644 --- a/frontend/test/lit/GraphDecomposition/TestMultiDecomp.mlir +++ b/frontend/test/lit/GraphDecomposition/TestMultiDecomp.mlir @@ -12,9 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRX=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes FIRST +// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRX=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes FIRST -// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRX=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"},graph-decomposition{gate-set=testRZ=1.0,testRY=1.0,GlobalPhase=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes SECOND +// RUN: catalyst --tool=opt --split-input-file --pass-pipeline='builtin.module(graph-decomposition{gate-set=testRX=1.0 bytecode-rules="%BYTECODE_PATH"},graph-decomposition{gate-set=testRZ=1.0,testRY=1.0 bytecode-rules="%BYTECODE_PATH"})' %s | FileCheck %s --check-prefixes SECOND func.func @circuit() -> !quantum.bit { %0 = quantum.alloc(2) : !quantum.reg diff --git a/frontend/test/pytest/test_QPD.py b/frontend/test/pytest/test_QPD.py index 3c73473740..0630c701e1 100644 --- a/frontend/test/pytest/test_QPD.py +++ b/frontend/test/pytest/test_QPD.py @@ -17,7 +17,7 @@ import pennylane as qp import pytest -from catalyst.decomposition.python_decompositions import python_decomposition_wrapper +from catalyst.decomposition.python_decompositions import compile_decomposition_rules_wrapper class TestQPD: @@ -25,7 +25,7 @@ class TestQPD: def test_paulirot_wrapper(self): """Test that the paulirot QPD wrapper correctly returns the IR as a string.""" - result = python_decomposition_wrapper( + result = compile_decomposition_rules_wrapper( "PauliRot", "PauliRot[f64][3]{pauli_word:XZZ}", [0.4], [3], {"pauli_word": "XZZ"} ) assert isinstance(result, str) @@ -46,7 +46,7 @@ def test_decomp(angle, wires, pauli_word): # pylint: disable=unused-argument qp.add_decomps(qp.PauliRot, test_decomp) - result = python_decomposition_wrapper( + result = compile_decomposition_rules_wrapper( "PauliRot", "PauliRot[f64][3]{pauli_word:XYX}", [float], [3], {"pauli_word": "XYX"} ) diff --git a/frontend/test/pytest/test_decomposition.py b/frontend/test/pytest/test_decomposition.py index 91d207de8c..8aa1d0ec53 100644 --- a/frontend/test/pytest/test_decomposition.py +++ b/frontend/test/pytest/test_decomposition.py @@ -24,7 +24,7 @@ get_abstract_args, precompile_decomp_rules, ) -from catalyst.decomposition.python_decompositions import python_decomposition_wrapper +from catalyst.decomposition.python_decompositions import compile_decomposition_rules_wrapper from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH @@ -33,7 +33,7 @@ class TestGenericUtilities: def test_paulirot(self): """Test that the QPD wrapper correctly returns the IR as a string.""" - result = python_decomposition_wrapper( + result = compile_decomposition_rules_wrapper( "PauliRot", "PauliRot[f64][3]{pauli_word:XZZ}", ["i32"], [3], {"pauli_word": "XZZ"} ) assert isinstance(result, str) @@ -55,7 +55,7 @@ def test_decomp(angle, wires, pauli_word): # pylint: disable=unused-argument qp.add_decomps(qp.PauliRot, test_decomp) - result = python_decomposition_wrapper( + result = compile_decomposition_rules_wrapper( "PauliRot", "PauliRot[f64][3]{pauli_word:XYX}", ["f64"], [3], {"pauli_word": "XYX"} ) diff --git a/mlir/lib/Quantum/Transforms/QuantumPythonDecompositions/PythonFunction.cpp b/mlir/lib/Quantum/Transforms/QuantumPythonDecompositions/PythonFunction.cpp index 222ab8ef1f..2dd01b5b02 100644 --- a/mlir/lib/Quantum/Transforms/QuantumPythonDecompositions/PythonFunction.cpp +++ b/mlir/lib/Quantum/Transforms/QuantumPythonDecompositions/PythonFunction.cpp @@ -77,7 +77,7 @@ std::string pythonRuleLowering(catalyst::quantum::DecomposableGate op) QuantumPythonDecompositions::PyInterpreterGuard guard; std::string mlirText = guard.withGil([&] -> std::string { const char *moduleName = "catalyst.decomposition.python_decompositions"; - const char *functionName = "python_decomposition_wrapper"; + const char *functionName = "compile_decomposition_rules"; try { nb::module_ wrapperModule = nb::module_::import_(moduleName); From 9b6d859c0ba595980b90edf0312d92e412d2281e Mon Sep 17 00:00:00 2001 From: paul0403 Date: Fri, 24 Jul 2026 09:24:22 -0400 Subject: [PATCH 24/36] move get_dummy_args to util file --- .../decomposition/decomposition_rules.py | 25 +------------------ frontend/catalyst/decomposition/type_utils.py | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+), 24 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index 2ff460d6fc..15e7314405 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -26,35 +26,12 @@ from catalyst.decomposition.type_utils import ( _MLIR_DTYPES_TO_PY_DTYPES, _PY_DTYPES_TO_MLIR_DTYPES, + get_dummy_values_for_container, mlir_stringify_type, ) from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval -def get_dummy_values_for_container(container): - """Given a container of python types, replace the types with corresponding dummy values.""" - dummy_args = [] - for dtype in container: - if isinstance(dtype, str): - if dtype in _MLIR_DTYPES_TO_PY_DTYPES: - count = 1 - dtype = _MLIR_DTYPES_TO_PY_DTYPES[dtype] - elif dtype.startswith("tensor"): - # tensor<{number}x{type}> - dtype = dtype.removeprefix("tensor<") - dtype = dtype.remove_suffice(">") - count, dtype = dtype.split("x") - else: - raise ValueError(f"Unknown dtype {dtype}.") - else: - count = 1 - dtype = jnp.dtype(dtype) - - dummy_args.append(jnp.zeros((count,), dtype=dtype)) - - return tuple(dummy_args) - - class GraphOpID: """ Return the graph operator id for the operator2 instance `op`. diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index 70d7f97ebc..1ba0407de8 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import jax.numpy as jnp import numpy as np import pennylane as qp @@ -47,3 +48,27 @@ def mlir_stringify_type(dtype: qp.typing.AbstractArray): return _PY_DTYPES_TO_MLIR_DTYPES[element_type] else: return _stringify_shaped_type(dtype.shape, 0, element_type) + + +def get_dummy_values_for_container(container): + """Given a container of python types, replace the types with corresponding dummy values.""" + dummy_args = [] + for dtype in container: + if isinstance(dtype, str): + if dtype in _MLIR_DTYPES_TO_PY_DTYPES: + count = 1 + dtype = _MLIR_DTYPES_TO_PY_DTYPES[dtype] + elif dtype.startswith("tensor"): + # tensor<{number}x{type}> + dtype = dtype.removeprefix("tensor<") + dtype = dtype.remove_suffice(">") + count, dtype = dtype.split("x") + else: + raise ValueError(f"Unknown dtype {dtype}.") + else: + count = 1 + dtype = jnp.dtype(dtype) + + dummy_args.append(jnp.zeros((count,), dtype=dtype)) + + return tuple(dummy_args) From 0dfd9b5cd2bd3605aa0875afaa0c62418e793335 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Fri, 24 Jul 2026 10:03:46 -0400 Subject: [PATCH 25/36] dummy arg maker works with lists instead of tensors --- frontend/catalyst/decomposition/type_utils.py | 32 +++++++------------ 1 file changed, 11 insertions(+), 21 deletions(-) diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index 1ba0407de8..26303f2d9b 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -51,24 +51,14 @@ def mlir_stringify_type(dtype: qp.typing.AbstractArray): def get_dummy_values_for_container(container): - """Given a container of python types, replace the types with corresponding dummy values.""" - dummy_args = [] - for dtype in container: - if isinstance(dtype, str): - if dtype in _MLIR_DTYPES_TO_PY_DTYPES: - count = 1 - dtype = _MLIR_DTYPES_TO_PY_DTYPES[dtype] - elif dtype.startswith("tensor"): - # tensor<{number}x{type}> - dtype = dtype.removeprefix("tensor<") - dtype = dtype.remove_suffice(">") - count, dtype = dtype.split("x") - else: - raise ValueError(f"Unknown dtype {dtype}.") - else: - count = 1 - dtype = jnp.dtype(dtype) - - dummy_args.append(jnp.zeros((count,), dtype=dtype)) - - return tuple(dummy_args) + """ + Converts a nested list of dtype strings into a matching nested list of jnp.zeros. + """ + if isinstance(container, (list, tuple)): + return [get_dummy_values_for_container(item) for item in container] + elif isinstance(container, str): + return jnp.zeros((), dtype=_MLIR_DTYPES_TO_PY_DTYPES[container]) + else: + raise TypeError( + f"Unexpected type in container when creating dummy values: {type(container)}" + ) From e86c1da1f2835fecf69a07758d090ed79fe753a0 Mon Sep 17 00:00:00 2001 From: River McCubbin Date: Fri, 24 Jul 2026 10:40:43 -0400 Subject: [PATCH 26/36] update and test get_dummy_values_for_container --- frontend/catalyst/decomposition/type_utils.py | 57 ++++++++++++------- frontend/test/pytest/test_decomposition.py | 46 ++++++++++++++- 2 files changed, 82 insertions(+), 21 deletions(-) diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index 26303f2d9b..3739715c45 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -12,21 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Type handling utilities for decomposition rule lowering.""" + import jax.numpy as jnp -import numpy as np import pennylane as qp +from jax.core import ShapedArray _MLIR_DTYPES_TO_PY_DTYPES = { - "i1": np.bool_, - "i8": np.int8, - "i16": np.int16, - "i32": np.int32, - "i64": np.int64, - "f16": np.float16, - "f32": np.float32, - "f64": np.float64, - "complex": np.complex64, - "complex": np.complex128, + "i1": jnp.bool_, + "i8": jnp.int8, + "i16": jnp.int16, + "i32": jnp.int32, + "i64": jnp.int64, + "f16": jnp.float16, + "f32": jnp.float32, + "f64": jnp.float64, + "complex": jnp.complex64, + "complex": jnp.complex128, } _PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} @@ -42,6 +44,7 @@ def _stringify_shaped_type(shape: tuple, dim: int, element_type): def mlir_stringify_type(dtype: qp.typing.AbstractArray): + """Return a string representation of the given data type.""" assert isinstance(dtype, qp.typing.AbstractArray) element_type = dtype.dtype.type if dtype.shape == (): @@ -52,13 +55,27 @@ def mlir_stringify_type(dtype: qp.typing.AbstractArray): def get_dummy_values_for_container(container): """ - Converts a nested list of dtype strings into a matching nested list of jnp.zeros. + Given a container of python or MLIR types, replace the types with corresponding dummy values. + + Each item in the container must be representible as an MLIR tensor with at most one layer of + nesting, i.e. cannot be nested and all elements must be of the same type. + Ex. + [[float, float], [int, int, int], [int32, int32, int32, int32]] """ - if isinstance(container, (list, tuple)): - return [get_dummy_values_for_container(item) for item in container] - elif isinstance(container, str): - return jnp.zeros((), dtype=_MLIR_DTYPES_TO_PY_DTYPES[container]) - else: - raise TypeError( - f"Unexpected type in container when creating dummy values: {type(container)}" - ) + + def handle_item(item): + if isinstance(item, (list, tuple)): + return jnp.zeros(len(item), dtype=handle_item(item[0]).dtype) + if isinstance(item, ShapedArray): + return jnp.zeros(item.shape[0], dtype=item.dtype) + elif isinstance(item, str): + return jnp.zeros((), dtype=_MLIR_DTYPES_TO_PY_DTYPES[item]) + elif isinstance(item, (type, jnp.dtype)): + try: + return jnp.zeros((), jnp.dtype(item)) + except TypeError: + raise TypeError( + f"Unexpected type in container when creating dummy values: {type(item)}" + ) + + return tuple(handle_item(item) for item in container) diff --git a/frontend/test/pytest/test_decomposition.py b/frontend/test/pytest/test_decomposition.py index 8aa1d0ec53..72870a3afa 100644 --- a/frontend/test/pytest/test_decomposition.py +++ b/frontend/test/pytest/test_decomposition.py @@ -16,21 +16,65 @@ from pathlib import Path +import jax.numpy as jnp import pennylane as qp import pytest +from jax.core import ShapedArray from catalyst.compiler import _quantum_opt +from catalyst.decomposition.decomposition_rules import ( + compile_decomposition_rules_wrapper, +) from catalyst.decomposition.precompile_decomposition_rules import ( get_abstract_args, precompile_decomp_rules, ) -from catalyst.decomposition.python_decompositions import compile_decomposition_rules_wrapper +from catalyst.decomposition.type_utils import get_dummy_values_for_container from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH class TestGenericUtilities: """Tests for common decomposition rule lowering utilities.""" + def test_get_dummy_values_types(self): + """Test that get_dummy_values_for_container handles MLIR and Python types correctly.""" + python_types = [int, float, jnp.dtype("int32"), bool, complex] + result = get_dummy_values_for_container(python_types) + + assert result[0].dtype == "int64" + assert result[1].dtype == "float64" + assert result[2].dtype == "int32" + assert result[3].dtype == "bool" + assert result[4].dtype == "complex128" + + mlir_types = ["i1", "i32", "f64", "complex", "complex"] + result = get_dummy_values_for_container(mlir_types) + + assert result[0].dtype == "bool" + assert result[1].dtype == "int32" + assert result[2].dtype == "float64" + assert result[3].dtype == "complex64" + assert result[4].dtype == "complex128" + + def test_get_dummy_values_shapes(self): + """Test that get_dummy_values_for_container handles MLIR and python shapes correctly.""" + python_shapes = [bool, [float, float], [int], ShapedArray((4,), "int32")] + result = get_dummy_values_for_container(python_shapes) + print(result) + + assert result[0].shape == () + assert result[1].shape == (2,) + assert result[2].shape == (1,) + assert result[3].shape == (4,) + + mlir_types = ["i32", ["f64", "f64"], ["i1", "i1", "i1"]] + result = get_dummy_values_for_container(mlir_types) + print(result) + + assert result[0].shape == () + assert result[1].shape == (2,) + assert result[2].shape == (3,) + def test_paulirot(self): """Test that the QPD wrapper correctly returns the IR as a string.""" result = compile_decomposition_rules_wrapper( From 6b82b7f874bfc577f16481b9fff158cc5e898ee6 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Fri, 24 Jul 2026 10:45:23 -0400 Subject: [PATCH 27/36] Decomp/lower time rules --- .../decomposition/decomposition_rules.py | 33 +++++++++++++++++++ frontend/catalyst/decomposition/type_utils.py | 27 +++++++++++---- 2 files changed, 53 insertions(+), 7 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index 15e7314405..fc784907df 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -18,6 +18,8 @@ # pylint: disable=protected-access,bare-except +from collections import deque + import jax.numpy as jnp import pennylane as qp from jax._src.lib.mlir import ir @@ -27,6 +29,7 @@ _MLIR_DTYPES_TO_PY_DTYPES, _PY_DTYPES_TO_MLIR_DTYPES, get_dummy_values_for_container, + listify_type, mlir_stringify_type, ) from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval @@ -85,6 +88,9 @@ def get_operator_name(self) -> str: """Return the name of the operator.""" return self.operator_name + def get_dynamic_shape_as_list(self) -> list[str]: + return [listify_type(t) for t in self.dynamic_shape] + def get_dynamic_shape_id_format(self) -> str: """Return the dynamic shape formatted for GraphOpId.""" return f"[{','.join(map(mlir_stringify_type, self.dynamic_shape))}]" @@ -204,3 +210,30 @@ def compile_decomposition_rules_wrapper( ) -> str: """Return a string MLIR module containing the decomposition rules for an operator instance.""" return str(compile_decomposition_rules(op_name, op_id, dynamic_shape, wire_lens, static_data)) + + +def fetch_all_reachable_decomposition_rules_from_op(op_name, dynamic_shape, wire_lens, static_data): + q = deque() + start = (op_name, dynamic_shape, wire_lens, static_data) + q.append(start) + visited = [start] + while len(q) != 0: + this_name, this_dynamic_shape, this_wire_lens, this_static_data = q.popleft() + dummy_wires = tuple(jnp.array(range(length), dtype=int) for length in this_wire_lens) + dummy_dynamic_args = get_dummy_values_for_container(this_dynamic_shape) + resources, _, _ = collect_resources_for_op( + this_name, dummy_dynamic_args, dummy_wires, this_static_data + ) + for _rule_name, resource in resources.items(): + for op, _count in resource.items(): + graph_op_id = GraphOpID(op) + probe = ( + graph_op_id.get_operator_name(), + graph_op_id.get_dynamic_shape_as_list(), + graph_op_id.wire_lens, + graph_op_id.static_data, + ) + if not probe in visited: + visited.append(probe) + q.append(probe) + return visited diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index 26303f2d9b..d519156f3b 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -32,22 +32,35 @@ _PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} -def _stringify_shaped_type(shape: tuple, dim: int, element_type): +def _process_shaped_type(shape: tuple, dim: int, element_type, formatter): + """Recursively processes the shape, applying the formatter at each dimension.""" if dim + 1 == len(shape): inner_content = _PY_DTYPES_TO_MLIR_DTYPES[element_type] else: - inner_content = _stringify_shaped_type(shape, dim + 1, element_type) - length = shape[dim] - return f"[{','.join([inner_content] * length)}]" + inner_content = _process_shaped_type(shape, dim + 1, element_type, formatter) + return formatter(inner_content, shape[dim]) -def mlir_stringify_type(dtype: qp.typing.AbstractArray): + +def _convert_type(dtype: qp.typing.AbstractArray, formatter): + """Base function to handle scalar checks before processing the shape.""" assert isinstance(dtype, qp.typing.AbstractArray) element_type = dtype.dtype.type + if dtype.shape == (): return _PY_DTYPES_TO_MLIR_DTYPES[element_type] - else: - return _stringify_shaped_type(dtype.shape, 0, element_type) + + return _process_shaped_type(dtype.shape, 0, element_type, formatter) + + +def mlir_stringify_type(dtype: qp.typing.AbstractArray): + string_formatter = lambda content, length: f"[{','.join([content] * length)}]" + return _convert_type(dtype, string_formatter) + + +def listify_type(dtype: qp.typing.AbstractArray): + list_formatter = lambda content, length: [content] * length + return _convert_type(dtype, list_formatter) def get_dummy_values_for_container(container): From bbf2125a510d84b0092ce9a23133c829e2179f2c Mon Sep 17 00:00:00 2001 From: paul0403 Date: Fri, 24 Jul 2026 12:54:10 -0400 Subject: [PATCH 28/36] fetch funcs --- .../decomposition/decomposition_rules.py | 52 ++++++++++++++++++- .../precompile_decomposition_rules.py | 40 ++------------ 2 files changed, 55 insertions(+), 37 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index fc784907df..056d437acc 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -121,6 +121,40 @@ def getID(self) -> str: return ID_string +def get_rule_funcs_from_module(module: ir.Module) -> list[ir.Operation]: + funcOps = [] + + def find_condition(op): + if op.name == "func.func": + if "target_gate" in op.attributes: + old_attr = op.attributes["sym_name"] + op.attributes["sym_name"] = ir.StringAttr.get( + "__builtin_" + old_attr.value.strip('"'), context=old_attr.context + ) + funcOps.append(op.detach_from_parent()) + return ir.WalkResult.SKIP + return ir.WalkResult.ADVANCE + + module.operation.walk(find_condition) + return funcOps + + +def get_rules_from_module(module: ir.Module) -> str: + """ + Parse and modify decomposition rules from a ModuleOp. + + Args: + module: an MLIR module object containing a FuncOp named `rule_wrapper` to be extracted + + Returns: + str: The string representation of any decomposition rules from `module`, pre-pending the + `__builtin_` prefix to their names. + """ + funcOps = get_rule_funcs_from_module(module) + + return "\n".join(str(funcOp) for funcOp in funcOps) if funcOps else "" + + def collect_resources_for_op(op_name, dummy_dynamic_args, dummy_wires, static_data): """ Return resource data for all decomposition rules associated to op_name. @@ -212,11 +246,18 @@ def compile_decomposition_rules_wrapper( return str(compile_decomposition_rules(op_name, op_id, dynamic_shape, wire_lens, static_data)) -def fetch_all_reachable_decomposition_rules_from_op(op_name, dynamic_shape, wire_lens, static_data): +def fetch_all_reachable_decomposition_rules_from_op( + op_name, op_id, dynamic_shape, wire_lens, static_data +): q = deque() start = (op_name, dynamic_shape, wire_lens, static_data) q.append(start) visited = [start] + rules = [ + *get_rule_funcs_from_module( + compile_decomposition_rules(op_name, op_id, dynamic_shape, wire_lens, static_data) + ) + ] while len(q) != 0: this_name, this_dynamic_shape, this_wire_lens, this_static_data = q.popleft() dummy_wires = tuple(jnp.array(range(length), dtype=int) for length in this_wire_lens) @@ -236,4 +277,11 @@ def fetch_all_reachable_decomposition_rules_from_op(op_name, dynamic_shape, wire if not probe in visited: visited.append(probe) q.append(probe) - return visited + rules.extend( + get_rule_funcs_from_module( + compile_decomposition_rules( + probe[0], graph_op_id.getID(), probe[1], probe[2], probe[3] + ) + ) + ) + return rules diff --git a/frontend/catalyst/decomposition/precompile_decomposition_rules.py b/frontend/catalyst/decomposition/precompile_decomposition_rules.py index 27405e0834..2b393400c2 100644 --- a/frontend/catalyst/decomposition/precompile_decomposition_rules.py +++ b/frontend/catalyst/decomposition/precompile_decomposition_rules.py @@ -21,7 +21,11 @@ from pennylane.operation import Operator, Operator2 from catalyst.compiler import _quantum_opt -from catalyst.decomposition.decomposition_rules import GraphOpID, compile_decomposition_rules +from catalyst.decomposition.decomposition_rules import ( + GraphOpID, + compile_decomposition_rules, + get_rules_from_module, +) from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH # TODO: Uncomment dynamic size wires ops once they are supported @@ -71,40 +75,6 @@ } -def get_rule_funcs_from_module(module: ir.Module) -> list[ir.Operation]: - funcOps = [] - - def find_condition(op): - if op.name == "func.func": - if "target_gate" in op.attributes: - old_attr = op.attributes["sym_name"] - op.attributes["sym_name"] = ir.StringAttr.get( - "__builtin_" + old_attr.value.strip('"'), context=old_attr.context - ) - funcOps.append(op) - return ir.WalkResult.SKIP - return ir.WalkResult.ADVANCE - - module.operation.walk(find_condition) - return funcOps - - -def get_rules_from_module(module: ir.Module) -> str: - """ - Parse and modify decomposition rules from a ModuleOp. - - Args: - module: an MLIR module object containing a FuncOp named `rule_wrapper` to be extracted - - Returns: - str: The string representation of any decomposition rules from `module`, pre-pending the - `__builtin_` prefix to their names. - """ - funcOps = get_rule_funcs_from_module(module) - - return "\n".join(str(funcOp) for funcOp in funcOps) if funcOps else "" - - def get_abstract_args(op_class: type[Operator]) -> list[type]: """ Create jax-compatible abstract args for catalyst DecompositionRules that apply to op_class. From 8f4fbd55466e9d2d96325b8e1a51a358c7e79e12 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 30 Jul 2026 13:32:56 -0400 Subject: [PATCH 29/36] add stuff --- .../decomposition/decomposition_rules.py | 33 +++++++++++-------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index b1f15c954f..881968578a 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -177,7 +177,7 @@ def find_condition(op): op.attributes["sym_name"] = ir.StringAttr.get( "__builtin_" + old_attr.value.strip('"'), context=old_attr.context ) - funcOps.append(op) + funcOps.append(op.detach_from_parent()) return ir.WalkResult.SKIP return ir.WalkResult.ADVANCE @@ -314,9 +314,9 @@ def compile_decomposition_rules_wrapper( def fetch_all_reachable_decomposition_rules_from_op( op_name, op_id, dynamic_shape, wire_lens, static_data, extra_data=None ): - q = deque() + queue = deque() start = (op_name, dynamic_shape, wire_lens, static_data, extra_data) - q.append(start) + queue.append(start) visited = [start] rules = [ *get_rule_funcs_from_module( @@ -326,9 +326,9 @@ def fetch_all_reachable_decomposition_rules_from_op( ) ] - while len(q) != 0: + while len(queue) != 0: this_name, this_dynamic_shape, this_wire_lens, this_static_data, this_extra_data = ( - q.popleft() + queue.popleft() ) this_extra_data = this_extra_data or {} this_kwargs = prepare_dynamic_op_kwargs(this_dynamic_shape, this_wire_lens) @@ -340,18 +340,25 @@ def fetch_all_reachable_decomposition_rules_from_op( graph_op_id = GraphOpID(op) probe = ( graph_op_id.get_operator_name(), - graph_op_id.get_dynamic_shape(), + { + name: mlir_stringify_type(shape) + for name, shape in graph_op_id.get_dynamic_shape().items() + }, graph_op_id.wire_lens, graph_op_id.static_data, + graph_op_id.extra_data, ) if not probe in visited: visited.append(probe) - q.append(probe) - rules.extend( - get_rule_funcs_from_module( - compile_decomposition_rules( - probe[0], graph_op_id.getID(), probe[1], probe[2], probe[3] - ) - ) + queue.append(probe) + + module = compile_decomposition_rules( + probe[0], + graph_op_id.getID(), + probe[1], + probe[2], + probe[3], + probe[4], ) + rules.extend(get_rule_funcs_from_module(module)) return rules From 83161d1ac997f06b3e17ccf6c77067eccacc744e Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 30 Jul 2026 15:18:37 -0400 Subject: [PATCH 30/36] save --- .../decomposition/decomposition_rules.py | 1 - .../from_plxpr/qref_operator2_primitives.py | 20 +++++++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index 075f1853f2..70a1b107ae 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -33,7 +33,6 @@ post_process_concretize_leaves, replace_abstract_wires_with_concrete_wires, ) -from catalyst.from_plxpr.qref_operator2_primitives import _is_custom_op from catalyst.from_plxpr.uid import generate_uid from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 0dc62b35d4..72feda4005 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -23,14 +23,16 @@ from jaxlib.mlir.dialects.stablehlo import ConvertOp as StableHLOConvertOp from pennylane.pytrees import unflatten -from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval - # TODO: remove after jax v0.7.2 upgrade # Mock _ods_cext.globals.register_traceback_file_exclusion due to API conflicts between # Catalyst's MLIR version and the MLIR version used by JAX. The current JAX version has not # yet updated to the latest MLIR, causing compatibility issues. This workaround will be removed # once JAX updates to a compatible MLIR version # pylint: disable=ungrouped-imports +from catalyst.decomposition.decomposition_rules import ( + fetch_all_reachable_decomposition_rules_from_op, +) +from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval from catalyst.jax_extras.patches import mock_attributes from catalyst.jax_primitives import ( extract_scalar, @@ -302,6 +304,20 @@ def _qref_operator_p_lowering(jax_ctx: mlir.LoweringRuleContext, *args, op_cls, qubit_map=qubit_map, ) + # Collect decomp rules + # op_name, op_id, dynamic_shape, wire_lens, static_data, extra_data=None + repack_wire_argnames = [] + for wire_argname in op_cls.wire_argnames: + if wire_argname not in op_cls.hybrid_argnames: + repack_wire_argnames.append(wire_argname) + decomp_rules = fetch_all_reachable_decomposition_rules_from_op( + op_name=op_cls.__name__, + op_id="Bob{phi:[f64],thetas:[f64,f64]}{wires:1,other_wires:2}{bob_word:blah}", + dynamic_shape={"phi": ["f64"], "thetas": ["f64", "f64"]}, + wire_lens={a: b for a, b in zip(repack_wire_argnames, wire_lens, strict=True)}, + static_data=repack_static_data, + ) + return [] From 5d3a8391bc7f4981be08ffc64c223cd1ae8de0a4 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 30 Jul 2026 16:02:16 -0400 Subject: [PATCH 31/36] . --- .../decomposition/decomposition_rules.py | 31 ++++--- frontend/catalyst/decomposition/type_utils.py | 88 ++++++++++++++----- 2 files changed, 86 insertions(+), 33 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index 70a1b107ae..5a874e9888 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -21,6 +21,7 @@ from collections import deque from functools import partial +import jax import jax.numpy as jnp import pennylane as qp from jax._src.lib.mlir import ir @@ -28,8 +29,9 @@ from pennylane.pytrees import flatten from catalyst.decomposition.type_utils import ( + convert_types_to_mlir_strings, + format_for_id, get_dummy_values_for_container, - mlir_stringify_type, post_process_concretize_leaves, replace_abstract_wires_with_concrete_wires, ) @@ -77,7 +79,15 @@ def __init__(self, op: qp.core.Operator2): def parse_dynamic_shape(self) -> dict: """Return the dynamic shape as a dictionary of dtypes from the dynamic arg names.""" - return {argname: argtype for argname, argtype in sorted(self.op.dynamic_args.items())} + # breakpoint() + return { + argname: ( + argtype + if isinstance(argtype, (jax.core.ShapedArray, qp.typing.AbstractArray)) + else [argtype] + ) + for argname, argtype in sorted(self.op.dynamic_args.items()) + } def parse_wire_lens(self) -> dict: """Return the length of each of the wire args as a dictionary from the wire arg names.""" @@ -133,13 +143,7 @@ def get_dynamic_shape(self) -> dict: def get_dynamic_shape_id_format(self) -> str: """Return the dynamic shape formatted for GraphOpId.""" - return ( - "{" - + ",".join( - f"{name}:{mlir_stringify_type(shape)}" for name, shape in self.dynamic_shape.items() - ) - + "}" - ) + return format_for_id(convert_types_to_mlir_strings(self.dynamic_shape)) def get_wire_lens_id_format(self) -> str: """Return the wire lengths formatted for GraphOpId.""" @@ -360,10 +364,11 @@ def fetch_all_reachable_decomposition_rules_from_op( graph_op_id = GraphOpID(op) probe = ( graph_op_id.get_operator_name(), - { - name: mlir_stringify_type(shape) - for name, shape in graph_op_id.get_dynamic_shape().items() - }, + # { + # name: convert_types_to_mlir_strings(shape) + # for name, shape in graph_op_id.get_dynamic_shape().items() + # }, + convert_types_to_mlir_strings(graph_op_id.get_dynamic_shape()), graph_op_id.wire_lens, graph_op_id.static_data, graph_op_id.extra_data, diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index 88b910fd46..acbb698edf 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -36,26 +36,74 @@ _PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} -def _stringify_shaped_type(shape: tuple, dim: int, element_type) -> str: - """Return a string representation of the given shaped data type.""" - if dim + 1 == len(shape): - inner_content = _PY_DTYPES_TO_MLIR_DTYPES[element_type] - else: - inner_content = _stringify_shaped_type(shape, dim + 1, element_type) - length = shape[dim] - return f"[{','.join([inner_content] * length)}]" - - -def mlir_stringify_type(dtype: qp.typing.AbstractArray): - """Return a string representation of the given data type.""" - assert isinstance( - dtype, qp.typing.AbstractArray - ), f"Expected an AbstractArray to stringify, got {dtype}" - element_type = dtype.dtype.type - if dtype.shape == (): - return f"[{_PY_DTYPES_TO_MLIR_DTYPES[element_type]}]" - else: - return _stringify_shaped_type(dtype.shape, 0, element_type) +# def _stringify_shaped_type(shape: tuple, dim: int, element_type) -> str: +# """Return a string representation of the given shaped data type.""" +# if dim + 1 == len(shape): +# inner_content = _PY_DTYPES_TO_MLIR_DTYPES[element_type] +# else: +# inner_content = _stringify_shaped_type(shape, dim + 1, element_type) +# length = shape[dim] +# return f"[{','.join([inner_content] * length)}]" + + +# def mlir_stringify_type(dtype: qp.typing.AbstractArray): +# """Return a string representation of the given data type.""" +# assert isinstance( +# dtype, qp.typing.AbstractArray +# ), f"Expected an AbstractArray to stringify, got {dtype}" +# element_type = dtype.dtype.type +# if dtype.shape == (): +# return f"[{_PY_DTYPES_TO_MLIR_DTYPES[element_type]}]" +# else: +# return _stringify_shaped_type(dtype.shape, 0, element_type) + + +def convert_shaped_type_to_mlir_string(shaped_type, current_dim=0): + """Convert a shape of arbitrary dimension to a string with MLIR type strings for values.""" + if current_dim == shaped_type.ndim: + return _PY_DTYPES_TO_MLIR_DTYPES[shaped_type.dtype.type] + + return [convert_shaped_type_to_mlir_string(shaped_type, current_dim + 1)] * shaped_type.shape[ + current_dim + ] + + +def convert_types_to_mlir_strings(d: dict) -> dict: + """Convert the values of a dictionary to MLIR type strings.""" + + def handle_item(item): + if isinstance(item, type): + if item in _PY_DTYPES_TO_MLIR_DTYPES: + return _PY_DTYPES_TO_MLIR_DTYPES[item] + raise TypeError( + f"encountered unknown type {type(item)} of item {item} when converting to mlir strings." + ) + elif type(item) in _PY_DTYPES_TO_MLIR_DTYPES: + return _PY_DTYPES_TO_MLIR_DTYPES[type(item)] + elif isinstance(item, str): + return item + elif isinstance(item, (list, tuple)): + return [handle_item(i) for i in item] + elif isinstance(item, (ShapedArray, qp.typing.AbstractArray)): + return convert_shaped_type_to_mlir_string(item) + else: + raise TypeError( + f"encountered unknown type {type(item)} of item {item} when converting to mlir strings." + ) + + return {k: handle_item(v) for k, v in d.items()} + + +def format_for_id(d): + """Format a structure for ID, after calling convert_types_to_mlir_string on it.""" + + def handle_item(item): + if isinstance(item, str): + return item + elif isinstance(item, list): + return "[" + ",".join(handle_item(i) for i in item) + "]" + + return "{" + ",".join(k + ":" + handle_item(v) for k, v in d.items()) + "}" def get_dummy_values_for_container(container): From ea345f63957b741e8633b01c17b3da7c29a9b8bb Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 30 Jul 2026 16:12:44 -0400 Subject: [PATCH 32/36] fix --- .../catalyst/decomposition/decomposition_rules.py | 14 +------------- frontend/catalyst/decomposition/type_utils.py | 2 ++ 2 files changed, 3 insertions(+), 13 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index 5a874e9888..3234b90e42 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -79,15 +79,7 @@ def __init__(self, op: qp.core.Operator2): def parse_dynamic_shape(self) -> dict: """Return the dynamic shape as a dictionary of dtypes from the dynamic arg names.""" - # breakpoint() - return { - argname: ( - argtype - if isinstance(argtype, (jax.core.ShapedArray, qp.typing.AbstractArray)) - else [argtype] - ) - for argname, argtype in sorted(self.op.dynamic_args.items()) - } + return {argname: argtype for argname, argtype in sorted(self.op.dynamic_args.items())} def parse_wire_lens(self) -> dict: """Return the length of each of the wire args as a dictionary from the wire arg names.""" @@ -364,10 +356,6 @@ def fetch_all_reachable_decomposition_rules_from_op( graph_op_id = GraphOpID(op) probe = ( graph_op_id.get_operator_name(), - # { - # name: convert_types_to_mlir_strings(shape) - # for name, shape in graph_op_id.get_dynamic_shape().items() - # }, convert_types_to_mlir_strings(graph_op_id.get_dynamic_shape()), graph_op_id.wire_lens, graph_op_id.static_data, diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index acbb698edf..fe32d2dc27 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -85,6 +85,8 @@ def handle_item(item): elif isinstance(item, (list, tuple)): return [handle_item(i) for i in item] elif isinstance(item, (ShapedArray, qp.typing.AbstractArray)): + if item.shape == (): + return [_PY_DTYPES_TO_MLIR_DTYPES[item.dtype.type]] return convert_shaped_type_to_mlir_string(item) else: raise TypeError( From 795b4160ebb499e8bcb68d602287f29cc27af960 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 30 Jul 2026 16:20:58 -0400 Subject: [PATCH 33/36] FIX! --- .../catalyst/decomposition/decomposition_rules.py | 4 +--- frontend/catalyst/decomposition/type_utils.py | 11 +++++------ 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index 3234b90e42..2faf94fea5 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -224,9 +224,7 @@ def prepare_dynamic_op_kwargs(dynamic_shape, wire_lens) -> dict: kwargs = {} for wire_name, wire_len in wire_lens.items(): kwargs[wire_name] = jnp.array(range(wire_len), dtype=int) - for arg_name, arg_shape in dynamic_shape.items(): - kwargs[arg_name] = get_dummy_values_for_container(arg_shape) - return kwargs + return kwargs | get_dummy_values_for_container(dynamic_shape) def compile_decomposition_rules( diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index fe32d2dc27..254c6db368 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -108,17 +108,15 @@ def handle_item(item): return "{" + ",".join(k + ":" + handle_item(v) for k, v in d.items()) + "}" -def get_dummy_values_for_container(container): +def get_dummy_values_for_container(dictionary): """ - Given a container of python or MLIR types, replace the types with corresponding dummy values. + Given a dictionary of python or MLIR types, replace the types with corresponding dummy values. - Each item in the container must be representible as an MLIR tensor with at most one layer of + Each item in the dictionary must be representable as an MLIR tensor with at most one layer of nesting, i.e. cannot be nested and all elements must be of the same type. Ex. [[float, float], [int, int, int], [int32, int32, int32, int32]] """ - if isinstance(container, str): - return jnp.zeros((), dtype=_MLIR_DTYPES_TO_PY_DTYPES[container]) def handle_item(item): if isinstance(item, (list, tuple)): @@ -135,7 +133,8 @@ def handle_item(item): f"Unexpected type in container when creating dummy values: {type(item)}" ) - return tuple(handle_item(item) for item in container) + # return tuple(handle_item(item) for item in container) + return {k: handle_item(v) for k, v in dictionary.items()} def replace_abstract_wires_with_concrete_wires(node): From 0a836bdfa69a2d66712a52300c22debe3183a15c Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 30 Jul 2026 16:29:11 -0400 Subject: [PATCH 34/36] burn --- .../decomposition/decomposition_rules.py | 5 ++-- frontend/catalyst/decomposition/type_utils.py | 24 +------------------ 2 files changed, 3 insertions(+), 26 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index 2faf94fea5..e29fb77475 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -21,7 +21,6 @@ from collections import deque from functools import partial -import jax import jax.numpy as jnp import pennylane as qp from jax._src.lib.mlir import ir @@ -30,7 +29,7 @@ from catalyst.decomposition.type_utils import ( convert_types_to_mlir_strings, - format_for_id, + format_dynamic_params_for_id, get_dummy_values_for_container, post_process_concretize_leaves, replace_abstract_wires_with_concrete_wires, @@ -135,7 +134,7 @@ def get_dynamic_shape(self) -> dict: def get_dynamic_shape_id_format(self) -> str: """Return the dynamic shape formatted for GraphOpId.""" - return format_for_id(convert_types_to_mlir_strings(self.dynamic_shape)) + return format_dynamic_params_for_id(convert_types_to_mlir_strings(self.dynamic_shape)) def get_wire_lens_id_format(self) -> str: """Return the wire lengths formatted for GraphOpId.""" diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index 254c6db368..5c981a137a 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -36,28 +36,6 @@ _PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} -# def _stringify_shaped_type(shape: tuple, dim: int, element_type) -> str: -# """Return a string representation of the given shaped data type.""" -# if dim + 1 == len(shape): -# inner_content = _PY_DTYPES_TO_MLIR_DTYPES[element_type] -# else: -# inner_content = _stringify_shaped_type(shape, dim + 1, element_type) -# length = shape[dim] -# return f"[{','.join([inner_content] * length)}]" - - -# def mlir_stringify_type(dtype: qp.typing.AbstractArray): -# """Return a string representation of the given data type.""" -# assert isinstance( -# dtype, qp.typing.AbstractArray -# ), f"Expected an AbstractArray to stringify, got {dtype}" -# element_type = dtype.dtype.type -# if dtype.shape == (): -# return f"[{_PY_DTYPES_TO_MLIR_DTYPES[element_type]}]" -# else: -# return _stringify_shaped_type(dtype.shape, 0, element_type) - - def convert_shaped_type_to_mlir_string(shaped_type, current_dim=0): """Convert a shape of arbitrary dimension to a string with MLIR type strings for values.""" if current_dim == shaped_type.ndim: @@ -96,7 +74,7 @@ def handle_item(item): return {k: handle_item(v) for k, v in d.items()} -def format_for_id(d): +def format_dynamic_params_for_id(d): """Format a structure for ID, after calling convert_types_to_mlir_string on it.""" def handle_item(item): From 699ab5ab72bf1475f40abcc349666991a73366f4 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Thu, 30 Jul 2026 16:31:39 -0400 Subject: [PATCH 35/36] . --- frontend/catalyst/decomposition/type_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index 5c981a137a..bfd8a7726a 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -111,7 +111,6 @@ def handle_item(item): f"Unexpected type in container when creating dummy values: {type(item)}" ) - # return tuple(handle_item(item) for item in container) return {k: handle_item(v) for k, v in dictionary.items()} From e5a4acfc72e6f5d912384d48ab8b5b40bbf4529d Mon Sep 17 00:00:00 2001 From: paul0403 Date: Tue, 4 Aug 2026 15:11:44 -0400 Subject: [PATCH 36/36] LOWERING TIME RULES LETS GOOOOOOO --- .../decomposition/decomposition_rules.py | 24 +++--- frontend/catalyst/decomposition/type_utils.py | 43 +++++++++-- .../from_plxpr/qref_operator2_primitives.py | 76 +++++++++++++++++-- 3 files changed, 119 insertions(+), 24 deletions(-) diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index e29fb77475..569aa1703a 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -172,7 +172,7 @@ def find_condition(op): op.attributes["sym_name"] = ir.StringAttr.get( "__builtin_" + old_attr.value.strip('"'), context=old_attr.context ) - funcOps.append(op.detach_from_parent()) + funcOps.append(op) return ir.WalkResult.SKIP return ir.WalkResult.ADVANCE @@ -180,6 +180,11 @@ def find_condition(op): return funcOps +def get_rules_from_module_as_list(module: ir.Module) -> list[str]: + funcOps = get_rule_funcs_from_module(module) + return [str(funcOp) for funcOp in funcOps] + + def get_rules_from_module(module: ir.Module) -> str: """ Parse and modify decomposition rules from a ModuleOp. @@ -192,7 +197,6 @@ def get_rules_from_module(module: ir.Module) -> str: `__builtin_` prefix to their names. """ funcOps = get_rule_funcs_from_module(module) - return "\n".join(str(funcOp) for funcOp in funcOps) if funcOps else "" @@ -242,7 +246,6 @@ def compile_decomposition_rules( """ kwargs = prepare_dynamic_op_kwargs(dynamic_shape, wire_lens) extra_data = extra_data or {} - device = qp.device("null.qubit", wires=sum(wire_lens.values())) _, name_to_resource_ids, decomp_rules = collect_resources_for_op( op_name, kwargs | static_data | extra_data, is_custom_op @@ -265,6 +268,8 @@ def decomp_rule(*_args, **_kwargs): subroutines = [rule_to_subroutine(rule) for rule in decomp_rules] + # TODO: reconcretify abstracted hybrid ops + @qp.qjit( target="mlir", capture=True, @@ -331,13 +336,11 @@ def fetch_all_reachable_decomposition_rules_from_op( start = (op_name, dynamic_shape, wire_lens, static_data, extra_data) queue.append(start) visited = [start] - rules = [ - *get_rule_funcs_from_module( - compile_decomposition_rules( - op_name, op_id, dynamic_shape, wire_lens, static_data, extra_data=extra_data - ) + rules = get_rules_from_module_as_list( + compile_decomposition_rules( + op_name, op_id, dynamic_shape, wire_lens, static_data, extra_data=extra_data ) - ] + ) while len(queue) != 0: this_name, this_dynamic_shape, this_wire_lens, this_static_data, this_extra_data = ( @@ -361,7 +364,6 @@ def fetch_all_reachable_decomposition_rules_from_op( if not probe in visited: visited.append(probe) queue.append(probe) - module = compile_decomposition_rules( probe[0], graph_op_id.getID(), @@ -370,5 +372,5 @@ def fetch_all_reachable_decomposition_rules_from_op( probe[3], probe[4], ) - rules.extend(get_rule_funcs_from_module(module)) + rules.extend(get_rules_from_module_as_list(module)) return rules diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index bfd8a7726a..09e8d3dccd 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -18,6 +18,7 @@ import jax.numpy as jnp import pennylane as qp +from jax._src.lib.mlir import ir from jax.core import ShapedArray _MLIR_DTYPES_TO_PY_DTYPES = { @@ -33,17 +34,43 @@ "complex": jnp.complex128, } -_PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} +_PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} | { + (ir.IntegerType, 1): "i1", + (ir.IntegerType, 8): "i8", + (ir.IntegerType, 16): "i16", + (ir.IntegerType, 32): "i32", + (ir.IntegerType, 64): "i64", + ir.F16Type: "f16", + ir.F32Type: "f32", + ir.F64Type: "f64", + (ir.ComplexType, ir.F64Type): "complex", +} + + +def get_mlir_tensor_type_map_key(mlir_type): + if isinstance(mlir_type, ir.ComplexType): + return (type(mlir_type), type(mlir_type.element_type)) + if isinstance(mlir_type, ir.IntegerType): + return (type(mlir_type), mlir_type.width) + return type(mlir_type) def convert_shaped_type_to_mlir_string(shaped_type, current_dim=0): """Convert a shape of arbitrary dimension to a string with MLIR type strings for values.""" - if current_dim == shaped_type.ndim: - return _PY_DTYPES_TO_MLIR_DTYPES[shaped_type.dtype.type] + if isinstance(shaped_type, (ShapedArray, qp.typing.AbstractArray)): + if current_dim == shaped_type.ndim: + return _PY_DTYPES_TO_MLIR_DTYPES[shaped_type.dtype.type] - return [convert_shaped_type_to_mlir_string(shaped_type, current_dim + 1)] * shaped_type.shape[ - current_dim - ] + return [ + convert_shaped_type_to_mlir_string(shaped_type, current_dim + 1) + ] * shaped_type.shape[current_dim] + elif isinstance(shaped_type, ir.RankedTensorType): + if current_dim == shaped_type.rank: + return _PY_DTYPES_TO_MLIR_DTYPES[get_mlir_tensor_type_map_key(shaped_type.element_type)] + + return [ + convert_shaped_type_to_mlir_string(shaped_type, current_dim + 1) + ] * shaped_type.shape[current_dim] def convert_types_to_mlir_strings(d: dict) -> dict: @@ -66,6 +93,10 @@ def handle_item(item): if item.shape == (): return [_PY_DTYPES_TO_MLIR_DTYPES[item.dtype.type]] return convert_shaped_type_to_mlir_string(item) + elif isinstance(item, ir.RankedTensorType): + if len(item.shape) == 0: + return [_PY_DTYPES_TO_MLIR_DTYPES[get_mlir_tensor_type_map_key(item.element_type)]] + return convert_shaped_type_to_mlir_string(item) else: raise TypeError( f"encountered unknown type {type(item)} of item {item} when converting to mlir strings." diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 72feda4005..a146249515 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -15,23 +15,31 @@ of quantum operations to reference semantics JAXPR. """ -# pylint: disable=unused-argument from jax._src.lib.mlir import ir + +# pylint: disable=unused-argument +from jax.core import ShapedArray from jax.extend.core import Primitive from jax.interpreters import mlir from jaxlib.mlir._mlir_libs import _mlir as _ods_cext from jaxlib.mlir.dialects.stablehlo import ConvertOp as StableHLOConvertOp +from pennylane.core.operator.utils import abstractify from pennylane.pytrees import unflatten +from pennylane.typing import AbstractArray +from pennylane.wires import AbstractQubit # TODO: remove after jax v0.7.2 upgrade # Mock _ods_cext.globals.register_traceback_file_exclusion due to API conflicts between # Catalyst's MLIR version and the MLIR version used by JAX. The current JAX version has not # yet updated to the latest MLIR, causing compatibility issues. This workaround will be removed # once JAX updates to a compatible MLIR version -# pylint: disable=ungrouped-imports from catalyst.decomposition.decomposition_rules import ( fetch_all_reachable_decomposition_rules_from_op, ) +from catalyst.decomposition.type_utils import ( + convert_types_to_mlir_strings, + format_dynamic_params_for_id, +) from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval from catalyst.jax_extras.patches import mock_attributes from catalyst.jax_primitives import ( @@ -206,6 +214,11 @@ def _process_qubits(*args, op_cls, wire_lens, hybrid_lens) -> tuple[list, dict[s return qubits, qubit_map +@abstractify.register(ShapedArray) +def _abstractify_jax_array(val): + return AbstractArray(val.shape, val.dtype) + + def _qref_operator_p_lowering(jax_ctx: mlir.LoweringRuleContext, *args, op_cls, **kwargs): ctx = jax_ctx.module_context.context ctx.allow_unregistered_dialects = True @@ -304,20 +317,69 @@ def _qref_operator_p_lowering(jax_ctx: mlir.LoweringRuleContext, *args, op_cls, qubit_map=qubit_map, ) - # Collect decomp rules - # op_name, op_id, dynamic_shape, wire_lens, static_data, extra_data=None + # Collect decomp rules reachable from the current op + dynamic_shape = {} + for dynamic_argname, param in zip(op_cls.dynamic_argnames, params, strict=True): + dynamic_shape[dynamic_argname] = param.type + dynamic_shape = convert_types_to_mlir_strings(dynamic_shape) + repack_wire_argnames = [] for wire_argname in op_cls.wire_argnames: if wire_argname not in op_cls.hybrid_argnames: repack_wire_argnames.append(wire_argname) + repack_wire_lens = {a: b for a, b in zip(repack_wire_argnames, wire_lens, strict=True)} + + extra_data = {} + non_hybrid_wire_len = 0 + for w in repack_wire_argnames: + non_hybrid_wire_len += len(qubit_map[w]) # pylint:disable=unsubscriptable-object + hybrid_arg_start_idx = len(params) + non_hybrid_wire_len + for hybrid_argname, hybrid_len, hybrid_tree in zip( + op_cls.hybrid_argnames, hybrid_lens, hybrid_trees + ): + replaced_leaves = [] + for leaf in jax_ctx.avals_in[hybrid_arg_start_idx : hybrid_arg_start_idx + hybrid_len]: + if isinstance(leaf, AbstractQubit): + replaced_leaves.append(ShapedArray((), dtype=int)) + else: + replaced_leaves.append(leaf) + + with Patcher( + (AbstractArray, "__hash__", lambda x: id(x)), + ): + replaced_leaves = abstractify(replaced_leaves) + unflattened = unflatten(replaced_leaves, hybrid_tree) + unflattened = abstractify(unflattened) + extra_data[hybrid_argname] = unflattened + hybrid_arg_start_idx += hybrid_len + + op_id = ( + op_cls.__name__ + + format_dynamic_params_for_id(dict(sorted(dynamic_shape.items()))) + + "{" + + ",".join(f"{name}:{shape}" for name, shape in sorted(repack_wire_lens.items())) + + "}" + + "{" + + ",".join(f"{k}:{v}" for k, v in sorted(repack_static_data.items())) + + "}" + + "[" + + str(uid) + + "]" + ) decomp_rules = fetch_all_reachable_decomposition_rules_from_op( op_name=op_cls.__name__, - op_id="Bob{phi:[f64],thetas:[f64,f64]}{wires:1,other_wires:2}{bob_word:blah}", - dynamic_shape={"phi": ["f64"], "thetas": ["f64", "f64"]}, - wire_lens={a: b for a, b in zip(repack_wire_argnames, wire_lens, strict=True)}, + op_id=op_id, + dynamic_shape=dynamic_shape, + wire_lens=repack_wire_lens, static_data=repack_static_data, + extra_data=extra_data, ) + with ir.InsertionPoint(jax_ctx.module_context.module.body): + for decomp_rule in decomp_rules: + if decomp_rule: + ir.Operation.parse(decomp_rule).clone() + return []