From 6b714ec2d52f25cece4ef896b4654c484410dcfd Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Wed, 27 May 2026 15:38:20 +0200 Subject: [PATCH 01/25] Add Graph class --- .../Dialect/QCO/Transforms/Mapping/Mapping.h | 3 +- .../mlir/Dialect/QCO/Utils/Algorithms.h | 34 +++++++++++- .../QCO/Transforms/Mapping/Mapping.cpp | 37 ++++--------- mlir/lib/Dialect/QCO/Utils/Algorithms.cpp | 55 ++++++++++++++++++- .../QCO/Transforms/Mapping/test_mapping.cpp | 44 +++++++-------- 5 files changed, 117 insertions(+), 56 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h b/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h index 0eead57bd9..7b76d6affa 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h @@ -25,7 +25,6 @@ namespace mlir::qco { * @brief Create a mapping pass instance with the given target architecture. * @returns a pass object. */ -std::unique_ptr createMappingPass(size_t nqubits, const Edges& coupling, - MappingPassOptions options); +std::unique_ptr createMappingPass(const EdgeSet&, MappingPassOptions); } // namespace mlir::qco diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h b/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h index 77c57585f3..113be2ed92 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h @@ -11,15 +11,43 @@ #pragma once #include +#include #include #include #include namespace mlir::qco { +template using Vector = SmallVector; +template using Matrix = Vector>; +using EdgeSet = llvm::DenseSet>; -using Matrix = SmallVector, 0>; -using Edges = llvm::DenseSet>; +class Graph { +public: + /// Construct an empty graph. + Graph() = default; + /// Construct graph from edge set. + explicit Graph(const EdgeSet& edges); + /// Add a node to the graph. + void addNode(size_t id); + /// Add an edge to the graph. + void addEdge(size_t id, size_t neighbourId); + /// Add an edge to the graph. + void addEdge(std::pair edge); + /// Add multiple edges to the graph. + void addEdges(SmallVector> edges); + /// Return a set of edges. + [[nodiscard]] EdgeSet getEdges() const; + /// Return the edges of a node. + [[nodiscard]] ArrayRef getEdges(size_t id) const; + /// Return the number of nodes. + [[nodiscard]] size_t getNumNodes() const { return nodes_.size(); } + /// Returns the max degree of the graph. + [[nodiscard]] size_t getMaxDegree() const; + +private: + llvm::DenseMap> nodes_; +}; /** * @brief Find all shortest paths between two nodes in a graph. @@ -33,6 +61,6 @@ using Edges = llvm::DenseSet>; * @returns The distance matrix dist, where dist[i, j] is defined as the * distance between node i and j. */ -Matrix findAllShortestPaths(size_t n, const Edges& edges); +Matrix findAllShortestPaths(const Graph& graph); } // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 29f80e74ab..3460242c45 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -216,24 +216,19 @@ struct MappingPass : impl::MappingPassBase { public: AugmentedDevice() = default; - AugmentedDevice(size_t nqubits, const Edges& coupling) - : nqubits_(nqubits), dist_(findAllShortestPaths(nqubits, coupling)), - coupling_(coupling), neighbours_(nqubits) { - for (const auto& [u, v] : coupling_) { - neighbours_[u].push_back(v); - } - } + explicit AugmentedDevice(const Graph& coupling) + : dist_(findAllShortestPaths(coupling)), coupling_(coupling) {} /** * @returns the device's number of qubits. */ - [[nodiscard]] size_t nqubits() const { return nqubits_; } + [[nodiscard]] size_t nqubits() const { return coupling_.getNumNodes(); } /** * @returns true if @p u and @p v are adjacent. */ [[nodiscard]] bool areAdjacent(size_t u, size_t v) const { - return coupling_.contains(std::make_pair(u, v)); + return dist_[u][v] == 1UL; } /** @@ -251,25 +246,17 @@ struct MappingPass : impl::MappingPassBase { * @returns all neighbours of @p u. */ [[nodiscard]] ArrayRef neighboursOf(size_t u) const { - return neighbours_[u]; + return coupling_.getEdges(u); } /** * @returns the max degree (connectivity) of any qubit of the device. */ - [[nodiscard]] size_t maxDegree() const { - size_t deg = 0; - for (const auto& nbrs : neighbours_) { - deg = std::max(deg, nbrs.size()); - } - return deg; - } + [[nodiscard]] size_t maxDegree() const { return coupling_.getMaxDegree(); } private: - size_t nqubits_{}; - Matrix dist_; - Edges coupling_; - Neighbours neighbours_; + Matrix dist_; + Graph coupling_; }; struct [[nodiscard]] Trial { @@ -372,9 +359,9 @@ struct MappingPass : impl::MappingPassBase { public: MappingPass() = default; explicit MappingPass(MappingPassOptions options) : MappingPassBase(options) {} - explicit MappingPass(size_t nqubits, const Edges& coupling, + explicit MappingPass(const EdgeSet& couplingSet, MappingPassOptions options = {}) - : MappingPassBase(options), device(nqubits, coupling) {} + : MappingPassBase(options), device(Graph(couplingSet)) {} protected: void runOnOperation() override { @@ -959,9 +946,9 @@ struct MappingPass : impl::MappingPassBase { } // namespace -std::unique_ptr createMappingPass(size_t nqubits, const Edges& coupling, +std::unique_ptr createMappingPass(const EdgeSet& couplingSet, MappingPassOptions options) { - return std::make_unique(nqubits, coupling, options); + return std::make_unique(couplingSet, options); } } // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp b/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp index c8da5817f0..3e9436049a 100644 --- a/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp +++ b/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp @@ -10,6 +10,8 @@ #include "mlir/Dialect/QCO/Utils/Algorithms.h" +#include +#include #include #include @@ -18,16 +20,63 @@ #include namespace mlir::qco { -Matrix findAllShortestPaths(size_t n, const Edges& edges) { - Matrix dist(n, SmallVector(n, UINT64_MAX)); +Graph::Graph(const EdgeSet& edges) { + for (const auto& [u, v] : edges) { + if (!nodes_.contains(u)) { + nodes_[u] = Vector(); + } + nodes_[u].emplace_back(v); + } +} +void Graph::addNode(size_t id) { + const auto r = nodes_.try_emplace(id, SmallVector{}); + assert(r.second && "addNode: didn't insert node"); +} + +void Graph::addEdge(size_t id, size_t neighbourId) { + assert(nodes_.contains(id) && "addEdge: missing node id"); + nodes_[id].emplace_back(neighbourId); +} + +void Graph::addEdge(std::pair edge) { + addEdge(edge.first, edge.second); +} + +void Graph::addEdges(SmallVector> edges) { + for_each(edges, [this](const auto& edge) { addEdge(edge); }); +} +EdgeSet Graph::getEdges() const { + EdgeSet set; + for (const auto& [u, nbrs] : nodes_) { + for (const auto& v : nbrs) { + set.insert(std::make_pair(u, v)); + } + } + return set; +} + +ArrayRef Graph::getEdges(size_t id) const { return nodes_.at(id); } + +size_t Graph::getMaxDegree() const { + size_t deg = 0; + for (const auto& [u, nbrs] : nodes_) { + deg = std::max(deg, nbrs.size()); + } + return deg; +} + +Matrix findAllShortestPaths(const Graph& graph) { + const size_t n = graph.getNumNodes(); + const auto edges = graph.getEdges(); + + Matrix dist(n, Vector(n, UINT64_MAX)); for (const auto& [u, v] : edges) { dist[u][v] = 1; } for (std::size_t v = 0; v < n; ++v) { dist[v][v] = 0; } - for (std::size_t k = 0; k < n; ++k) { for (std::size_t i = 0; i < n; ++i) { for (std::size_t j = 0; j < n; ++j) { diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index fdab758e68..291dcd02ad 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -41,13 +41,11 @@ using namespace mlir; using namespace mlir::qco; -using DeviceSpec = std::pair; - /** * @returns llvm::success() if all two-qubit gates inside @p region * fulfill the given coupling constraints. llvm::failure(), otherwise. */ -static LogicalResult isExecutable(Region& region, const Edges& coupling) { +static LogicalResult isExecutable(Region& region, const EdgeSet& couplingSet) { return walkProgram(region, [&](Operation* curr, const Qubits& qubits) { if (auto op = dyn_cast(curr)) { if (isa(op)) { @@ -63,7 +61,7 @@ static LogicalResult isExecutable(Region& region, const Edges& coupling) { const auto i0 = qubits.getIndex(q0); const auto i1 = qubits.getIndex(q1); - if (!coupling.contains(std::make_pair(i0, i1))) { + if (!couplingSet.contains(std::make_pair(i0, i1))) { return WalkResult::interrupt(); } } @@ -74,20 +72,19 @@ static LogicalResult isExecutable(Region& region, const Edges& coupling) { } /** - * @returns a 9x9 square-grid device. + * @returns a 9x9 square-grid coupling set; */ -static DeviceSpec getNineQubitSquareGrid() { - const static Edges COUPLING{{0, 3}, {3, 0}, {0, 1}, {1, 0}, {1, 4}, {4, 1}, - {1, 2}, {2, 1}, {2, 5}, {5, 2}, {3, 6}, {6, 3}, - {3, 4}, {4, 3}, {4, 7}, {7, 4}, {4, 5}, {5, 4}, - {5, 8}, {8, 5}, {6, 7}, {7, 6}, {7, 8}, {8, 7}}; - return std::make_pair(9, COUPLING); +static EdgeSet getNineQubitSquareGrid() { + return EdgeSet{{0, 3}, {3, 0}, {0, 1}, {1, 0}, {1, 4}, {4, 1}, + {1, 2}, {2, 1}, {2, 5}, {5, 2}, {3, 6}, {6, 3}, + {3, 4}, {4, 3}, {4, 7}, {7, 4}, {4, 5}, {5, 4}, + {5, 8}, {8, 5}, {6, 7}, {7, 6}, {7, 8}, {8, 7}}; } namespace { class MappingPassTest : public testing::Test, - public testing::WithParamInterface { + public testing::WithParamInterface { protected: void SetUp() override { DialectRegistry registry; @@ -97,10 +94,10 @@ class MappingPassTest : public testing::Test, context->loadAllAvailableDialects(); } - static LogicalResult runPass(ModuleOp m, const DeviceSpec& device, + static LogicalResult runPass(ModuleOp m, const EdgeSet& couplingSet, const MappingPassOptions& options) { PassManager pm(m->getContext()); - pm.addPass(createMappingPass(device.first, device.second, options)); + pm.addPass(createMappingPass(couplingSet, options)); return pm.run(m); } @@ -192,12 +189,13 @@ TEST_P(MappingPassTest, NoExtractAfterInsert) { } TEST_P(MappingPassTest, TooManyQubitsForArch) { - const auto& device = GetParam(); + const auto& couplingSet = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); - int64_t nqubits = static_cast(device.first) + 1; + const int64_t nqubits = + static_cast(Graph(couplingSet).getNumNodes()) + 1; Value tensor = builder.qtensorAlloc(nqubits); SmallVector qubits(nqubits); for (int64_t i = 0; i < nqubits; ++i) { @@ -214,13 +212,13 @@ TEST_P(MappingPassTest, TooManyQubitsForArch) { builder.qtensorDealloc(tensor); auto m = builder.finalize(); - auto res = runPass(m.get(), device, MappingPassOptions{}); + auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); ASSERT_TRUE(res.failed()); } TEST_P(MappingPassTest, GHZ) { - const auto& device = GetParam(); + const auto& couplingSet = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -246,15 +244,15 @@ TEST_P(MappingPassTest, GHZ) { builder.qtensorDealloc(tensor); auto m = builder.finalize(); - auto res = runPass(m.get(), device, MappingPassOptions{}); + auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE(isExecutable(entry.getFunctionBody(), device.second).succeeded()); + EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); } TEST_P(MappingPassTest, Sabre) { - const auto& device = GetParam(); + const auto& couplingSet = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -338,11 +336,11 @@ TEST_P(MappingPassTest, Sabre) { builder.qtensorDealloc(tensor); auto m = builder.finalize(); - auto res = runPass(m.get(), device, MappingPassOptions{}); + auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE(isExecutable(entry.getFunctionBody(), device.second).succeeded()); + EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); } INSTANTIATE_TEST_SUITE_P(NineQubitSquareGrid, MappingPassTest, From 6431cfd9049974209201f033eebcd8ad08704989 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Thu, 28 May 2026 08:57:02 +0200 Subject: [PATCH 02/25] Update graph data structure --- .../Dialect/QCO/Transforms/Mapping/Mapping.h | 3 +- .../mlir/Dialect/QCO/Utils/Algorithms.h | 37 +++----- .../QCO/Transforms/Mapping/Mapping.cpp | 13 +-- mlir/lib/Dialect/QCO/Utils/Algorithms.cpp | 18 ++-- .../QCO/Transforms/Mapping/test_mapping.cpp | 87 ++++++++++++++----- 5 files changed, 99 insertions(+), 59 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h b/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h index 7b76d6affa..2f1ac23972 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h @@ -25,6 +25,7 @@ namespace mlir::qco { * @brief Create a mapping pass instance with the given target architecture. * @returns a pass object. */ -std::unique_ptr createMappingPass(const EdgeSet&, MappingPassOptions); +std::unique_ptr createMappingPass(const Graph::EdgeSet&, + MappingPassOptions); } // namespace mlir::qco diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h b/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h index 113be2ed92..f021a11d3b 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h @@ -20,47 +20,38 @@ namespace mlir::qco { template using Vector = SmallVector; template using Matrix = Vector>; -using EdgeSet = llvm::DenseSet>; class Graph { public: + using IdT = size_t; + using EdgeSet = llvm::DenseSet>; + /// Construct an empty graph. Graph() = default; /// Construct graph from edge set. explicit Graph(const EdgeSet& edges); /// Add a node to the graph. - void addNode(size_t id); + void addNode(IdT id); /// Add an edge to the graph. - void addEdge(size_t id, size_t neighbourId); + void addEdge(IdT id, IdT neighbourId); /// Add an edge to the graph. - void addEdge(std::pair edge); + void addEdge(std::pair edge); /// Add multiple edges to the graph. - void addEdges(SmallVector> edges); + void addEdges(SmallVector> edges); /// Return a set of edges. [[nodiscard]] EdgeSet getEdges() const; /// Return the edges of a node. - [[nodiscard]] ArrayRef getEdges(size_t id) const; + [[nodiscard]] ArrayRef getEdges(size_t id) const; /// Return the number of nodes. [[nodiscard]] size_t getNumNodes() const { return nodes_.size(); } - /// Returns the max degree of the graph. + /// Return the max degree of the graph. [[nodiscard]] size_t getMaxDegree() const; + /// Return the minimum distance matrix of the graph by implementing the + /// Floyd-Warshall Algorithm (https://en.wikipedia.org/wiki/Floyd–Warshall_algorithm) + /// where dist[i][j] denotes the distance between i and j. + [[nodiscard]] Matrix getDistMatrix() const; private: - llvm::DenseMap> nodes_; + llvm::DenseMap> nodes_; }; - -/** - * @brief Find all shortest paths between two nodes in a graph. - * @details Has a time complexity of O(n^3). - * - * @link Adapted from https://en.wikipedia.org/wiki/Floyd–Warshall_algorithm - * - * @param n The number of nodes in the graph. - * @param edges The set of edges (i, j). - * - * @returns The distance matrix dist, where dist[i, j] is defined as the - * distance between node i and j. - */ -Matrix findAllShortestPaths(const Graph& graph); - } // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 3460242c45..005225d8ee 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -216,8 +216,8 @@ struct MappingPass : impl::MappingPassBase { public: AugmentedDevice() = default; - explicit AugmentedDevice(const Graph& coupling) - : dist_(findAllShortestPaths(coupling)), coupling_(coupling) {} + explicit AugmentedDevice(const Graph::EdgeSet& couplingSet) + : coupling_(couplingSet), dist_(coupling_.getDistMatrix()) {} /** * @returns the device's number of qubits. @@ -255,8 +255,8 @@ struct MappingPass : impl::MappingPassBase { [[nodiscard]] size_t maxDegree() const { return coupling_.getMaxDegree(); } private: - Matrix dist_; Graph coupling_; + Matrix dist_; }; struct [[nodiscard]] Trial { @@ -359,9 +359,10 @@ struct MappingPass : impl::MappingPassBase { public: MappingPass() = default; explicit MappingPass(MappingPassOptions options) : MappingPassBase(options) {} - explicit MappingPass(const EdgeSet& couplingSet, + /// @brief Construct Mapping Pass from coupling-set. + explicit MappingPass(const Graph::EdgeSet& couplingSet, MappingPassOptions options = {}) - : MappingPassBase(options), device(Graph(couplingSet)) {} + : MappingPassBase(options), device(couplingSet) {} protected: void runOnOperation() override { @@ -946,7 +947,7 @@ struct MappingPass : impl::MappingPassBase { } // namespace -std::unique_ptr createMappingPass(const EdgeSet& couplingSet, +std::unique_ptr createMappingPass(const Graph::EdgeSet& couplingSet, MappingPassOptions options) { return std::make_unique(couplingSet, options); } diff --git a/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp b/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp index 3e9436049a..289553f7a1 100644 --- a/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp +++ b/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp @@ -46,7 +46,7 @@ void Graph::addEdges(SmallVector> edges) { for_each(edges, [this](const auto& edge) { addEdge(edge); }); } -EdgeSet Graph::getEdges() const { +Graph::EdgeSet Graph::getEdges() const { EdgeSet set; for (const auto& [u, nbrs] : nodes_) { for (const auto& v : nbrs) { @@ -66,25 +66,25 @@ size_t Graph::getMaxDegree() const { return deg; } -Matrix findAllShortestPaths(const Graph& graph) { - const size_t n = graph.getNumNodes(); - const auto edges = graph.getEdges(); +Matrix Graph::getDistMatrix() const { + const size_t n = getNumNodes(); + const auto edges = getEdges(); Matrix dist(n, Vector(n, UINT64_MAX)); for (const auto& [u, v] : edges) { dist[u][v] = 1; } - for (std::size_t v = 0; v < n; ++v) { + for (size_t v = 0; v < n; ++v) { dist[v][v] = 0; } - for (std::size_t k = 0; k < n; ++k) { - for (std::size_t i = 0; i < n; ++i) { - for (std::size_t j = 0; j < n; ++j) { + for (size_t k = 0; k < n; ++k) { + for (size_t i = 0; i < n; ++i) { + for (size_t j = 0; j < n; ++j) { if (dist[i][k] == UINT64_MAX || dist[k][j] == UINT64_MAX) { continue; // Avoid overflow with "infinite" distances. } - const std::size_t sum = dist[i][k] + dist[k][j]; + const size_t sum = dist[i][k] + dist[k][j]; dist[i][j] = std::min(dist[i][j], sum); } } diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 291dcd02ad..085f4068e8 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -30,6 +31,7 @@ #include #include #include +#include #include #include @@ -45,7 +47,8 @@ using namespace mlir::qco; * @returns llvm::success() if all two-qubit gates inside @p region * fulfill the given coupling constraints. llvm::failure(), otherwise. */ -static LogicalResult isExecutable(Region& region, const EdgeSet& couplingSet) { +static LogicalResult isExecutable(Region& region, + const Graph::EdgeSet& couplingSet) { return walkProgram(region, [&](Operation* curr, const Qubits& qubits) { if (auto op = dyn_cast(curr)) { if (isa(op)) { @@ -72,29 +75,30 @@ static LogicalResult isExecutable(Region& region, const EdgeSet& couplingSet) { } /** - * @returns a 9x9 square-grid coupling set; + * @returns a 9x9 square-grid coupling set. */ -static EdgeSet getNineQubitSquareGrid() { - return EdgeSet{{0, 3}, {3, 0}, {0, 1}, {1, 0}, {1, 4}, {4, 1}, - {1, 2}, {2, 1}, {2, 5}, {5, 2}, {3, 6}, {6, 3}, - {3, 4}, {4, 3}, {4, 7}, {7, 4}, {4, 5}, {5, 4}, - {5, 8}, {8, 5}, {6, 7}, {7, 6}, {7, 8}, {8, 7}}; +static Graph::EdgeSet getNineQubitSquareGrid() { + return Graph::EdgeSet{{0, 3}, {3, 0}, {0, 1}, {1, 0}, {1, 4}, {4, 1}, + {1, 2}, {2, 1}, {2, 5}, {5, 2}, {3, 6}, {6, 3}, + {3, 4}, {4, 3}, {4, 7}, {7, 4}, {4, 5}, {5, 4}, + {5, 8}, {8, 5}, {6, 7}, {7, 6}, {7, 8}, {8, 7}}; } namespace { class MappingPassTest : public testing::Test, - public testing::WithParamInterface { + public testing::WithParamInterface { protected: void SetUp() override { DialectRegistry registry; - registry.insert(); + registry.insert(); context = std::make_unique(); context->appendDialectRegistry(registry); context->loadAllAvailableDialects(); } - static LogicalResult runPass(ModuleOp m, const EdgeSet& couplingSet, + static LogicalResult runPass(ModuleOp m, const Graph::EdgeSet& couplingSet, const MappingPassOptions& options) { PassManager pm(m->getContext()); pm.addPass(createMappingPass(couplingSet, options)); @@ -107,17 +111,15 @@ class MappingPassTest : public testing::Test, }; // namespace TEST_P(MappingPassTest, NoEntryPoint) { - const auto& device = GetParam(); + const auto& couplingSet = GetParam(); OwningOpRef m = ModuleOp::create(UnknownLoc::get(context.get())); - - auto res = runPass(m.get(), device, MappingPassOptions{}); - + auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); ASSERT_TRUE(res.failed()); } TEST_P(MappingPassTest, NoQubitAllocations) { - const auto& device = GetParam(); + const auto& couplingSet = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -127,13 +129,13 @@ TEST_P(MappingPassTest, NoQubitAllocations) { builder.sink(q0); auto m = builder.finalize(); - auto res = runPass(m.get(), device, MappingPassOptions{}); + auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); ASSERT_TRUE(res.failed()); } TEST_P(MappingPassTest, NoTwoTensors) { - const auto& device = GetParam(); + const auto& couplingSet = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -158,13 +160,13 @@ TEST_P(MappingPassTest, NoTwoTensors) { builder.qtensorDealloc(tensor1); auto m = builder.finalize(); - auto res = runPass(m.get(), device, MappingPassOptions{}); + auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); ASSERT_TRUE(res.failed()); } TEST_P(MappingPassTest, NoExtractAfterInsert) { - const auto& device = GetParam(); + const auto& couplingSet = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -183,7 +185,7 @@ TEST_P(MappingPassTest, NoExtractAfterInsert) { builder.qtensorDealloc(tensor0); auto m = builder.finalize(); - auto res = runPass(m.get(), device, MappingPassOptions{}); + auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); ASSERT_TRUE(res.failed()); } @@ -251,6 +253,51 @@ TEST_P(MappingPassTest, GHZ) { EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); } +TEST_P(MappingPassTest, GHZUnrolled) { + const auto& couplingSet = GetParam(); + + PassManager pm(context.get()); + pm.addNestedPass(createQuantumLoopUnroll()); + pm.addPass(createCSEPass()); + pm.addPass(createCanonicalizerPass()); + // pm.addPass(createMappingPass(couplingSet, MappingPassOptions{})); + + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + Value tensor = builder.qtensorAlloc(3); + Value q0; + std::tie(tensor, q0) = builder.qtensorExtract(tensor, 0); + q0 = builder.h(q0); + tensor = builder.qtensorInsert(q0, tensor, 0); + tensor = builder.scfFor( + 1, 3, 1, {tensor}, [&builder](Value iv, ValueRange iterArgs) { + Value loopTensor = iterArgs[0]; + Value ctrl; + Value targ; + + std::tie(loopTensor, ctrl) = builder.qtensorExtract(loopTensor, 0); + std::tie(loopTensor, targ) = builder.qtensorExtract(loopTensor, iv); + + std::tie(ctrl, targ) = builder.cx(ctrl, targ); + + loopTensor = builder.qtensorInsert(ctrl, loopTensor, 0); + loopTensor = builder.qtensorInsert(targ, loopTensor, iv); + + return SmallVector{loopTensor}; + })[0]; + builder.qtensorDealloc(tensor); + + auto m = builder.finalize(); + auto res = pm.run(m.get()); + auto entry = getEntryPoint(m.get()); + + m->dump(); + + ASSERT_TRUE(res.succeeded()); + EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); +} + TEST_P(MappingPassTest, Sabre) { const auto& couplingSet = GetParam(); From a2982d133376fd13b65fe529fc0457bea43007d8 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Thu, 28 May 2026 09:04:05 +0200 Subject: [PATCH 03/25] Add GHZUnrolled test --- .../QCO/Transforms/Mapping/test_mapping.cpp | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 085f4068e8..721cb5dd6f 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -192,22 +192,22 @@ TEST_P(MappingPassTest, NoExtractAfterInsert) { TEST_P(MappingPassTest, TooManyQubitsForArch) { const auto& couplingSet = GetParam(); + const size_t nqubits = Graph(couplingSet).getNumNodes(); + const auto n = static_cast(nqubits) + 1; QCOProgramBuilder builder(context.get()); builder.initialize(); - const int64_t nqubits = - static_cast(Graph(couplingSet).getNumNodes()) + 1; - Value tensor = builder.qtensorAlloc(nqubits); - SmallVector qubits(nqubits); - for (int64_t i = 0; i < nqubits; ++i) { + Value tensor = builder.qtensorAlloc(n); + SmallVector qubits(n); + for (int64_t i = 0; i < n; ++i) { Value qi; std::tie(tensor, qi) = builder.qtensorExtract(tensor, i); qi = builder.h(qi); qubits[i] = qi; } - for (int64_t i = 0; i < nqubits; ++i) { + for (int64_t i = 0; i < n; ++i) { tensor = builder.qtensorInsert(qubits[i], tensor, i); } @@ -255,23 +255,25 @@ TEST_P(MappingPassTest, GHZ) { TEST_P(MappingPassTest, GHZUnrolled) { const auto& couplingSet = GetParam(); + const size_t nqubits = Graph(couplingSet).getNumNodes(); + const auto n = static_cast(nqubits); PassManager pm(context.get()); pm.addNestedPass(createQuantumLoopUnroll()); pm.addPass(createCSEPass()); pm.addPass(createCanonicalizerPass()); - // pm.addPass(createMappingPass(couplingSet, MappingPassOptions{})); + pm.addPass(createMappingPass(couplingSet, MappingPassOptions{})); QCOProgramBuilder builder(context.get()); builder.initialize(); - Value tensor = builder.qtensorAlloc(3); + Value tensor = builder.qtensorAlloc(n); Value q0; std::tie(tensor, q0) = builder.qtensorExtract(tensor, 0); q0 = builder.h(q0); tensor = builder.qtensorInsert(q0, tensor, 0); tensor = builder.scfFor( - 1, 3, 1, {tensor}, [&builder](Value iv, ValueRange iterArgs) { + 1, n, 1, {tensor}, [&builder](Value iv, ValueRange iterArgs) { Value loopTensor = iterArgs[0]; Value ctrl; Value targ; @@ -292,8 +294,6 @@ TEST_P(MappingPassTest, GHZUnrolled) { auto res = pm.run(m.get()); auto entry = getEntryPoint(m.get()); - m->dump(); - ASSERT_TRUE(res.succeeded()); EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); } From 2d2438d2ece2850e427e4f786af5f08d7d40d771 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 29 May 2026 08:42:06 +0200 Subject: [PATCH 04/25] Add findCycle method --- .../mlir/Dialect/QCO/Utils/Algorithms.h | 14 +-- mlir/lib/Dialect/QCO/Utils/Algorithms.cpp | 90 ++++++++++++++++--- 2 files changed, 89 insertions(+), 15 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h b/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h index f021a11d3b..f87fce50de 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h @@ -15,6 +15,7 @@ #include #include +#include #include namespace mlir::qco { @@ -43,15 +44,18 @@ class Graph { /// Return the edges of a node. [[nodiscard]] ArrayRef getEdges(size_t id) const; /// Return the number of nodes. - [[nodiscard]] size_t getNumNodes() const { return nodes_.size(); } + [[nodiscard]] size_t getNumNodes() const { return adj_.size(); } /// Return the max degree of the graph. [[nodiscard]] size_t getMaxDegree() const; - /// Return the minimum distance matrix of the graph by implementing the - /// Floyd-Warshall Algorithm (https://en.wikipedia.org/wiki/Floyd–Warshall_algorithm) - /// where dist[i][j] denotes the distance between i and j. + /// Return the minimum distance matrix of the graph by implementing the + /// Floyd-Warshall Algorithm + /// (https://en.wikipedia.org/wiki/Floyd–Warshall_algorithm) where dist[i][j] + /// denotes the distance between i and j. [[nodiscard]] Matrix getDistMatrix() const; + /// Return cycle in graph or std::nullopt if none exists. + [[nodiscard]] std::optional> findCycle() const; private: - llvm::DenseMap> nodes_; + llvm::DenseMap> adj_; }; } // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp b/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp index 289553f7a1..7b81c5e1e0 100644 --- a/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp +++ b/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp @@ -18,24 +18,25 @@ #include #include #include +#include namespace mlir::qco { Graph::Graph(const EdgeSet& edges) { for (const auto& [u, v] : edges) { - if (!nodes_.contains(u)) { - nodes_[u] = Vector(); + if (!adj_.contains(u)) { + adj_[u] = Vector(); } - nodes_[u].emplace_back(v); + adj_[u].emplace_back(v); } } void Graph::addNode(size_t id) { - const auto r = nodes_.try_emplace(id, SmallVector{}); - assert(r.second && "addNode: didn't insert node"); + const auto r = adj_.try_emplace(id, SmallVector{}); + assert(r.second && "expected to insert node"); } void Graph::addEdge(size_t id, size_t neighbourId) { - assert(nodes_.contains(id) && "addEdge: missing node id"); - nodes_[id].emplace_back(neighbourId); + assert(adj_.contains(id) && "addEdge: missing node id"); + adj_[id].emplace_back(neighbourId); } void Graph::addEdge(std::pair edge) { @@ -48,7 +49,7 @@ void Graph::addEdges(SmallVector> edges) { Graph::EdgeSet Graph::getEdges() const { EdgeSet set; - for (const auto& [u, nbrs] : nodes_) { + for (const auto& [u, nbrs] : adj_) { for (const auto& v : nbrs) { set.insert(std::make_pair(u, v)); } @@ -56,11 +57,11 @@ Graph::EdgeSet Graph::getEdges() const { return set; } -ArrayRef Graph::getEdges(size_t id) const { return nodes_.at(id); } +ArrayRef Graph::getEdges(size_t id) const { return adj_.at(id); } size_t Graph::getMaxDegree() const { size_t deg = 0; - for (const auto& [u, nbrs] : nodes_) { + for (const auto& [u, nbrs] : adj_) { deg = std::max(deg, nbrs.size()); } return deg; @@ -92,4 +93,73 @@ Matrix Graph::getDistMatrix() const { return dist; } + +[[nodiscard]] std::optional> Graph::findCycle() const { + enum struct State : uint8_t { Unseen, Seen, Finished }; + + struct Frame { + IdT id; + size_t neighbourIdx; + }; + + SmallVector stack; + llvm::DenseMap parents; + llvm::DenseMap states; + + // Preparation step: Mark all nodes as unseen. + for_each(adj_.keys(), [&](IdT id) { states[id] = State::Unseen; }); + + for (const auto initId : adj_.keys()) { + // Only start from unseen nodes. + if (states[initId] != State::Unseen) { + continue; + } + + stack.emplace_back(initId, 0); + + while (!stack.empty()) { + Frame& top = stack.back(); + + // If we haven't seen this node before, mark it as seen. + if (states[top.id] == State::Unseen) { + states[top.id] = State::Seen; + } + + auto it = adj_.find(top.id); + assert(it != adj_.end() && "expected node id in adjacency map"); + const auto nbrs = it->getSecond(); + + // Once all neighbours have been visited (indicated by the index exceeding + // the number of neighbours - 1), set the frame on node to finished and + // pop it from the stack. + if (top.neighbourIdx >= nbrs.size()) { + states[top.id] = State::Finished; + stack.pop_back(); + continue; + } + + // Collect the neighbour and advance the index on the + // frame for the next iteration. + const auto nbrId = nbrs[top.neighbourIdx]; + ++top.neighbourIdx; + + if (states[nbrId] == State::Unseen) { + parents[nbrId] = top.id; + stack.emplace_back(nbrId, 0); + } else if (states[nbrId] == State::Seen) { + SmallVector path; + for (auto curr = top.id; curr != nbrId; curr = parents[curr]) { + path.emplace_back(curr); + } + path.emplace_back(nbrId); + return path; + } + } + + // Preparse stack for next iteration. + stack.clear(); + } + + return std::nullopt; +} } // namespace mlir::qco From f1261ced4dd5db6f01efeb4ae7dc42e9c3ce192d Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 29 May 2026 08:47:07 +0200 Subject: [PATCH 05/25] Rename Algorithms.{h,cpp} to Graph.{h, cpp} --- mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h | 2 +- mlir/include/mlir/Dialect/QCO/Utils/{Algorithms.h => Graph.h} | 4 +++- mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp | 2 +- mlir/lib/Dialect/QCO/Utils/{Algorithms.cpp => Graph.cpp} | 2 +- 4 files changed, 6 insertions(+), 4 deletions(-) rename mlir/include/mlir/Dialect/QCO/Utils/{Algorithms.h => Graph.h} (92%) rename mlir/lib/Dialect/QCO/Utils/{Algorithms.cpp => Graph.cpp} (98%) diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h b/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h index 2f1ac23972..96e4db0333 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h @@ -11,7 +11,7 @@ #pragma once #include "mlir/Dialect/QCO/Transforms/Passes.h" -#include "mlir/Dialect/QCO/Utils/Algorithms.h" +#include "mlir/Dialect/QCO/Utils/Graph.h" #include #include diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h b/mlir/include/mlir/Dialect/QCO/Utils/Graph.h similarity index 92% rename from mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h rename to mlir/include/mlir/Dialect/QCO/Utils/Graph.h index f87fce50de..2e78fd7ecf 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Algorithms.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Graph.h @@ -52,7 +52,9 @@ class Graph { /// (https://en.wikipedia.org/wiki/Floyd–Warshall_algorithm) where dist[i][j] /// denotes the distance between i and j. [[nodiscard]] Matrix getDistMatrix() const; - /// Return cycle in graph or std::nullopt if none exists. + /// Return reverse cycle in graph or `std::nullopt` if none exists. + /// Implements an iterative depth-first search inspired by LLVM's SCC + /// utilities. [[nodiscard]] std::optional> findCycle() const; private: diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 005225d8ee..3c39f3ecde 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -12,7 +12,7 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" -#include "mlir/Dialect/QCO/Utils/Algorithms.h" +#include "mlir/Dialect/QCO/Utils/Graph.h" #include "mlir/Dialect/QCO/Utils/Drivers.h" #include "mlir/Dialect/QCO/Utils/WireIterator.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" diff --git a/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp b/mlir/lib/Dialect/QCO/Utils/Graph.cpp similarity index 98% rename from mlir/lib/Dialect/QCO/Utils/Algorithms.cpp rename to mlir/lib/Dialect/QCO/Utils/Graph.cpp index 7b81c5e1e0..d25e16324f 100644 --- a/mlir/lib/Dialect/QCO/Utils/Algorithms.cpp +++ b/mlir/lib/Dialect/QCO/Utils/Graph.cpp @@ -8,7 +8,7 @@ * Licensed under the MIT License */ -#include "mlir/Dialect/QCO/Utils/Algorithms.h" +#include "mlir/Dialect/QCO/Utils/Graph.h" #include #include From a539116feeb221b97131295f610d462db608886e Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 29 May 2026 08:49:03 +0200 Subject: [PATCH 06/25] Add GroverLike unit test --- .../QCO/Transforms/Mapping/test_mapping.cpp | 89 ++++++++++++++++++- 1 file changed, 88 insertions(+), 1 deletion(-) diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 721cb5dd6f..5df33c5924 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -14,7 +14,7 @@ #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/Transforms/Mapping/Mapping.h" #include "mlir/Dialect/QCO/Transforms/Passes.h" -#include "mlir/Dialect/QCO/Utils/Algorithms.h" +#include "mlir/Dialect/QCO/Utils/Graph.h" #include "mlir/Dialect/QCO/Utils/Drivers.h" #include "mlir/Dialect/QCO/Utils/Qubits.h" @@ -298,6 +298,93 @@ TEST_P(MappingPassTest, GHZUnrolled) { EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); } +TEST_P(MappingPassTest, GroverLike) { + const auto& couplingSet = GetParam(); + + PassManager pm(context.get()); + pm.addPass(createMappingPass(couplingSet, MappingPassOptions{})); + + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + Value tensor = builder.qtensorAlloc(4); + Value flagTensor = builder.qtensorAlloc(1); + Value q0; + Value q1; + Value q2; + Value q3; + Value flag; + + std::tie(tensor, q0) = builder.qtensorExtract(tensor, 0); + std::tie(tensor, q1) = builder.qtensorExtract(tensor, 1); + std::tie(tensor, q2) = builder.qtensorExtract(tensor, 2); + std::tie(tensor, q3) = builder.qtensorExtract(tensor, 3); + std::tie(flagTensor, flag) = builder.qtensorExtract(flagTensor, 0); + + q0 = builder.h(q0); + q1 = builder.h(q1); + q2 = builder.h(q2); + q3 = builder.h(q3); + flag = builder.x(flag); + + const auto forResults = builder.scfFor( + 1, 3, 1, {q0, q1, q2, q3, flag}, [&builder](Value, ValueRange iterArgs) { + Value iterQ0 = iterArgs[0]; + Value iterQ1 = iterArgs[1]; + Value iterQ2 = iterArgs[2]; + Value iterQ3 = iterArgs[3]; + Value iterFlag = iterArgs[4]; + + std::tie(iterQ0, iterQ2) = builder.cx(iterQ0, iterQ2); + std::tie(iterQ2, iterQ3) = builder.cx(iterQ2, iterQ3); + std::tie(iterQ3, iterQ0) = builder.cx(iterQ3, iterQ0); + std::tie(iterQ0, iterFlag) = builder.cx(iterQ0, iterFlag); + + return SmallVector{iterQ0, iterQ1, iterQ2, iterQ3, iterFlag}; + }); + + q0 = forResults[0]; + q1 = forResults[1]; + q2 = forResults[2]; + q3 = forResults[3]; + flag = forResults[4]; + + const auto barrierResults = builder.barrier({q0, q1, q2, q3, flag}); + q0 = barrierResults[0]; + q1 = barrierResults[1]; + q2 = barrierResults[2]; + q3 = barrierResults[3]; + flag = barrierResults[4]; + + Value c0; + Value c1; + Value c2; + Value c3; + Value c4; + + std::tie(q0, c0) = builder.measure(q0); + std::tie(q1, c1) = builder.measure(q1); + std::tie(q2, c2) = builder.measure(q2); + std::tie(q3, c3) = builder.measure(q3); + std::tie(flag, c4) = builder.measure(flag); + + tensor = builder.qtensorInsert(q0, tensor, 1); + tensor = builder.qtensorInsert(q1, tensor, 2); + tensor = builder.qtensorInsert(q2, tensor, 3); + tensor = builder.qtensorInsert(q3, tensor, 4); + flagTensor = builder.qtensorInsert(flag, flagTensor, 0); + + builder.qtensorDealloc(tensor); + builder.qtensorDealloc(flagTensor); + + auto m = builder.finalize(); + auto res = pm.run(m.get()); + auto entry = getEntryPoint(m.get()); + + ASSERT_TRUE(res.succeeded()); + EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); +} + TEST_P(MappingPassTest, Sabre) { const auto& couplingSet = GetParam(); From 6c555121c87e325d21763d3cd46017b47ba926b5 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 29 May 2026 09:18:59 +0200 Subject: [PATCH 07/25] Add scf.for and qco.if support to wireiterator --- .../mlir/Dialect/QCO/Utils/WireIterator.h | 5 +- mlir/lib/Dialect/QCO/Utils/WireIterator.cpp | 49 +++++++++++++++++- .../Dialect/QCO/Utils/test_wireiterator.cpp | 51 ++++++++++++++++--- 3 files changed, 96 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h b/mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h index 70ba4b26ff..b6fc522fb2 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h @@ -32,9 +32,9 @@ class [[nodiscard]] WireIterator { using difference_type = std::ptrdiff_t; using value_type = Operation*; - WireIterator() : op_(nullptr), qubit_(nullptr), isSentinel_(false) {} + WireIterator() : op_(nullptr), qubit_(nullptr), isFinal_(false), isSentinel_(false) {} explicit WireIterator(Value qubit) - : op_(qubit.getDefiningOp()), qubit_(qubit), isSentinel_(false) {} + : op_(qubit.getDefiningOp()), qubit_(qubit), isFinal_(false), isSentinel_(false) {} /// @returns the operation the iterator points to. [[nodiscard]] Operation* operation() const { return op_; } @@ -85,6 +85,7 @@ class [[nodiscard]] WireIterator { Operation* op_; Value qubit_; + bool isFinal_; bool isSentinel_; }; diff --git a/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp b/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp index 45522cd8bf..031c13bf86 100644 --- a/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp +++ b/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp @@ -16,6 +16,7 @@ #include #include +#include #include #include #include @@ -38,13 +39,19 @@ void WireIterator::forward() { return; } + // After the final operation comes the sentinel. + if (isFinal_) { + isSentinel_ = true; + return; + } + // Find the user-operation of the qubit SSA value. assert(qubit_.hasOneUse() && "expected linear typing"); op_ = *(qubit_.user_begin()); // A sink/insert defines the end of the qubit wire (dynamic and static). if (isa(op_)) { - isSentinel_ = true; + isFinal_ = true; return; } @@ -56,6 +63,15 @@ void WireIterator::forward() { }) .Case([&](MeasureOp op) { qubit_ = op.getQubitOut(); }) .Case([&](ResetOp op) { qubit_ = op.getQubitOut(); }) + .Case([&](scf::ForOp op) { + qubit_ = op.getTiedLoopResult(&*(qubit_.use_begin())); + }) + .Case([&](qco::IfOp op) { + auto it = llvm::find(op.getQubits(), qubit_); + assert(it != op.getQubits().end()); + const auto idx = std::distance(op.getQubits().begin(), it); + qubit_ = op.getResults()[idx]; + }) .Default([&](Operation* op) { llvm::reportFatalInternalError("unknown op in def-use chain: " + op->getName().getStringRef()); @@ -67,6 +83,13 @@ void WireIterator::backward() { // If the iterator is a sentinel, reactivate the iterator. if (isSentinel_) { isSentinel_ = false; + isFinal_ = true; + return; + } + + // If the op is a nullptr, the qubit value is a block argument and thus the + // beginning of the qubit wire. + if (op_ == nullptr) { return; } @@ -74,6 +97,7 @@ void WireIterator::backward() { // the def-op. if (isa(op_)) { op_ = qubit_.getDefiningOp(); + isFinal_ = false; return; } @@ -89,6 +113,28 @@ void WireIterator::backward() { [&](UnitaryOpInterface op) { qubit_ = op.getInputForOutput(qubit_); }) .Case([&](MeasureOp op) { qubit_ = op.getQubitIn(); }) .Case([&](ResetOp op) { qubit_ = op.getQubitIn(); }) + .Case([&](scf::ForOp op) { + if (auto res = dyn_cast(qubit_)) { + OpOperand* operand = op.getTiedLoopInit(res); + qubit_ = operand->get(); + return; + } + + llvm::reportFatalInternalError( + "expected scf.for result for tied init lookup"); + }) + .Case([&](qco::IfOp op) { + if (auto res = dyn_cast(qubit_)) { + auto it = llvm::find(op.getResults(), res); + assert(it != op->result_end()); + const auto idx = std::distance(op.result_begin(), it); + qubit_ = op.getQubits()[idx]; + return; + } + + llvm::reportFatalInternalError( + "expected scf.for result for tied init lookup"); + }) .Default([&](Operation* op) { llvm::reportFatalInternalError("unknown op in def-use chain: " + op->getName().getStringRef()); @@ -98,6 +144,7 @@ void WireIterator::backward() { // If the current qubit SSA value is a BlockArgument (no defining op), the // operation will be a nullptr. op_ = qubit_.getDefiningOp(); + isFinal_ = false; } static_assert(std::bidirectional_iterator); diff --git a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp index 7eb59db509..891f03cf39 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include @@ -29,7 +30,8 @@ class WireIteratorTest : public testing::TestWithParam { protected: void SetUp() override { DialectRegistry registry; - registry.insert(); + registry.insert(); context = std::make_unique(); context->appendDialectRegistry(registry); @@ -40,20 +42,37 @@ class WireIteratorTest : public testing::TestWithParam { }; } // namespace -TEST_P(WireIteratorTest, MixedUse) { +TEST_P(WireIteratorTest, Traversal) { const bool isDynamic = GetParam(); // Build circuit. qco::QCOProgramBuilder builder(context.get()); builder.initialize(); + const auto q00 = isDynamic ? builder.allocQubit() : builder.staticQubit(0); const auto q10 = isDynamic ? builder.allocQubit() : builder.staticQubit(1); const auto q01 = builder.h(q00); const auto [q02, q11] = builder.cx(q01, q10); const auto [q03, c0] = builder.measure(q02); const auto q04 = builder.reset(q03); - builder.sink(q04); - builder.sink(q11); + const auto loopOut = builder.scfFor( + 1, 4, 1, {q04, q11}, [&builder](Value, ValueRange iterArgs) { + const auto iterQ00 = iterArgs[0]; + const auto iterQ10 = iterArgs[1]; + const auto iterQ01 = builder.h(iterQ00); + const auto [iterQ02, iterQ11] = builder.cx(iterQ01, iterQ10); + return SmallVector{iterQ02, iterQ11}; + }); + const auto q05 = loopOut[0]; + const auto q12 = loopOut[1]; + const auto ifOut = builder.qcoIf( + true, {q05, q12}, + [&](ValueRange args) { return SmallVector{args[0], args[1]}; }, + [&](ValueRange args) { return SmallVector{args[0], args[1]}; }); + const auto q06 = ifOut[0]; + const auto q13 = ifOut[1]; + builder.sink(q06); + builder.sink(q13); [[maybe_unused]] auto module = builder.finalize(); // Setup WireIterator. @@ -83,7 +102,15 @@ TEST_P(WireIteratorTest, MixedUse) { ASSERT_EQ(it.qubit(), q04); ++it; - ASSERT_EQ(it.operation(), *(q04.getUsers().begin())); // qco.sink + ASSERT_EQ(it.operation(), q05.getDefiningOp()); // scf.for + ASSERT_EQ(it.qubit(), q05); + + ++it; + ASSERT_EQ(it.operation(), q06.getDefiningOp()); // qco.if + ASSERT_EQ(it.qubit(), q06); + + ++it; + ASSERT_EQ(it.operation(), *(q06.getUsers().begin())); // qco.sink ASSERT_EQ(it.qubit(), nullptr); ++it; @@ -97,9 +124,17 @@ TEST_P(WireIteratorTest, MixedUse) { // --it; - ASSERT_EQ(it.operation(), *(q04.getUsers().begin())); // qco.sink + ASSERT_EQ(it.operation(), *(q06.getUsers().begin())); // qco.sink ASSERT_EQ(it.qubit(), nullptr); + --it; + ASSERT_EQ(it.operation(), q06.getDefiningOp()); // qco.if + ASSERT_EQ(it.qubit(), q06); + + --it; + ASSERT_EQ(it.operation(), q05.getDefiningOp()); // scf.for + ASSERT_EQ(it.qubit(), q05); + --it; ASSERT_EQ(it.operation(), q04.getDefiningOp()); // qco.reset ASSERT_EQ(it.qubit(), q04); @@ -123,6 +158,10 @@ TEST_P(WireIteratorTest, MixedUse) { --it; ASSERT_EQ(it.operation(), q00.getDefiningOp()); // qco.alloc or qco.static ASSERT_EQ(it.qubit(), q00); + + // + // Test: Recursive use with block-argument. + // } INSTANTIATE_TEST_SUITE_P(DynamicAndStatic, WireIteratorTest, ::testing::Bool(), From db099031d6bf46aadbed7653b1b0d53fbfa9c165 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 29 May 2026 12:31:14 +0200 Subject: [PATCH 08/25] Fix recursive wire iterator behavior --- mlir/lib/Dialect/QCO/Utils/WireIterator.cpp | 14 ++--- .../Dialect/QCO/Utils/test_wireiterator.cpp | 59 +++++++++++++++++-- 2 files changed, 60 insertions(+), 13 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp b/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp index 031c13bf86..0968039957 100644 --- a/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp +++ b/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp @@ -26,8 +26,9 @@ namespace mlir::qco { Value WireIterator::qubit() const { - // A sink/deallocation/insert doesn't have an OpResult. - if (op_ != nullptr && (isa(op_))) { + // The following operations don't have an OpResult. + if (op_ != nullptr && + (isa(op_))) { return nullptr; } return qubit_; @@ -49,8 +50,8 @@ void WireIterator::forward() { assert(qubit_.hasOneUse() && "expected linear typing"); op_ = *(qubit_.user_begin()); - // A sink/insert defines the end of the qubit wire (dynamic and static). - if (isa(op_)) { + // The following operations define the end of the qubit wire. + if (isa(op_)) { isFinal_ = true; return; } @@ -93,9 +94,8 @@ void WireIterator::backward() { return; } - // For sinks/deallocations/inserts, qubit_ is an OpOperand. Hence, only get - // the def-op. - if (isa(op_)) { + // For these operations, qubit_ is an OpOperand. Hence, only get the def-op. + if (isa(op_)) { op_ = qubit_.getDefiningOp(); isFinal_ = false; return; diff --git a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp index 891f03cf39..aac823bf2e 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp @@ -13,6 +13,7 @@ #include "mlir/Dialect/QCO/Utils/WireIterator.h" #include +#include #include #include #include @@ -21,6 +22,7 @@ #include #include +#include #include using namespace mlir; @@ -55,12 +57,19 @@ TEST_P(WireIteratorTest, Traversal) { const auto [q02, q11] = builder.cx(q01, q10); const auto [q03, c0] = builder.measure(q02); const auto q04 = builder.reset(q03); - const auto loopOut = builder.scfFor( - 1, 4, 1, {q04, q11}, [&builder](Value, ValueRange iterArgs) { - const auto iterQ00 = iterArgs[0]; - const auto iterQ10 = iterArgs[1]; - const auto iterQ01 = builder.h(iterQ00); - const auto [iterQ02, iterQ11] = builder.cx(iterQ01, iterQ10); + + Value iterQ00; + Value iterQ01; + Value iterQ02; + Value iterQ10; + Value iterQ11; + + const auto loopOut = + builder.scfFor(1, 4, 1, {q04, q11}, [&](Value, ValueRange iterArgs) { + iterQ00 = iterArgs[0]; + iterQ10 = iterArgs[1]; + iterQ01 = builder.h(iterQ00); + std::tie(iterQ02, iterQ11) = builder.cx(iterQ01, iterQ10); return SmallVector{iterQ02, iterQ11}; }); const auto q05 = loopOut[0]; @@ -162,6 +171,44 @@ TEST_P(WireIteratorTest, Traversal) { // // Test: Recursive use with block-argument. // + + qco::WireIterator recIt(iterQ00); + ASSERT_EQ(recIt.operation(), nullptr); // Blockargument + ASSERT_EQ(recIt.qubit(), iterQ00); + + ++recIt; + ASSERT_EQ(recIt.operation(), iterQ01.getDefiningOp()); // qco.h + ASSERT_EQ(recIt.qubit(), iterQ01); + + ++recIt; + ASSERT_EQ(recIt.operation(), iterQ02.getDefiningOp()); // qco.ctrl + ASSERT_EQ(recIt.qubit(), iterQ02); + + ++recIt; + ASSERT_EQ(recIt.operation(), *(iterQ02.getUsers().begin())); // scf.yield + ASSERT_EQ(recIt.qubit(), nullptr); + + ++recIt; + ASSERT_EQ(recIt, std::default_sentinel); + + ++recIt; + ASSERT_EQ(recIt, std::default_sentinel); + + --recIt; + ASSERT_EQ(recIt.operation(), *(iterQ02.getUsers().begin())); // scf.yield + ASSERT_EQ(recIt.qubit(), nullptr); + + --recIt; + ASSERT_EQ(recIt.operation(), iterQ02.getDefiningOp()); // qco.ctrl + ASSERT_EQ(recIt.qubit(), iterQ02); + + --recIt; + ASSERT_EQ(recIt.operation(), iterQ01.getDefiningOp()); // qco.h + ASSERT_EQ(recIt.qubit(), iterQ01); + + --recIt; + ASSERT_EQ(recIt.operation(), nullptr); // Blockargument + ASSERT_EQ(recIt.qubit(), iterQ00); } INSTANTIATE_TEST_SUITE_P(DynamicAndStatic, WireIteratorTest, ::testing::Bool(), From aff7da2c178df63452d22436d43bed5487502281 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Mon, 1 Jun 2026 15:31:16 +0200 Subject: [PATCH 09/25] WiP: recursive for-loop mapping --- .../Dialect/QCO/Transforms/Mapping/Mapping.h | 6 +- mlir/include/mlir/Dialect/QCO/Utils/Drivers.h | 67 +++- mlir/include/mlir/Dialect/QCO/Utils/Graph.h | 188 ++++++++- .../QCO/Transforms/Mapping/Mapping.cpp | 357 +++++++++++++----- mlir/lib/Dialect/QCO/Utils/Graph.cpp | 165 -------- .../Dialect/QTensor/Utils/TensorIterator.cpp | 6 +- .../QCO/Transforms/Mapping/test_mapping.cpp | 82 ++-- 7 files changed, 536 insertions(+), 335 deletions(-) delete mode 100644 mlir/lib/Dialect/QCO/Utils/Graph.cpp diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h b/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h index 96e4db0333..544f1872c7 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h @@ -11,7 +11,6 @@ #pragma once #include "mlir/Dialect/QCO/Transforms/Passes.h" -#include "mlir/Dialect/QCO/Utils/Graph.h" #include #include @@ -25,7 +24,8 @@ namespace mlir::qco { * @brief Create a mapping pass instance with the given target architecture. * @returns a pass object. */ -std::unique_ptr createMappingPass(const Graph::EdgeSet&, - MappingPassOptions); +std::unique_ptr +createMappingPass(const llvm::DenseSet>&, + MappingPassOptions); } // namespace mlir::qco diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h index 0a59e7d83e..e4f941bd7c 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h @@ -18,6 +18,8 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include +#include +#include #include #include #include @@ -88,13 +90,22 @@ LogicalResult walkProgram(Region& region, const WalkProgramFn& fn) { return success(); } -using ReleasedOps = SmallVector; -using PendingWiresMap = - DenseMap>; +using ReleasedOps = SmallVector; +using PendingWiresMap = DenseMap>; struct IsReady { bool operator()(PendingWiresMap::value_type& kv) const { - return kv.second.size() == kv.first.getNumQubits(); + const auto npending = kv.second.size(); + return TypeSwitch(kv.first) + .Case( + [&](auto& op) { return op.getNumQubits() == npending; }) + .template Case( + [&](auto& op) { return op.getInits().size() == npending; }) + .Default([&](Operation* op) { + const auto name = op->getName().getStringRef(); + reportFatalInternalError("unknown pending op: " + name); + return false; + }); } }; @@ -138,20 +149,20 @@ LogicalResult walkProgramGraph(MutableArrayRef wires, PendingWiresMap pending; pending.reserve(wires.size()); - SmallVector curr(wires.size()); + SmallVector curr(wires.size()); std::iota(curr.begin(), curr.end(), 0UL); - SmallVector next; + SmallVector next; next.reserve(wires.size()); while (!curr.empty()) { - for (std::size_t i : curr) { + for (size_t i : curr) { auto& it = wires[i]; while (Traits::isActive(it)) { const auto res = TypeSwitch(it.operation()) - .template Case([&](UnitaryOpInterface& op) { - // If there are fewer wires than the qubit requires inputs, + .template Case([&](auto& op) { + // If there are fewer wires than the unitary requires inputs, // it's impossible to release the operation. Hence, fail. if (op.getNumQubits() > wires.size()) { return WalkResult::interrupt(); @@ -173,19 +184,45 @@ LogicalResult walkProgramGraph(MutableArrayRef wires, indices.emplace_back(i); - return WalkResult::skip(); // Stop at multi-qubit gate. + return WalkResult::skip(); // Stop at multi-qubit unitary. + }) + .template Case([&](scf::ForOp& op) { + // If there are fewer wires than the loop requires inputs, + // it's impossible to release the operation. Hence, fail. + if (op.getInits().size() > wires.size()) { + return WalkResult::interrupt(); + } + + if (op.getInits().size() == 1) { + std::ranges::advance(it, Traits::stride()); + return WalkResult::advance(); + } + + // Insert the loop to the pending map. + // The caller decides if this op should be released. + const auto [mapIt, inserted] = pending.try_emplace(op); + auto& indices = mapIt->second; + + if (inserted) { + indices.reserve(op.getInits().size()); + } + + indices.emplace_back(i); + + return WalkResult::skip(); // Stop at multi-qubit loop. }) // AllocOp, StaticOp, and qtensor::ExtractOp are only reachable // on the forward path; backward isActive() halts before // reaching them (decrementing at a source op is a no-op). - .template Case([&](auto) { + .template Case([&](auto) { std::ranges::advance(it, Traits::stride()); return WalkResult::advance(); }) .Default([&](Operation* op) -> WalkResult { const auto name = op->getName().getStringRef(); - report_fatal_error("unknown op encountered: " + name); + reportFatalInternalError("unknown op encountered: " + name); }); if (res.wasSkipped()) { @@ -212,11 +249,11 @@ LogicalResult walkProgramGraph(MutableArrayRef wires, } } - for (const UnitaryOpInterface& op : released) { + for (Operation* op : released) { const auto mapIt = pending.find(op); assert(mapIt != pending.end()); - for (std::size_t i : mapIt->second) { + for (size_t i : mapIt->second) { std::ranges::advance(wires[i], Traits::stride()); next.emplace_back(i); } diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Graph.h b/mlir/include/mlir/Dialect/QCO/Utils/Graph.h index 2e78fd7ecf..1409158f79 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Graph.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Graph.h @@ -10,10 +10,12 @@ #pragma once +#include #include #include #include +#include #include #include #include @@ -22,42 +24,194 @@ namespace mlir::qco { template using Vector = SmallVector; template using Matrix = Vector>; -class Graph { -public: - using IdT = size_t; - using EdgeSet = llvm::DenseSet>; +enum class GraphType : bool { Directed, Undirected }; +template class Graph { +public: /// Construct an empty graph. Graph() = default; + /// Construct graph from edge set. - explicit Graph(const EdgeSet& edges); - /// Add a node to the graph. - void addNode(IdT id); - /// Add an edge to the graph. - void addEdge(IdT id, IdT neighbourId); + explicit Graph(const llvm::DenseSet>& edges) { + for_each(edges, [this](const auto& e) { addEdge(e); }); + } + + /// Add an edge to the graph. Implicitly adds nodes. + void addEdge(IdT u, IdT v) { + addDirectedEdge(u, v); + if constexpr (Type == GraphType::Undirected) { + addDirectedEdge(v, u); + } + } + /// Add an edge to the graph. - void addEdge(std::pair edge); - /// Add multiple edges to the graph. - void addEdges(SmallVector> edges); + void addEdge(std::pair edge) { addEdge(edge.first, edge.second); } + /// Return a set of edges. - [[nodiscard]] EdgeSet getEdges() const; + [[nodiscard]] llvm::DenseSet> getEdges() const { + llvm::DenseSet> edges; + for (const auto& [u, nbrs] : adj_) { + for (const auto& v : nbrs) { + if constexpr (Type == GraphType::Directed) { + edges.insert(std::make_pair(u, v)); + } else { + edges.insert(std::minmax(u, v)); + } + } + } + return edges; + } + /// Return the edges of a node. - [[nodiscard]] ArrayRef getEdges(size_t id) const; + [[nodiscard]] ArrayRef getEdges(size_t id) const { return adj_.at(id); } + + /// Return the nodes. + [[nodiscard]] ArrayRef getNodes() const { + return to_vector(adj_.keys()); + } /// Return the number of nodes. [[nodiscard]] size_t getNumNodes() const { return adj_.size(); } + + /// Return the degree of a node. + [[nodiscard]] size_t getDegree(size_t id) { return adj_.at(id).size(); } + /// Return the max degree of the graph. - [[nodiscard]] size_t getMaxDegree() const; + [[nodiscard]] size_t getMaxDegree() const { + size_t deg = 0; + for (const auto& [u, nbrs] : adj_) { + deg = std::max(deg, nbrs.size()); + } + return deg; + } + + /// Return true if the graph has no nodes and edges. + [[nodiscard]] bool empty() const { return adj_.empty(); } + + /// Clear the graph. + [[nodiscard]] void clear() { adj_.clear(); } + /// Return the minimum distance matrix of the graph by implementing the /// Floyd-Warshall Algorithm /// (https://en.wikipedia.org/wiki/Floyd–Warshall_algorithm) where dist[i][j] /// denotes the distance between i and j. - [[nodiscard]] Matrix getDistMatrix() const; + [[nodiscard]] Matrix getDistMatrix() const { + const auto n = getNumNodes(); + + SmallVector> edges; + for (const auto& [u, nbrs] : adj_) { + for (const auto& v : nbrs) { + edges.emplace_back(u, v); + } + } + + Matrix dist(n, Vector(n, UINT64_MAX)); + for (const auto& [u, v] : edges) { + dist[u][v] = 1; + } + for (size_t v = 0; v < n; ++v) { + dist[v][v] = 0; + } + + for (size_t k = 0; k < n; ++k) { + for (size_t i = 0; i < n; ++i) { + for (size_t j = 0; j < n; ++j) { + if (dist[i][k] == UINT64_MAX || dist[k][j] == UINT64_MAX) { + continue; // Avoid overflow with "infinite" distances. + } + + const size_t sum = dist[i][k] + dist[k][j]; + dist[i][j] = std::min(dist[i][j], sum); + } + } + } + + return dist; + } + /// Return reverse cycle in graph or `std::nullopt` if none exists. /// Implements an iterative depth-first search inspired by LLVM's SCC /// utilities. - [[nodiscard]] std::optional> findCycle() const; + [[nodiscard]] std::optional> findCycle() const { + enum struct State : uint8_t { Unseen, Seen, Finished }; + + struct Frame { + IdT id; + size_t neighbourIdx; + }; + + SmallVector stack; + llvm::DenseMap parents; + llvm::DenseMap states; + + // Preparation step: Mark all nodes as unseen. + for_each(adj_.keys(), [&](IdT id) { states[id] = State::Unseen; }); + + for (const auto initId : adj_.keys()) { + // Only start from unseen nodes. + if (states[initId] != State::Unseen) { + continue; + } + + stack.emplace_back(initId, 0); + + while (!stack.empty()) { + Frame& top = stack.back(); + + // If we haven't seen this node before, mark it as seen. + if (states[top.id] == State::Unseen) { + states[top.id] = State::Seen; + } + + auto it = adj_.find(top.id); + assert(it != adj_.end() && "expected node id in adjacency map"); + const auto nbrs = it->getSecond(); + + // Once all neighbours have been visited (indicated by the index + // exceeding the number of neighbours - 1), set the frame on node to + // finished and pop it from the stack. + if (top.neighbourIdx >= nbrs.size()) { + states[top.id] = State::Finished; + stack.pop_back(); + continue; + } + + // Collect the neighbour and advance the index on the + // frame for the next iteration. + const auto nbrId = nbrs[top.neighbourIdx]; + ++top.neighbourIdx; + + if (states[nbrId] == State::Unseen) { + parents[nbrId] = top.id; + stack.emplace_back(nbrId, 0); + } else if (states[nbrId] == State::Seen) { + SmallVector path; + for (auto curr = top.id; curr != nbrId; curr = parents[curr]) { + path.emplace_back(curr); + } + path.emplace_back(nbrId); + return path; + } + } + + // Preparse stack for next iteration. + stack.clear(); + } + + return std::nullopt; + } private: + /// Adds a directed edge to the internal representation of the graph. + void addDirectedEdge(IdT u, IdT v) { + if (!adj_.contains(u)) { + adj_[u] = Vector(); + } + if (!adj_.contains(v)) { + adj_[v] = Vector(); + } + adj_[u].emplace_back(v); + } + llvm::DenseMap> adj_; }; } // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index cfeb0b5bf7..6fef2ba282 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -11,9 +11,10 @@ #include "mlir/Dialect/QCO/Transforms/Mapping/Mapping.h" #include "mlir/Dialect/QCO/IR/QCODialect.h" +#include "mlir/Dialect/QCO/IR/QCOInterfaces.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" -#include "mlir/Dialect/QCO/Utils/Graph.h" #include "mlir/Dialect/QCO/Utils/Drivers.h" +#include "mlir/Dialect/QCO/Utils/Graph.h" #include "mlir/Dialect/QCO/Utils/WireIterator.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include "mlir/Dialect/QTensor/Utils/TensorIterator.h" @@ -23,9 +24,12 @@ #include #include #include +#include #include +#include #include #include +#include #include #include #include @@ -217,7 +221,8 @@ struct MappingPass : impl::MappingPassBase { public: AugmentedDevice() = default; - explicit AugmentedDevice(const Graph::EdgeSet& couplingSet) + explicit AugmentedDevice( + const llvm::DenseSet>& couplingSet) : coupling_(couplingSet), dist_(coupling_.getDistMatrix()) {} /** @@ -250,13 +255,27 @@ struct MappingPass : impl::MappingPassBase { return coupling_.getEdges(u); } + /** + * @returns the qubit identifiers. + */ + [[nodiscard]] ArrayRef qubits() const { + return coupling_.getNodes(); + } + + /** + * @returns the links of the device. + */ + [[nodiscard]] llvm::DenseSet> links() const { + return coupling_.getEdges(); + } + /** * @returns the max degree (connectivity) of any qubit of the device. */ [[nodiscard]] size_t maxDegree() const { return coupling_.getMaxDegree(); } private: - Graph coupling_; + Graph coupling_; Matrix dist_; }; @@ -357,12 +376,21 @@ struct MappingPass : impl::MappingPassBase { } }; + /// A unit is a (to be) routeable region. + struct Unit { + SmallVector wires; + Layout layout; + }; + public: + /// Construct default mapping pass. MappingPass() = default; + /// Construct default mapping pass with options. explicit MappingPass(MappingPassOptions options) : MappingPassBase(options) {} - /// @brief Construct Mapping Pass from coupling-set. - explicit MappingPass(const Graph::EdgeSet& couplingSet, - MappingPassOptions options = {}) + /// Construct mapping from coupling set. + explicit MappingPass( + const llvm::DenseSet>& couplingSet, + MappingPassOptions options = {}) : MappingPassBase(options), device(couplingSet) {} protected: @@ -398,38 +426,45 @@ struct MappingPass : impl::MappingPassBase { // Create trials for initial layout refining. Currently, this includes // `ntrials` many random layouts. - SmallVector trials; - trials.reserve(ntrials); - for (size_t i = 0; i < ntrials; ++i) { - trials.emplace_back(Layout::random(device.nqubits(), rng())); - } + // SmallVector trials; + // trials.reserve(ntrials); + // for (size_t i = 0; i < ntrials; ++i) { + // trials.emplace_back(Layout::random(device.nqubits(), rng())); + // } // Execute each of the trials (possibly in parallel). Collect the results // and find the one with the fewest SWAPs on the final backwards pass. - parallelForEach(&getContext(), trials, [&, this](Trial& trial) { - if (const auto res = refineLayout(*comp, trial.layout); succeeded(res)) { - trial.success = true; - trial.nswaps = *res; - } - }); - - Trial* best = findBestTrial(trials); - if (best == nullptr) { - func.emitError() << "failed to find the best layout trial"; - signalPassFailure(); - return; - } - + // parallelForEach(&getContext(), trials, [&, this](Trial& trial) { + // if (const auto res = refineLayout(*comp, trial.layout); succeeded(res)) + // { + // trial.success = true; + // trial.nswaps = *res; + // } + // }); + + // Trial* best = findBestTrial(trials); + // if (best == nullptr) { + // func.emitError() << "failed to find the best layout trial"; + // signalPassFailure(); + // return; + // } + + Layout l = Layout::identity(device.nqubits()); // Perform placement and hot routing by inserting SWAPs into the IR. - auto placedWires = place(func, best->layout, rewriter); - const auto res = route( - placedWires, best->layout, &rewriter); + auto placedWires = place(func, l, rewriter); + Unit u{.wires = placedWires, .layout = l}; + const auto res = + route(u, &rewriter); if (failed(res)) { func.emitError() << "failed to map the " << func.getName() << " function"; signalPassFailure(); return; } + assert(llvm::all_of(u.wires, + [](auto& it) { return isa(it.operation()); + })); + // Collect statistics. numSwaps += *res; @@ -560,7 +595,7 @@ struct MappingPass : impl::MappingPassBase { * @brief Find the best trial result in terms of the number of SWAPs. * @returns the best trial result or nullptr if no result is valid. */ - [[nodiscard]] static Trial* findBestTrial(MutableArrayRef trials) { + static Trial* findBestTrial(MutableArrayRef trials) { Trial* best = nullptr; for (auto& trial : trials) { if (!trial.success) { @@ -583,15 +618,14 @@ struct MappingPass : impl::MappingPassBase { * along the way. Repeat this procedure "niterations" times. * @returns failure() if routing fails. */ - FailureOr refineLayout(SmallVector wires, - Layout& layout) { + FailureOr refineLayout(Unit& u) { size_t nswaps{0}; for (size_t i = 0; i < niterations; ++i) { - if (failed(route(wires, layout))) { + if (failed(route(u))) { return failure(); } - const auto resB = route(wires, layout); + const auto resB = route(u); if (failed(resB)) { return failure(); } @@ -618,8 +652,8 @@ struct MappingPass : impl::MappingPassBase { * @returns a vector of hardware-index pairs (each denoting a SWAP) or * failure() if A* fails. */ - [[nodiscard]] FailureOr> - search(const Window& window, const Layout& layout) { + FailureOr> search(const Window& window, + const Layout& layout) { constexpr size_t cap = 25'000'000UL; const size_t b = device.maxDegree() * ((device.nqubits() + 1) / 2); @@ -768,11 +802,13 @@ struct MappingPass : impl::MappingPassBase { } for (const auto& [op, progs] : ready) { - if (isa(op)) { + if (isa(op)) { released.emplace_back(op); continue; } + assert(isa(op)); + const auto p0 = progs[0]; const auto p1 = progs[1]; window.emplace_back(p0, p1); @@ -791,27 +827,49 @@ struct MappingPass : impl::MappingPassBase { } /** + * @todo UPDATE THIS DESCRIPTION * @brief Advance past all executable gates. * @details Traverses the multi-qubit gates of the circuit until no more * executable gates are found. */ template - void skipExecutableGates(MutableArrayRef wires, - Layout& layout) { + std::pair> advanceFrame(Unit& unit) { + Operation* subOp = nullptr; + SmallVector subs; + std::ignore = walkProgramGraph( - wires, [&](const ReadyRange& ready, ReleasedOps& released) { + unit.wires, [&](const ReadyRange& ready, ReleasedOps& released) { if (ready.empty()) { return WalkResult::advance(); } for (const auto& [op, progs] : ready) { + /// TODO: TypeSwitch if (isa(op)) { released.emplace_back(op); continue; } + if (auto forOp = dyn_cast(op)) { + SmallVector rWires; + for (auto prog : progs) { + const auto res = cast(unit.wires[prog].qubit()); + auto& it = + rWires.emplace_back(forOp.getTiedLoopRegionIterArg(res)); + // The iterator points at a block argument. Move the iterator to + // a position where it.operation() != nullptr. + std::ranges::advance(it, + WireTraversalTraits::stride()); + } + subOp = op; + subs.emplace_back(rWires, unit.layout); + return WalkResult::interrupt(); + } + + assert(isa(op)); + const auto [hw0, hw1] = - layout.getHardwareIndices(progs[0], progs[1]); + unit.layout.getHardwareIndices(progs[0], progs[1]); if (device.areAdjacent(hw0, hw1)) { released.emplace_back(op); @@ -825,6 +883,101 @@ struct MappingPass : impl::MappingPassBase { return WalkResult::advance(); }); + + return std::make_pair(subOp, subs); + } + + SmallVector restore(const Layout& from, const Layout& to) { + SmallVector swaps; + + Layout curr(from); + + Graph g; + const auto constructEdge = [&](size_t hwX, size_t hwY) { + const auto prog = curr.getProgramIndex(hwX); + + const auto hwGoal = to.getHardwareIndex(prog); + const auto distPre = device.distanceBetween(hwX, hwGoal); + const auto distPost = device.distanceBetween(hwY, hwGoal); + + if (distPost < distPre) { + llvm::dbgs() << "prog=" << prog << " edge=(" << hwX << ", " << hwY + << ")" + << " dist(pre)=" << distPre << " dist(post)=" << distPost + << '\n'; + g.addEdge(hwX, hwY); + } + }; + + do { + // Construct 'F' graph. + g.clear(); + for (const auto& [hwA, hwB] : device.links()) { + constructEdge(hwA, hwB); + constructEdge(hwB, hwA); + } + + // Find happy swap chain or unhappy swap. + if (const auto cycle = g.findCycle(); cycle) { + // Apply happy SWAP chain. + for (size_t i = 0; i < cycle->size() - 1; ++i) { + llvm::dbgs() << "happySWAP=(" << (*cycle)[i] << ", " + << (*cycle)[i + 1] << ")\n"; + curr.swap((*cycle)[i], (*cycle)[i + 1]); + swaps.emplace_back((*cycle)[i], (*cycle)[i + 1]); + } + + continue; + } + + for (const auto e : g.getEdges()) { + if (g.getDegree(e.second) == 0) { + llvm::dbgs() << "unhappySWAP=(" << e.first << ", " << e.second + << ")\n"; + curr.swap(e.first, e.second); + swaps.emplace_back(e.first, e.second); + break; + } + } + } while (!g.empty()); + + return swaps; + } + + template + void applySWAPs(Unit& unit, ArrayRef swaps, + IRRewriter* rewriter) { + for (const auto& [hw0, hw1] : swaps) { + if constexpr (Mode == RoutingMode::Hot) { + const auto& [prog0, prog1] = unit.layout.getProgramIndices(hw0, hw1); + const auto& w0 = unit.wires[prog0]; + const auto& w1 = unit.wires[prog1]; + + assert(!isa(w0.operation())); + assert(!isa(w1.operation())); + + const auto in0 = w0.qubit(); + const auto in1 = w1.qubit(); + + rewriter->setInsertionPointAfter(in0.getDefiningOp()); + auto swapOp = SWAPOp::create(*rewriter, in0.getLoc(), in0, in1); + + const auto out0 = swapOp.getQubit0Out(); + const auto out1 = swapOp.getQubit1Out(); + + rewriter->replaceAllUsesExcept(in0, out1, swapOp); + rewriter->replaceAllUsesExcept(in1, out0, swapOp); + + // Preserve program-indexed wire semantics. + unit.wires[prog0] = WireIterator(out1); + unit.wires[prog1] = WireIterator(out0); + + assert(isa(w0.operation())); + assert(isa(w1.operation())); + } + + unit.layout.swap(hw0, hw1); + } } /** @@ -836,82 +989,95 @@ struct MappingPass : impl::MappingPassBase { * @returns failure() if A* search isn't able to find a solution, the number * of SWAPs otherwise. */ - template - FailureOr route(SmallVector& wires, Layout& layout, - IRRewriter* rewriter = nullptr) { + template + FailureOr route(Unit& unit, IRRewriter* rewriter = nullptr) { + assert(Mode == RoutingMode::Cold || + (Mode == RoutingMode::Hot && Direction == WireDirection::Forward)); + using Traits = WireTraversalTraits; size_t nswaps{0}; + while (true) { - skipExecutableGates(wires, layout); + auto [op, subUnits] = advanceFrame(unit); + if (op != nullptr) { + for (auto& subUnit : subUnits) { + const auto res = route(subUnit, rewriter); + if (failed(res)) { + return failure(); + } + nswaps += *res; + } + + assert(isa(op)); + + // An scf.for operation has exactly one sub-unit. + // To finalize the mapping of the sub-unit, we append an appendix + // of SWAPs to restore the base-layout. + + auto& subUnit = subUnits.front(); + const auto swaps = restore(subUnit.layout, unit.layout); + + if constexpr (Mode == RoutingMode::Hot) { + + // At this point the wire iterators point to scf.yield ops. Thus, + // decrement the iterator to point at a valid insertion point. Because + // the wire iterators are not required any more, there is no need to + // increment them again after applying the SWAP sequence. - const auto window = getWindow(wires); + for_each(subUnit.wires, [](auto& it) { + std::ranges::advance(it, -Traits::stride()); + }); + } + + applySWAPs(subUnit, swaps, rewriter); + nswaps += swaps.size(); + } + + const auto window = getWindow(unit.wires); if (window.empty()) { break; } - if constexpr (mode == RoutingMode::Hot) { + for (const auto [i0, i1] : window) { + llvm::dbgs() << "(" << i0 << ", " << i1 << ") "; + } + llvm::dbgs() << '\n'; + + const auto swaps = search(window, unit.layout); + if (failed(swaps)) { + return failure(); + } + + if constexpr (Mode == RoutingMode::Hot) { // At this point the wire iterators either point to // std::default_sentinel or a multi-qubit gate (including barriers) of // the current or subsequent layers. The former must be decremented - // twice (sentinel -> sink -> unitary/static). For the latter we simply - // must ensure the insertion point is before the multi-qubit gates. + // twice (sentinel -> sink -> unitary/static). For the latter we + // simply must ensure the insertion point is before the multi-qubit + // gates. - for (auto& it : wires) { + for (auto& it : unit.wires) { std::ranges::advance(it, it == std::default_sentinel ? -2 * Traits::stride() : -Traits::stride()); } } - const auto swaps = search(window, layout); - if (failed(swaps)) { - return failure(); - } - - for (const auto& [hw0, hw1] : *swaps) { - if constexpr (mode == RoutingMode::Hot) { - const auto& [prog0, prog1] = layout.getProgramIndices(hw0, hw1); - const auto& w0 = wires[prog0]; - const auto& w1 = wires[prog1]; - - assert(!isa(w0.operation())); - assert(!isa(w1.operation())); - - const auto in0 = w0.qubit(); - const auto in1 = w1.qubit(); - - rewriter->setInsertionPointAfter(in0.getDefiningOp()); - auto swapOp = SWAPOp::create(*rewriter, in0.getLoc(), in0, in1); - - const auto out0 = swapOp.getQubit0Out(); - const auto out1 = swapOp.getQubit1Out(); - - rewriter->replaceAllUsesExcept(in0, out1, swapOp); - rewriter->replaceAllUsesExcept(in1, out0, swapOp); - - // Preserve program-indexed wire semantics. - wires[prog0] = WireIterator(out1); - wires[prog1] = WireIterator(out0); - - assert(isa(w0.operation())); - assert(isa(w1.operation())); - } - layout.swap(hw0, hw1); - } + applySWAPs(unit, *swaps, rewriter); - if constexpr (mode == RoutingMode::Hot) { + if constexpr (Mode == RoutingMode::Hot) { // After SWAP insertion, a wire is either untouched by the SWAP - // insertion or pointing at a SWAP operation. If the former is the case, - // incrementing the wire iterator will undo the previous decrement, - // leaving it at the same position as before the SWAP insertion. - // Otherwise, an increment will move the iterator to the multi-qubit op - // of the current or subsequent layer or to a sink (and thus - // std::default_sentinel). - - for_each(wires, + // insertion or pointing at a SWAP operation. If the former is the + // case, incrementing the wire iterator will undo the previous + // decrement, leaving it at the same position as before the SWAP + // insertion. Otherwise, an increment will move the iterator to the + // multi-qubit op of the current or subsequent layer or to a sink (and + // thus std::default_sentinel). + + for_each(unit.wires, [](auto& it) { std::ranges::advance(it, Traits::stride()); }); } @@ -926,8 +1092,9 @@ struct MappingPass : impl::MappingPassBase { } // namespace -std::unique_ptr createMappingPass(const Graph::EdgeSet& couplingSet, - MappingPassOptions options) { +std::unique_ptr +createMappingPass(const llvm::DenseSet>& couplingSet, + MappingPassOptions options) { return std::make_unique(couplingSet, options); } diff --git a/mlir/lib/Dialect/QCO/Utils/Graph.cpp b/mlir/lib/Dialect/QCO/Utils/Graph.cpp deleted file mode 100644 index d25e16324f..0000000000 --- a/mlir/lib/Dialect/QCO/Utils/Graph.cpp +++ /dev/null @@ -1,165 +0,0 @@ -/* - * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM - * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH - * All rights reserved. - * - * SPDX-License-Identifier: MIT - * - * Licensed under the MIT License - */ - -#include "mlir/Dialect/QCO/Utils/Graph.h" - -#include -#include -#include -#include - -#include -#include -#include -#include - -namespace mlir::qco { -Graph::Graph(const EdgeSet& edges) { - for (const auto& [u, v] : edges) { - if (!adj_.contains(u)) { - adj_[u] = Vector(); - } - adj_[u].emplace_back(v); - } -} -void Graph::addNode(size_t id) { - const auto r = adj_.try_emplace(id, SmallVector{}); - assert(r.second && "expected to insert node"); -} - -void Graph::addEdge(size_t id, size_t neighbourId) { - assert(adj_.contains(id) && "addEdge: missing node id"); - adj_[id].emplace_back(neighbourId); -} - -void Graph::addEdge(std::pair edge) { - addEdge(edge.first, edge.second); -} - -void Graph::addEdges(SmallVector> edges) { - for_each(edges, [this](const auto& edge) { addEdge(edge); }); -} - -Graph::EdgeSet Graph::getEdges() const { - EdgeSet set; - for (const auto& [u, nbrs] : adj_) { - for (const auto& v : nbrs) { - set.insert(std::make_pair(u, v)); - } - } - return set; -} - -ArrayRef Graph::getEdges(size_t id) const { return adj_.at(id); } - -size_t Graph::getMaxDegree() const { - size_t deg = 0; - for (const auto& [u, nbrs] : adj_) { - deg = std::max(deg, nbrs.size()); - } - return deg; -} - -Matrix Graph::getDistMatrix() const { - const size_t n = getNumNodes(); - const auto edges = getEdges(); - - Matrix dist(n, Vector(n, UINT64_MAX)); - for (const auto& [u, v] : edges) { - dist[u][v] = 1; - } - for (size_t v = 0; v < n; ++v) { - dist[v][v] = 0; - } - for (size_t k = 0; k < n; ++k) { - for (size_t i = 0; i < n; ++i) { - for (size_t j = 0; j < n; ++j) { - if (dist[i][k] == UINT64_MAX || dist[k][j] == UINT64_MAX) { - continue; // Avoid overflow with "infinite" distances. - } - - const size_t sum = dist[i][k] + dist[k][j]; - dist[i][j] = std::min(dist[i][j], sum); - } - } - } - - return dist; -} - -[[nodiscard]] std::optional> Graph::findCycle() const { - enum struct State : uint8_t { Unseen, Seen, Finished }; - - struct Frame { - IdT id; - size_t neighbourIdx; - }; - - SmallVector stack; - llvm::DenseMap parents; - llvm::DenseMap states; - - // Preparation step: Mark all nodes as unseen. - for_each(adj_.keys(), [&](IdT id) { states[id] = State::Unseen; }); - - for (const auto initId : adj_.keys()) { - // Only start from unseen nodes. - if (states[initId] != State::Unseen) { - continue; - } - - stack.emplace_back(initId, 0); - - while (!stack.empty()) { - Frame& top = stack.back(); - - // If we haven't seen this node before, mark it as seen. - if (states[top.id] == State::Unseen) { - states[top.id] = State::Seen; - } - - auto it = adj_.find(top.id); - assert(it != adj_.end() && "expected node id in adjacency map"); - const auto nbrs = it->getSecond(); - - // Once all neighbours have been visited (indicated by the index exceeding - // the number of neighbours - 1), set the frame on node to finished and - // pop it from the stack. - if (top.neighbourIdx >= nbrs.size()) { - states[top.id] = State::Finished; - stack.pop_back(); - continue; - } - - // Collect the neighbour and advance the index on the - // frame for the next iteration. - const auto nbrId = nbrs[top.neighbourIdx]; - ++top.neighbourIdx; - - if (states[nbrId] == State::Unseen) { - parents[nbrId] = top.id; - stack.emplace_back(nbrId, 0); - } else if (states[nbrId] == State::Seen) { - SmallVector path; - for (auto curr = top.id; curr != nbrId; curr = parents[curr]) { - path.emplace_back(curr); - } - path.emplace_back(nbrId); - return path; - } - } - - // Preparse stack for next iteration. - stack.clear(); - } - - return std::nullopt; -} -} // namespace mlir::qco diff --git a/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp b/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp index 064cdbd624..6af8c5bff7 100644 --- a/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp +++ b/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp @@ -26,12 +26,8 @@ namespace mlir::qtensor { TypedValue TensorIterator::tensor() const { - if (op_ == nullptr) { - return tensor_; - } - // The following operations don't have an OpResult. - if (isa(op_)) { + if (op_ != nullptr && isa(op_)) { return nullptr; } diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 5a9ba46e20..6fc38d37a9 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -14,7 +14,6 @@ #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/Transforms/Mapping/Mapping.h" #include "mlir/Dialect/QCO/Transforms/Passes.h" -#include "mlir/Dialect/QCO/Utils/Graph.h" #include "mlir/Dialect/QCO/Utils/Drivers.h" #include "mlir/Dialect/QCO/Utils/Qubits.h" @@ -43,12 +42,18 @@ using namespace mlir; using namespace mlir::qco; +struct Device { + size_t nqubits{}; + llvm::DenseSet> couplingSet; +}; + /** * @returns llvm::success() if all two-qubit gates inside @p region * fulfill the given coupling constraints. llvm::failure(), otherwise. */ -static LogicalResult isExecutable(Region& region, - const Graph::EdgeSet& couplingSet) { +static LogicalResult +isExecutable(Region& region, + const llvm::DenseSet>& couplingSet) { return walkProgram(region, [&](Operation* curr, const Qubits& qubits) { if (auto op = dyn_cast(curr)) { if (isa(op)) { @@ -77,17 +82,18 @@ static LogicalResult isExecutable(Region& region, /** * @returns a 9x9 square-grid coupling set. */ -static Graph::EdgeSet getNineQubitSquareGrid() { - return Graph::EdgeSet{{0, 3}, {3, 0}, {0, 1}, {1, 0}, {1, 4}, {4, 1}, - {1, 2}, {2, 1}, {2, 5}, {5, 2}, {3, 6}, {6, 3}, - {3, 4}, {4, 3}, {4, 7}, {7, 4}, {4, 5}, {5, 4}, - {5, 8}, {8, 5}, {6, 7}, {7, 6}, {7, 8}, {8, 7}}; +static Device getNineQubitSquareGrid() { + return {.nqubits = 9, + .couplingSet = {{0, 3}, {3, 0}, {0, 1}, {1, 0}, {1, 4}, {4, 1}, + {1, 2}, {2, 1}, {2, 5}, {5, 2}, {3, 6}, {6, 3}, + {3, 4}, {4, 3}, {4, 7}, {7, 4}, {4, 5}, {5, 4}, + {5, 8}, {8, 5}, {6, 7}, {7, 6}, {7, 8}, {8, 7}}}; } namespace { class MappingPassTest : public testing::Test, - public testing::WithParamInterface { + public testing::WithParamInterface { protected: void SetUp() override { DialectRegistry registry; @@ -98,8 +104,10 @@ class MappingPassTest : public testing::Test, context->loadAllAvailableDialects(); } - static LogicalResult runPass(ModuleOp m, const Graph::EdgeSet& couplingSet, - const MappingPassOptions& options) { + static LogicalResult + runPass(ModuleOp m, + const llvm::DenseSet>& couplingSet, + const MappingPassOptions& options) { PassManager pm(m->getContext()); pm.addPass(createMappingPass(couplingSet, options)); return pm.run(m); @@ -111,15 +119,15 @@ class MappingPassTest : public testing::Test, }; // namespace TEST_P(MappingPassTest, NoEntryPoint) { - const auto& couplingSet = GetParam(); + const auto& device = GetParam(); OwningOpRef m = ModuleOp::create(UnknownLoc::get(context.get())); - auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); + auto res = runPass(m.get(), device.couplingSet, MappingPassOptions{}); ASSERT_TRUE(res.failed()); } TEST_P(MappingPassTest, NoQubitAllocations) { - const auto& couplingSet = GetParam(); + const auto& device = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -129,13 +137,13 @@ TEST_P(MappingPassTest, NoQubitAllocations) { builder.sink(q0); auto m = builder.finalize(); - auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); + auto res = runPass(m.get(), device.couplingSet, MappingPassOptions{}); ASSERT_TRUE(res.failed()); } TEST_P(MappingPassTest, NoExtractAfterInsert) { - const auto& couplingSet = GetParam(); + const auto& device = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -154,15 +162,14 @@ TEST_P(MappingPassTest, NoExtractAfterInsert) { builder.qtensorDealloc(tensor0); auto m = builder.finalize(); - auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); + auto res = runPass(m.get(), device.couplingSet, MappingPassOptions{}); ASSERT_TRUE(res.failed()); } TEST_P(MappingPassTest, TooManyQubitsForArch) { - const auto& couplingSet = GetParam(); - const size_t nqubits = Graph(couplingSet).getNumNodes(); - const auto n = static_cast(nqubits) + 1; + const auto& device = GetParam(); + const auto n = static_cast(device.nqubits) + 1; QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -183,13 +190,13 @@ TEST_P(MappingPassTest, TooManyQubitsForArch) { builder.qtensorDealloc(tensor); auto m = builder.finalize(); - auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); + auto res = runPass(m.get(), device.couplingSet, MappingPassOptions{}); ASSERT_TRUE(res.failed()); } TEST_P(MappingPassTest, GHZ) { - const auto& couplingSet = GetParam(); + const auto& device = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -215,23 +222,23 @@ TEST_P(MappingPassTest, GHZ) { builder.qtensorDealloc(tensor); auto m = builder.finalize(); - auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); + auto res = runPass(m.get(), device.couplingSet, MappingPassOptions{}); auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); + EXPECT_TRUE( + isExecutable(entry.getFunctionBody(), device.couplingSet).succeeded()); } TEST_P(MappingPassTest, GHZUnrolled) { - const auto& couplingSet = GetParam(); - const size_t nqubits = Graph(couplingSet).getNumNodes(); - const auto n = static_cast(nqubits); + const auto& device = GetParam(); + const auto n = static_cast(device.nqubits); PassManager pm(context.get()); pm.addNestedPass(createQuantumLoopUnroll()); pm.addPass(createCSEPass()); pm.addPass(createCanonicalizerPass()); - pm.addPass(createMappingPass(couplingSet, MappingPassOptions{})); + pm.addPass(createMappingPass(device.couplingSet, MappingPassOptions{})); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -264,14 +271,17 @@ TEST_P(MappingPassTest, GHZUnrolled) { auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); + EXPECT_TRUE( + isExecutable(entry.getFunctionBody(), device.couplingSet).succeeded()); } TEST_P(MappingPassTest, GroverLike) { - const auto& couplingSet = GetParam(); + const auto& device = GetParam(); PassManager pm(context.get()); - pm.addPass(createMappingPass(couplingSet, MappingPassOptions{})); + pm.addPass(createMappingPass( + device.couplingSet, + MappingPassOptions{.ntrials = 1, .niterations = 1, .nlookahead = 2})); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -351,11 +361,12 @@ TEST_P(MappingPassTest, GroverLike) { auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); + EXPECT_TRUE( + isExecutable(entry.getFunctionBody(), device.couplingSet).succeeded()); } TEST_P(MappingPassTest, Sabre) { - const auto& couplingSet = GetParam(); + const auto& device = GetParam(); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -441,11 +452,12 @@ TEST_P(MappingPassTest, Sabre) { builder.qtensorDealloc(tensorDown); auto m = builder.finalize(); - auto res = runPass(m.get(), couplingSet, MappingPassOptions{}); + auto res = runPass(m.get(), device.couplingSet, MappingPassOptions{}); auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE(isExecutable(entry.getFunctionBody(), couplingSet).succeeded()); + EXPECT_TRUE( + isExecutable(entry.getFunctionBody(), device.couplingSet).succeeded()); } INSTANTIATE_TEST_SUITE_P(NineQubitSquareGrid, MappingPassTest, From 20243a5e61d968a66cf1f1f1a9a3c5e96e441493 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Mon, 1 Jun 2026 15:55:17 +0200 Subject: [PATCH 10/25] Update walkProgram driver --- mlir/include/mlir/Dialect/QCO/Utils/Drivers.h | 14 +++- .../QCO/Transforms/Mapping/test_mapping.cpp | 44 ++++++------ .../Dialect/QCO/Utils/test_drivers.cpp | 69 ++++++++++++------- 3 files changed, 79 insertions(+), 48 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h index e4f941bd7c..5e18fe7e18 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h @@ -46,12 +46,13 @@ using WalkProgramFn = function_ref; * Depending on the template parameter, the callback is executed before or after * updating the Qubits state. * @param region The targeted region. + * @param qubits An empty or filled qubits object. * @param fn The callback function. * @returns success(), if all operations have been visited. */ template -LogicalResult walkProgram(Region& region, const WalkProgramFn& fn) { - Qubits qubits; +LogicalResult walkProgram(Region& region, Qubits& qubits, + const WalkProgramFn& fn) { for (Operation& curr : region.getOps()) { if constexpr (Order == WalkOrder::PreOrder) { if (fn(&curr, qubits).wasInterrupted()) { @@ -71,6 +72,14 @@ LogicalResult walkProgram(Region& region, const WalkProgramFn& fn) { qubits.remap(prevQ, nextQ); } }) + .template Case([&](scf::ForOp op) { + for (OpOperand& operand : op.getInitsMutable()) { + auto result = op.getTiedLoopResult(&operand); + const auto prevQ = cast>(operand.get()); + const auto nextQ = cast>(result); + qubits.remap(prevQ, nextQ); + } + }) .template Case([&](ResetOp op) { qubits.remap(op.getQubitIn(), op.getQubitOut()); }) @@ -86,7 +95,6 @@ LogicalResult walkProgram(Region& region, const WalkProgramFn& fn) { } } } - return success(); } diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 6fc38d37a9..5b2a9fbc45 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -54,29 +54,31 @@ struct Device { static LogicalResult isExecutable(Region& region, const llvm::DenseSet>& couplingSet) { - return walkProgram(region, [&](Operation* curr, const Qubits& qubits) { - if (auto op = dyn_cast(curr)) { - if (isa(op)) { - return WalkResult::advance(); - } - - assert(op.getNumQubits() <= 2 && - "isExecutable: expected two-qubit gate decomposition"); - - if (op.getNumQubits() > 1) { - const auto q0 = cast>(op.getInputQubit(0)); - const auto q1 = cast>(op.getInputQubit(1)); - const auto i0 = qubits.getIndex(q0); - const auto i1 = qubits.getIndex(q1); - - if (!couplingSet.contains(std::make_pair(i0, i1))) { - return WalkResult::interrupt(); + Qubits qubits; + return walkProgram( + region, qubits, [&](Operation* curr, const Qubits& qubits) { + if (auto op = dyn_cast(curr)) { + if (isa(op)) { + return WalkResult::advance(); + } + + assert(op.getNumQubits() <= 2 && + "isExecutable: expected two-qubit gate decomposition"); + + if (op.getNumQubits() > 1) { + const auto q0 = cast>(op.getInputQubit(0)); + const auto q1 = cast>(op.getInputQubit(1)); + const auto i0 = qubits.getIndex(q0); + const auto i1 = qubits.getIndex(q1); + + if (!couplingSet.contains(std::make_pair(i0, i1))) { + return WalkResult::interrupt(); + } + } } - } - } - return WalkResult::advance(); - }); + return WalkResult::advance(); + }); } /** diff --git a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp index 616e08d66d..9606398018 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -36,7 +37,8 @@ class DriversTest : public testing::Test { protected: void SetUp() override { DialectRegistry registry; - registry.insert(); + registry.insert(); context = std::make_unique(); context->appendDialectRegistry(registry); @@ -50,27 +52,45 @@ class DriversTest : public testing::Test { TEST_F(DriversTest, ProgramWalk) { qco::QCOProgramBuilder builder(context.get()); builder.initialize(); - const auto q00 = builder.allocQubit(); - const auto q10 = builder.allocQubit(); - const auto q20 = builder.allocQubit(); - const auto q30 = builder.allocQubit(); - const auto q01 = builder.h(q00); - const auto [q02, q11] = builder.cx(q01, q10); - const auto [q21, q31] = builder.cx(q20, q30); + Value q0 = builder.allocQubit(); + Value q1 = builder.allocQubit(); + Value q2 = builder.allocQubit(); + Value q3 = builder.allocQubit(); - const auto [q03, c0] = builder.measure(q02); - const auto [q12, c1] = builder.measure(q11); - const auto [q22, c2] = builder.measure(q21); - const auto [q32, c3] = builder.measure(q31); + q0 = builder.h(q0); + std::tie(q0, q1) = builder.cx(q0, q1); + std::tie(q2, q3) = builder.cx(q2, q3); - builder.sink(q03); - builder.sink(q12); - builder.sink(q22); - builder.sink(q32); + const auto forOut = builder.scfFor( + 1, 3, 1, {q0, q1, q2, q3}, [&builder](Value, ValueRange iterArgs) { + return SmallVector{iterArgs[0], iterArgs[1], iterArgs[2], iterArgs[3]}; + }); - auto mod = builder.finalize(); - auto func = *(mod->getOps().begin()); + q0 = forOut[0]; + q1 = forOut[1]; + q2 = forOut[2]; + q3 = forOut[3]; + + Value c0; + Value c1; + Value c2; + Value c3; + + std::tie(q0, c0) = builder.measure(q0); + Operation* firstMeasure = q0.getDefiningOp(); + + std::tie(q1, c1) = builder.measure(q1); + std::tie(q2, c2) = builder.measure(q2); + std::tie(q3, c3) = builder.measure(q3); + + builder.sink(q0); + builder.sink(q1); + builder.sink(q2); + builder.sink(q3); + + auto m = builder.finalize(); + auto func = qco::getEntryPoint(m.get()); Value ex0 = nullptr; Value ex1 = nullptr; @@ -81,9 +101,10 @@ TEST_F(DriversTest, ProgramWalk) { // Since WalkOrder::PreOrder is used here, the state of the qubits is not yet // updated with the SSA values of the measurement op. // Consequently, the program qubits point at the outputs of the controlled-Xs. - std::ignore = qco::walkProgram(func.getBody(), + qco::Qubits qubits; + std::ignore = qco::walkProgram(func.getBody(), qubits, [&](Operation* op, const qco::Qubits& qubits) { - if (op == q03.getDefiningOp()) { + if (op == firstMeasure) { ex0 = qubits.getProgramQubit(0); ex1 = qubits.getProgramQubit(1); ex2 = qubits.getProgramQubit(2); @@ -93,10 +114,10 @@ TEST_F(DriversTest, ProgramWalk) { return WalkResult::advance(); }); - ASSERT_EQ(ex0, q02); - ASSERT_EQ(ex1, q11); - ASSERT_EQ(ex2, q21); - ASSERT_EQ(ex3, q31); + ASSERT_EQ(ex0, forOut[0]); + ASSERT_EQ(ex1, forOut[1]); + ASSERT_EQ(ex2, forOut[2]); + ASSERT_EQ(ex3, forOut[3]); } TEST_F(DriversTest, ProgramGraphWalk) { From 5547a1cff9bee16d253cd5ae60cce101d538c314 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Wed, 3 Jun 2026 08:34:20 +0200 Subject: [PATCH 11/25] WiP: Extend for loop iter_args --- mlir/include/mlir/Dialect/QCO/Utils/Drivers.h | 2 +- mlir/include/mlir/Dialect/QCO/Utils/Qubits.h | 41 +++------- .../QCO/Transforms/Mapping/Mapping.cpp | 76 +++++++++++++++++-- mlir/lib/Dialect/QCO/Utils/Qubits.cpp | 36 +++++++++ .../QCO/Transforms/Mapping/test_mapping.cpp | 7 +- 5 files changed, 120 insertions(+), 42 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h index 5e18fe7e18..a68b7e2699 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h @@ -52,7 +52,7 @@ using WalkProgramFn = function_ref; */ template LogicalResult walkProgram(Region& region, Qubits& qubits, - const WalkProgramFn& fn) { + const WalkProgramFn& fn) { for (Operation& curr : region.getOps()) { if constexpr (Order == WalkOrder::PreOrder) { if (fn(&curr, qubits).wasInterrupted()) { diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Qubits.h b/mlir/include/mlir/Dialect/QCO/Utils/Qubits.h index b08952aa46..ccd9d5a577 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Qubits.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Qubits.h @@ -20,46 +20,27 @@ #include namespace mlir::qco { + class Qubits { - /** - * @brief Specifies the qubit "location" (hardware or program). - */ + /// Specifies the qubit "location" (hardware or program). enum class QubitLocation : std::uint8_t { Hardware, Program }; public: - /** - * @brief Add qubit with automatically assigned dynamic index. - */ + /// Update the qubits map based on the given op. + void update(Operation*); + /// Add qubit with automatically assigned dynamic index. void add(TypedValue q); - - /** - * @brief Add qubit with static index. - */ + /// Add qubit with static index. void add(TypedValue q, std::size_t hw); - - /** - * @brief Remap the qubit value from prev to next. - */ + /// Remap the qubit value from prev to next. void remap(TypedValue prev, TypedValue next); - - /** - * @brief Remove the qubit value. - */ + /// Remove the qubit value. void remove(TypedValue q); - - /** - * @returns the qubit value assigned to a program index. - */ + /// Return the qubit value assigned to a program index. [[nodiscard]] TypedValue getProgramQubit(std::size_t index) const; - - /** - * @returns the qubit value assigned to a hardware index. - */ + /// Return the qubit value assigned to a hardware index. [[nodiscard]] TypedValue getHardwareQubit(std::size_t index) const; - - /** - * @returns the index assigned to the qubit value. - */ + /// Return the index assigned to the qubit value. [[nodiscard]] std::size_t getIndex(TypedValue q) const; private: diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 6fef2ba282..6ce7e6a3a1 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -33,6 +33,8 @@ #include #include #include +#include +#include #include #include #include @@ -462,8 +464,7 @@ struct MappingPass : impl::MappingPassBase { } assert(llvm::all_of(u.wires, - [](auto& it) { return isa(it.operation()); - })); + [](auto& it) { return isa(it.operation()); })); // Collect statistics. numSwaps += *res; @@ -531,8 +532,8 @@ struct MappingPass : impl::MappingPassBase { * @returns a vector of wire iterators, where the i-th wire points at the i-th * static program qubit. */ - static SmallVector - place(func::FuncOp func, const Layout& layout, IRRewriter& rewriter) { + SmallVector place(func::FuncOp func, const Layout& layout, + IRRewriter& rewriter) { SmallVector staticOps; staticOps.reserve(layout.nqubits()); @@ -588,6 +589,61 @@ struct MappingPass : impl::MappingPassBase { SinkOp::create(rewriter, func->getLoc(), qubit); } + for (auto forOp : make_early_inc_range(func.getOps())) { + rewriter.setInsertionPoint(forOp); + + const auto& domInfo = getAnalysis(); + + // TODO: Idea use the IRMapping to map old to new, use nxt and "tied" + // method to find block arguments. use below and replace. Use then results + // of for loop for wires in next iteration. + + SmallVector newInitArgs(placedWires.size()); + SmallVector wires{placedWires}; + + IRMapping m; + for (const auto& [i, it] : enumerate(wires)) { + for (; it != std::default_sentinel; ++it) { + const auto nxt = std::next(it); + if (nxt.operation() == forOp || + !domInfo.dominates(nxt.operation(), forOp)) { + newInitArgs[i] = it.qubit(); + break; + } + } + } + + for (Value q : newInitArgs) { + llvm::dbgs() << q << '\n'; + } + + // Replace old for with a new one with updated init arguments. + // TODO: argValues in mergeBlocks is probably not right. This needs some + // additional map built above to map the correct old with the correct new + // value. + const size_t nargs = forOp.getBody()->getNumArguments(); + + auto newForOp = rewriter.create( + forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), + forOp.getStep(), newInitArgs); + rewriter.mergeBlocks( + forOp.getBody(), newForOp.getBody(), + newForOp.getBody()->getArguments().take_front(nargs)); + auto newYieldOp = cast(newForOp.getBody()->getTerminator()); + + // Once the new for loop has been created, we can use it to update the + // yield operation inside. Particularly, we extend the yielded results + // with the newly added, unused block arguments. + + SmallVector newResults(newYieldOp.getResults()); + for (BlockArgument& arg : newForOp.getRegionIterArgs()) { + if (arg.use_empty()) { + newResults.emplace_back(arg); + } + } + rewriter.replaceOpWithNewOp(newYieldOp, newResults); + } + return placedWires; } @@ -1001,6 +1057,10 @@ struct MappingPass : impl::MappingPassBase { while (true) { auto [op, subUnits] = advanceFrame(unit); if (op != nullptr) { + assert(isa(op)); + assert( + all_of(unit.wires, [&](auto& it) { return it.operation() == op; })); + for (auto& subUnit : subUnits) { const auto res = route(subUnit, rewriter); if (failed(res)) { @@ -1009,13 +1069,11 @@ struct MappingPass : impl::MappingPassBase { nswaps += *res; } - assert(isa(op)); - // An scf.for operation has exactly one sub-unit. // To finalize the mapping of the sub-unit, we append an appendix // of SWAPs to restore the base-layout. - auto& subUnit = subUnits.front(); + Unit& subUnit = subUnits.front(); const auto swaps = restore(subUnit.layout, unit.layout); if constexpr (Mode == RoutingMode::Hot) { @@ -1025,6 +1083,10 @@ struct MappingPass : impl::MappingPassBase { // the wire iterators are not required any more, there is no need to // increment them again after applying the SWAP sequence. + assert(all_of(subUnit.wires, [](auto& it) { + return isa(it.operation()); + })); + for_each(subUnit.wires, [](auto& it) { std::ranges::advance(it, -Traits::stride()); }); diff --git a/mlir/lib/Dialect/QCO/Utils/Qubits.cpp b/mlir/lib/Dialect/QCO/Utils/Qubits.cpp index c01187f625..12787b0edc 100644 --- a/mlir/lib/Dialect/QCO/Utils/Qubits.cpp +++ b/mlir/lib/Dialect/QCO/Utils/Qubits.cpp @@ -11,7 +11,12 @@ #include "mlir/Dialect/QCO/Utils/Qubits.h" #include "mlir/Dialect/QCO/IR/QCODialect.h" +#include "mlir/Dialect/QCO/IR/QCOInterfaces.h" +#include "mlir/Dialect/QCO/IR/QCOOps.h" +#include +#include +#include #include #include @@ -19,6 +24,37 @@ #include namespace mlir::qco { +void Qubits::update(Operation* op) { + TypeSwitch(op) + .Case([&](StaticOp op) { add(op.getQubit(), op.getIndex()); }) + .Case([&](AllocOp op) { add(op.getResult()); }) + .Case([&](UnitaryOpInterface& op) { + for (const auto& [prevV, nextV] : + llvm::zip(op.getInputQubits(), op.getOutputQubits())) { + const auto prevQ = cast>(prevV); + const auto nextQ = cast>(nextV); + remap(prevQ, nextQ); + } + }) + .Case([&](scf::ForOp op) { + for (OpOperand& operand : op.getInitsMutable()) { + auto result = op.getTiedLoopResult(&operand); + const auto prevQ = cast>(operand.get()); + const auto nextQ = cast>(result); + remap(prevQ, nextQ); + } + }) + .Case( + [&](ResetOp op) { remap(op.getQubitIn(), op.getQubitOut()); }) + .Case( + [&](MeasureOp op) { remap(op.getQubitIn(), op.getQubitOut()); }) + .Case([&](SinkOp op) { remove(op.getQubit()); }) + .Default([&](Operation* op) -> WalkResult { + const auto name = op->getName().getStringRef(); + reportFatalInternalError("unexpected op" + name); + }); +} + void Qubits::add(TypedValue q) { const auto index = programToValue_.size(); programToValue_.try_emplace(index, q); diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 5b2a9fbc45..f6284443d1 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -44,7 +44,7 @@ using namespace mlir::qco; struct Device { size_t nqubits{}; - llvm::DenseSet> couplingSet; + DenseSet> couplingSet; }; /** @@ -53,7 +53,7 @@ struct Device { */ static LogicalResult isExecutable(Region& region, - const llvm::DenseSet>& couplingSet) { + const DenseSet>& couplingSet) { Qubits qubits; return walkProgram( region, qubits, [&](Operation* curr, const Qubits& qubits) { @@ -107,8 +107,7 @@ class MappingPassTest : public testing::Test, } static LogicalResult - runPass(ModuleOp m, - const llvm::DenseSet>& couplingSet, + runPass(ModuleOp m, const DenseSet>& couplingSet, const MappingPassOptions& options) { PassManager pm(m->getContext()); pm.addPass(createMappingPass(couplingSet, options)); From 506101604a975885aea88c1b4ef7d9a3496696c2 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 19 Jun 2026 10:09:22 +0200 Subject: [PATCH 12/25] Implement working prototype --- mlir/include/mlir/Dialect/QCO/Utils/Drivers.h | 11 +- .../QCO/Transforms/Mapping/Mapping.cpp | 647 +++++++++--------- 2 files changed, 339 insertions(+), 319 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h index a68b7e2699..e8663567b0 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h @@ -25,6 +25,7 @@ #include #include +#include #include #include #include @@ -52,7 +53,7 @@ using WalkProgramFn = function_ref; */ template LogicalResult walkProgram(Region& region, Qubits& qubits, - const WalkProgramFn& fn) { + const WalkProgramFn& fn) { for (Operation& curr : region.getOps()) { if constexpr (Order == WalkOrder::PreOrder) { if (fn(&curr, qubits).wasInterrupted()) { @@ -167,6 +168,14 @@ LogicalResult walkProgramGraph(MutableArrayRef wires, for (size_t i : curr) { auto& it = wires[i]; while (Traits::isActive(it)) { + + if (it.qubit() != nullptr && isa(it.qubit())) { + std::ranges::advance(it, Traits::stride()); + continue; + } + + assert(it.operation() != nullptr); + const auto res = TypeSwitch(it.operation()) .template Case([&](auto& op) { diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 6ce7e6a3a1..1ea807e404 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -39,6 +40,7 @@ #include #include #include +#include #include #include #include @@ -73,7 +75,6 @@ struct MappingPass : impl::MappingPassBase { using IndexType = size_t; using IndexPairType = std::pair; using Window = SmallVector; - using Neighbours = SmallVector>; enum class RoutingMode : std::uint8_t { Cold, Hot }; @@ -281,12 +282,9 @@ struct MappingPass : impl::MappingPassBase { Matrix dist_; }; - struct [[nodiscard]] Trial { - explicit Trial(Layout layout) : layout(std::move(layout)) {} - - Layout layout; - size_t nswaps{}; - bool success{false}; + /// Statistics collected while routing. + struct Statistics { + size_t nswaps{0}; }; /** @@ -378,12 +376,6 @@ struct MappingPass : impl::MappingPassBase { } }; - /// A unit is a (to be) routeable region. - struct Unit { - SmallVector wires; - Layout layout; - }; - public: /// Construct default mapping pass. MappingPass() = default; @@ -401,80 +393,107 @@ struct MappingPass : impl::MappingPassBase { assert(niterations > 0 && "runOnOperation: expected niterations > 0"); assert(ntrials > 0 && "runOnOperation: expected ntrials > 0"); - std::mt19937_64 rng{seed}; IRRewriter rewriter(&getContext()); - ModuleOp m = getOperation(); - auto func = getEntryPoint(m); + auto mod = getOperation(); + auto func = getEntryPoint(mod); if (!func) { - m.emitError() << "does not contain an entry point function"; + mod.emitError() << "does not contain an entry point function"; signalPassFailure(); return; } - auto comp = getComputation(func); - if (failed(comp)) { + auto& body = func.getFunctionBody(); + + auto wires = getComputation(func); + if (failed(wires)) { signalPassFailure(); return; } - if (comp->size() > device.nqubits()) { - m.emitError() << "requires " + Twine(comp.value().size()) + - " qubits. However, the architecture only supports " + - Twine(device.nqubits()) + "qubits."; + if (wires->size() > device.nqubits()) { + func.emitError() + << "requires " + Twine(wires.value().size()) + + " qubits. However, the architecture only supports " + + Twine(device.nqubits()) + "qubits."; signalPassFailure(); return; } - // Create trials for initial layout refining. Currently, this includes - // `ntrials` many random layouts. - // SmallVector trials; - // trials.reserve(ntrials); - // for (size_t i = 0; i < ntrials; ++i) { - // trials.emplace_back(Layout::random(device.nqubits(), rng())); - // } - - // Execute each of the trials (possibly in parallel). Collect the results - // and find the one with the fewest SWAPs on the final backwards pass. - // parallelForEach(&getContext(), trials, [&, this](Trial& trial) { - // if (const auto res = refineLayout(*comp, trial.layout); succeeded(res)) - // { - // trial.success = true; - // trial.nswaps = *res; - // } - // }); - - // Trial* best = findBestTrial(trials); - // if (best == nullptr) { - // func.emitError() << "failed to find the best layout trial"; - // signalPassFailure(); - // return; - // } - - Layout l = Layout::identity(device.nqubits()); - // Perform placement and hot routing by inserting SWAPs into the IR. - auto placedWires = place(func, l, rewriter); - Unit u{.wires = placedWires, .layout = l}; - const auto res = - route(u, &rewriter); - if (failed(res)) { - func.emitError() << "failed to map the " << func.getName() << " function"; + auto layout = generateLayout(*wires); + if (failed(layout)) { + func->emitError() << "failed to refine random initial layouts."; + signalPassFailure(); + } + + wires = std::move(place(body, *layout, rewriter)); + + Statistics s; + + const auto res = route( + *wires, *layout, s, &rewriter); + if (res.failed()) { + func.emitError() << "failed to map the function"; signalPassFailure(); return; } - assert(llvm::all_of(u.wires, - [](auto& it) { return isa(it.operation()); })); + func->dumpPretty(); // Collect statistics. - numSwaps += *res; + numSwaps += s.nswaps; // Fix SSA Dominance issues. - for_each(func.getFunctionBody().getBlocks(), - [](Block& b) { sortTopologically(&b); }); + llvm::for_each(func.getFunctionBody().getBlocks(), + [](Block& b) { sortTopologically(&b); }); } private: + static scf::ForOp extend(scf::ForOp loop, ValueRange addons, + IRRewriter& rewriter) { + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(loop); + + const auto naddons = addons.size(); + + SmallVector inits; + llvm::append_range(inits, loop.getInits()); + llvm::append_range(inits, addons); + + auto newLoop = rewriter.create( + loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(), + loop.getStep(), inits); + + Block* loopBody = loop.getBody(); + Block* newLoopBody = newLoop.getBody(); + + rewriter.mergeBlocks( + loopBody, newLoopBody, + newLoopBody->getArguments().take_front(loopBody->getNumArguments())); + + for (const auto [before, after] : + llvm::zip_first(loop.getResults(), newLoop.getResults())) { + rewriter.replaceAllUsesWith(before, after); + } + + for (const auto [before, after] : + llvm::zip_equal(addons, newLoop.getResults().take_back(naddons))) { + rewriter.replaceAllUsesExcept(before, after, newLoop); + } + + auto yield = cast(newLoopBody->getTerminator()); + + SmallVector results; + llvm::append_range(results, yield.getResults()); + llvm::append_range(results, newLoop.getRegionIterArgs().take_back(naddons)); + rewriter.setInsertionPoint(yield); + rewriter.replaceOpWithNewOp(yield, results); + + rewriter.eraseOp(loop); + + return newLoop; + } + /** * @brief Collect wires of the quantum computation before placement. * @details @@ -496,8 +515,7 @@ struct MappingPass : impl::MappingPassBase { static FailureOr> getComputation(func::FuncOp func) { if (!func.getOps().empty()) { - func.emitError() << "must not contain qco.alloc operations"; - return failure(); + return func.emitError() << "must not contain qco.alloc operations"; } SmallVector wires; @@ -507,8 +525,8 @@ struct MappingPass : impl::MappingPassBase { for (; it != std::default_sentinel; ++it) { if (auto extract = dyn_cast(it.operation())) { if (!isInitPhase) { - func.emitError() << "must extract and insert all qubits at once."; - return failure(); + return func.emitError() + << "must extract and insert all qubits at once."; } wires.emplace_back(extract.getResult()); continue; @@ -532,24 +550,24 @@ struct MappingPass : impl::MappingPassBase { * @returns a vector of wire iterators, where the i-th wire points at the i-th * static program qubit. */ - SmallVector place(func::FuncOp func, const Layout& layout, - IRRewriter& rewriter) { + static SmallVector place(Region& body, const Layout& layout, + IRRewriter& rewriter) { SmallVector staticOps; staticOps.reserve(layout.nqubits()); // Create and save static qubit operations. - rewriter.setInsertionPointToStart(&func.getFunctionBody().front()); + rewriter.setInsertionPointToStart(&body.front()); for (size_t i = 0; i < layout.nqubits(); ++i) { - const auto op = StaticOp::create(rewriter, func.getLoc(), i); + const auto op = StaticOp::create(rewriter, body.getLoc(), i); staticOps.emplace_back(op); rewriter.setInsertionPointAfter(op); } // Replace extract ops and collect in program-qubit order. - SmallVector placedWires(layout.nqubits()); + SmallVector wires(layout.nqubits()); size_t prog = 0UL; - for (auto alloc : make_early_inc_range(func.getOps())) { + for (auto alloc : make_early_inc_range(body.getOps())) { TensorIterator it(alloc.getResult()); while (it != std::default_sentinel) { // Get the operation and early increment to avoid issues after erasure. @@ -565,7 +583,7 @@ struct MappingPass : impl::MappingPassBase { rewriter.replaceAllUsesWith(op.getOutTensor(), op.getTensor()); rewriter.eraseOp(op); - placedWires[prog] = WireIterator(qubit); + wires[prog] = WireIterator(qubit); ++prog; }) .Case([&](auto op) { @@ -581,114 +599,113 @@ struct MappingPass : impl::MappingPassBase { } // Create sinks for remaining, unused, static qubits. - rewriter.setInsertionPoint(func.getFunctionBody().back().getTerminator()); + rewriter.setInsertionPoint(body.back().getTerminator()); for (; prog < layout.nqubits(); ++prog) { const auto hw = layout.getHardwareIndex(prog); const auto qubit = staticOps[hw].getQubit(); - placedWires[prog] = WireIterator(qubit); - SinkOp::create(rewriter, func->getLoc(), qubit); + wires[prog] = WireIterator(qubit); + SinkOp::create(rewriter, body.getLoc(), qubit); } - for (auto forOp : make_early_inc_range(func.getOps())) { - rewriter.setInsertionPoint(forOp); - - const auto& domInfo = getAnalysis(); - - // TODO: Idea use the IRMapping to map old to new, use nxt and "tied" - // method to find block arguments. use below and replace. Use then results - // of for loop for wires in next iteration. - - SmallVector newInitArgs(placedWires.size()); - SmallVector wires{placedWires}; - - IRMapping m; - for (const auto& [i, it] : enumerate(wires)) { - for (; it != std::default_sentinel; ++it) { - const auto nxt = std::next(it); - if (nxt.operation() == forOp || - !domInfo.dominates(nxt.operation(), forOp)) { - newInitArgs[i] = it.qubit(); - break; - } - } - } + // Finally, update the SCF operations such that they take all static qubits + // as input. To handle recursively nested SCF operations, use a stack of + // (region, mapping) pairs. - for (Value q : newInitArgs) { - llvm::dbgs() << q << '\n'; - } + SmallVector>> stack; + stack.emplace_back(body, DenseSet{}); - // Replace old for with a new one with updated init arguments. - // TODO: argValues in mergeBlocks is probably not right. This needs some - // additional map built above to map the correct old with the correct new - // value. - const size_t nargs = forOp.getBody()->getNumArguments(); - - auto newForOp = rewriter.create( - forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), - forOp.getStep(), newInitArgs); - rewriter.mergeBlocks( - forOp.getBody(), newForOp.getBody(), - newForOp.getBody()->getArguments().take_front(nargs)); - auto newYieldOp = cast(newForOp.getBody()->getTerminator()); - - // Once the new for loop has been created, we can use it to update the - // yield operation inside. Particularly, we extend the yielded results - // with the newly added, unused block arguments. - - SmallVector newResults(newYieldOp.getResults()); - for (BlockArgument& arg : newForOp.getRegionIterArgs()) { - if (arg.use_empty()) { - newResults.emplace_back(arg); - } - } - rewriter.replaceOpWithNewOp(newYieldOp, newResults); - } + while (!stack.empty()) { + auto [region, m] = stack.pop_back_val(); - return placedWires; - } + for (Operation& op : llvm::make_early_inc_range(region.getOps())) { + TypeSwitch(&op) + .Case([&](StaticOp op) { m.insert(op.getQubit()); }) + .Case([&](UnitaryOpInterface& op) { + for (const auto [pred, succ] : + llvm::zip_equal(op.getInputQubits(), op.getOutputQubits())) { + m.insert(succ); + m.erase(pred); + } + }) + .Case([&](scf::ForOp loop) { + assert(m.size() == layout.nqubits()); - /** - * @brief Find the best trial result in terms of the number of SWAPs. - * @returns the best trial result or nullptr if no result is valid. - */ - static Trial* findBestTrial(MutableArrayRef trials) { - Trial* best = nullptr; - for (auto& trial : trials) { - if (!trial.success) { - continue; - } + DenseSet addons(m); + llvm::for_each(loop.getInits(), [&](auto v) { addons.erase(v); }); + auto newLoop = extend(loop, to_vector(addons), rewriter); + + for (OpOperand& operand : newLoop.getInitsMutable()) { + m.insert(newLoop.getTiedLoopResult(&operand)); + m.erase(operand.get()); + } - if (best == nullptr || best->nswaps > trial.nswaps) { - best = &trial; + stack.emplace_back( + newLoop.getRegion(), + DenseSet(newLoop.getRegionIterArgs().begin(), + newLoop.getRegionIterArgs().end())); + }) + .Case([&](auto op) { + m.insert(op.getQubitOut()); + m.erase(op.getQubitIn()); + }) + .Case([&](auto) { + llvm::reportFatalInternalError("unexpected dynamic qubit alloc"); + }); } } - return best; + return wires; } - /** - * @brief Refine the trial's layout and count #swaps for the final backwards - * pass. - * @details Use the SABRE Approach to improve the initial layout: - * Traverse the layers from left-to-right-to-left and cold-route - * along the way. Repeat this procedure "niterations" times. - * @returns failure() if routing fails. - */ - FailureOr refineLayout(Unit& u) { - size_t nswaps{0}; - for (size_t i = 0; i < niterations; ++i) { - if (failed(route(u))) { - return failure(); + /// Execute `ntrials` many (parallel) initial layout refinement trials and + /// return the heuristically best one. + /// + /// The function uses the SABRE Approach to improve the initial layout: + /// Traverse the layers of the program from left-to-right-to-left and + /// cold-route along the way. Repeat this procedure "niterations" times and + /// finally find the trial with the fewest SWAPs on the final backwards pass + /// and return the respective layout. + FailureOr generateLayout(ArrayRef wires) { + std::mt19937_64 rng{seed}; + + struct Trial { + Layout layout; + Statistics stats{}; + bool success{false}; + }; + + SmallVector trials; + trials.reserve(ntrials); + for (size_t i = 0; i < ntrials; ++i) { + trials.emplace_back(Layout::random(device.nqubits(), rng())); + } + + parallelForEach(&getContext(), trials, [&, this](Trial& t) { + SmallVector local(wires); + for (size_t i = 0; i < niterations; ++i) { + if (route(local, t.layout, t.stats).failed()) { + return; + } + if (route(local, t.layout, t.stats).failed()) { + return; + } } + t.success = true; + }); - const auto resB = route(u); - if (failed(resB)) { - return failure(); + Trial* best = nullptr; + for (Trial& t : trials) { + if (t.success && + (best == nullptr || best->stats.nswaps > t.stats.nswaps)) { + best = &t; } - nswaps = *resB; } - return nswaps; + if (best == nullptr) { + return failure(); + } + + return best->layout; } /** @@ -754,8 +771,8 @@ struct MappingPass : impl::MappingPassBase { it->second = curr->depth; } - // If the currently visited node is a goal node, reconstruct the sequence - // of SWAPs from this node to the root. + // If the currently visited node is a goal node, reconstruct the + // sequence of SWAPs from this node to the root. if (curr->isGoal(window.front(), device)) { SmallVector seq(curr->depth); @@ -826,6 +843,10 @@ struct MappingPass : impl::MappingPassBase { return; } + if (curr0.operation() == nullptr) { + return; + } + // Handle two-qubit barrier edge case explicitly. if (auto barrier = dyn_cast(curr0.operation())) { if (barrier.getNumQubits() != 2) { @@ -841,106 +862,44 @@ struct MappingPass : impl::MappingPassBase { /** * @brief Build and return window of layers. * @details Traverses the circuit-layers until the desired window sizes is - * reached. Assumes that wires[i] = i-th program qubit. The size of the window - * is 1 + nlookahead. + * reached. Assumes that wires[i] = i-th program qubit. The size of the + * window is 1 + nlookahead. * @returns window of layers. */ template - Window getWindow(ArrayRef baseWires) { + Window getWindow(ArrayRef wires) { Window window; window.reserve(1 + nlookahead); - SmallVector wires(baseWires); - std::ignore = walkProgramGraph( - wires, [&](const ReadyRange& ready, ReleasedOps& released) { - if (ready.empty()) { - return WalkResult::advance(); - } - - for (const auto& [op, progs] : ready) { - if (isa(op)) { - released.emplace_back(op); - continue; - } - - assert(isa(op)); - - const auto p0 = progs[0]; - const auto p1 = progs[1]; - window.emplace_back(p0, p1); - if (window.size() == 1 + nlookahead) { - return WalkResult::interrupt(); - } - - skipQubitPairBlock(wires[p0], wires[p1]); - released.emplace_back(wires[p0].operation()); - } - - return WalkResult::advance(); - }); - - return window; - } - - /** - * @todo UPDATE THIS DESCRIPTION - * @brief Advance past all executable gates. - * @details Traverses the multi-qubit gates of the circuit until no more - * executable gates are found. - */ - template - std::pair> advanceFrame(Unit& unit) { - Operation* subOp = nullptr; - SmallVector subs; - - std::ignore = walkProgramGraph( - unit.wires, [&](const ReadyRange& ready, ReleasedOps& released) { + SmallVector local(wires); + walkProgramGraph( + local, [&](const ReadyRange& ready, ReleasedOps& released) { if (ready.empty()) { return WalkResult::advance(); } for (const auto& [op, progs] : ready) { - /// TODO: TypeSwitch - if (isa(op)) { - released.emplace_back(op); - continue; - } - - if (auto forOp = dyn_cast(op)) { - SmallVector rWires; - for (auto prog : progs) { - const auto res = cast(unit.wires[prog].qubit()); - auto& it = - rWires.emplace_back(forOp.getTiedLoopRegionIterArg(res)); - // The iterator points at a block argument. Move the iterator to - // a position where it.operation() != nullptr. - std::ranges::advance(it, - WireTraversalTraits::stride()); + if (auto u = dyn_cast(op)) { + const auto p0 = progs[0]; + const auto p1 = progs[1]; + window.emplace_back(p0, p1); + if (window.size() == 1 + nlookahead) { + return WalkResult::interrupt(); } - subOp = op; - subs.emplace_back(rWires, unit.layout); - return WalkResult::interrupt(); - } - - assert(isa(op)); - const auto [hw0, hw1] = - unit.layout.getHardwareIndices(progs[0], progs[1]); - - if (device.areAdjacent(hw0, hw1)) { - released.emplace_back(op); + skipQubitPairBlock(local[p0], local[p1]); + released.emplace_back(u); + return WalkResult::advance(); } - } - // Stop, if there are no more ready AND executable gates. - if (released.empty()) { - return WalkResult::interrupt(); + released.emplace_back(op); + return WalkResult::advance(); } return WalkResult::advance(); }); - return std::make_pair(subOp, subs); + return window; } SmallVector restore(const Layout& from, const Layout& to) { @@ -1000,22 +959,73 @@ struct MappingPass : impl::MappingPassBase { return swaps; } - template - void applySWAPs(Unit& unit, ArrayRef swaps, - IRRewriter* rewriter) { + /** + * @brief Advance past all executable gates and recurse into nested regions, + * if necessary. + * @details Traverses the multi-qubit gates of the circuit until no more + * executable gates are found. + */ + template + void advanceAndRecurse(MutableArrayRef wires, Layout& layout, + Statistics& stats, IRRewriter* rewriter) { + walkProgramGraph(wires, [&](const ReadyRange& ready, + ReleasedOps& released) { + if (ready.empty()) { + return WalkResult::advance(); + } + + for (const auto& [readyOp, progs] : ready) { + TypeSwitch(readyOp) + .Case([&](UnitaryOpInterface op) { + const auto [hw0, hw1] = + layout.getHardwareIndices(progs[0], progs[1]); + + if (device.areAdjacent(hw0, hw1)) { + released.emplace_back(op); + } + }) + .template Case( + [&](BarrierOp op) { released.emplace_back(op); }) + .template Case([&](scf::ForOp op) { + // TODO: Don't ignore result here. + std::ignore = + routeForLoop(op, layout, stats, rewriter); + + released.emplace_back(op); + }); + } + + // Stop, if there are no more ready AND executable gates. + if (released.empty()) { + return WalkResult::interrupt(); + } + + return WalkResult::advance(); + }); + } + + /// Insert SWAP operations, exchanging two qubits, virtually (Mode::Cold) or + /// into the IR (Mode::Hot). The function expects that the i-th wire points + /// at the i-th program qubit and that each wire also points at the correct + /// insertion point. The function preserves the program qubit order, by + /// exchanging wires after a SWAP. + template + static void insertSWAPs(ArrayRef swaps, + MutableArrayRef wires, Layout& layout, + IRRewriter* rewriter) { for (const auto& [hw0, hw1] : swaps) { if constexpr (Mode == RoutingMode::Hot) { - const auto& [prog0, prog1] = unit.layout.getProgramIndices(hw0, hw1); - const auto& w0 = unit.wires[prog0]; - const auto& w1 = unit.wires[prog1]; - - assert(!isa(w0.operation())); - assert(!isa(w1.operation())); + const auto& [prog0, prog1] = layout.getProgramIndices(hw0, hw1); + const auto& w0 = wires[prog0]; + const auto& w1 = wires[prog1]; const auto in0 = w0.qubit(); const auto in1 = w1.qubit(); - rewriter->setInsertionPointAfter(in0.getDefiningOp()); + // Hot routing is only called on the forward direction. Hence, + // setInsertionPoint*After*Value is valid here. + + rewriter->setInsertionPointAfterValue(in0); auto swapOp = SWAPOp::create(*rewriter, in0.getLoc(), in0, in1); const auto out0 = swapOp.getQubit0Out(); @@ -1025,88 +1035,39 @@ struct MappingPass : impl::MappingPassBase { rewriter->replaceAllUsesExcept(in1, out0, swapOp); // Preserve program-indexed wire semantics. - unit.wires[prog0] = WireIterator(out1); - unit.wires[prog1] = WireIterator(out0); + wires[prog0] = WireIterator(out1); + wires[prog1] = WireIterator(out0); assert(isa(w0.operation())); assert(isa(w1.operation())); } - unit.layout.swap(hw0, hw1); + layout.swap(hw0, hw1); } } /** * @brief Route via SWAP insertion. - * @details Iterates over a dynamically computed window of layers and uses A* - * search to find a sequence of SWAPs that makes that layer executable. + * @details Iterates over a dynamically computed window of layers and uses + * A* search to find a sequence of SWAPs that makes that layer executable. * Depending on the template parameter, this function only updates * (and hence modifies) the layout or also inserts the SWAPs into the IR. - * @returns failure() if A* search isn't able to find a solution, the number - * of SWAPs otherwise. + * @returns failure() if A* search isn't able to find a solution. */ template - FailureOr route(Unit& unit, IRRewriter* rewriter = nullptr) { - assert(Mode == RoutingMode::Cold || - (Mode == RoutingMode::Hot && Direction == WireDirection::Forward)); - + LogicalResult route(SmallVector& wires, Layout& layout, + Statistics& stats, IRRewriter* rewriter = nullptr) { using Traits = WireTraversalTraits; - size_t nswaps{0}; - while (true) { - auto [op, subUnits] = advanceFrame(unit); - if (op != nullptr) { - assert(isa(op)); - assert( - all_of(unit.wires, [&](auto& it) { return it.operation() == op; })); - - for (auto& subUnit : subUnits) { - const auto res = route(subUnit, rewriter); - if (failed(res)) { - return failure(); - } - nswaps += *res; - } - - // An scf.for operation has exactly one sub-unit. - // To finalize the mapping of the sub-unit, we append an appendix - // of SWAPs to restore the base-layout. + advanceAndRecurse(wires, layout, stats, rewriter); - Unit& subUnit = subUnits.front(); - const auto swaps = restore(subUnit.layout, unit.layout); - - if constexpr (Mode == RoutingMode::Hot) { - - // At this point the wire iterators point to scf.yield ops. Thus, - // decrement the iterator to point at a valid insertion point. Because - // the wire iterators are not required any more, there is no need to - // increment them again after applying the SWAP sequence. - - assert(all_of(subUnit.wires, [](auto& it) { - return isa(it.operation()); - })); - - for_each(subUnit.wires, [](auto& it) { - std::ranges::advance(it, -Traits::stride()); - }); - } - - applySWAPs(subUnit, swaps, rewriter); - nswaps += swaps.size(); - } - - const auto window = getWindow(unit.wires); + const auto window = getWindow(wires); if (window.empty()) { break; } - for (const auto [i0, i1] : window) { - llvm::dbgs() << "(" << i0 << ", " << i1 << ") "; - } - llvm::dbgs() << '\n'; - - const auto swaps = search(window, unit.layout); + const auto swaps = search(window, layout); if (failed(swaps)) { return failure(); } @@ -1120,14 +1081,14 @@ struct MappingPass : impl::MappingPassBase { // simply must ensure the insertion point is before the multi-qubit // gates. - for (auto& it : unit.wires) { + for (auto& it : wires) { std::ranges::advance(it, it == std::default_sentinel ? -2 * Traits::stride() : -Traits::stride()); } } - applySWAPs(unit, *swaps, rewriter); + insertSWAPs(*swaps, wires, layout, rewriter); if constexpr (Mode == RoutingMode::Hot) { @@ -1139,14 +1100,64 @@ struct MappingPass : impl::MappingPassBase { // multi-qubit op of the current or subsequent layer or to a sink (and // thus std::default_sentinel). - for_each(unit.wires, - [](auto& it) { std::ranges::advance(it, Traits::stride()); }); + llvm::for_each(wires, [](auto& it) { + std::ranges::advance(it, Traits::stride()); + }); } - nswaps += swaps->size(); + stats.nswaps += swaps->size(); + } + + return success(); + } + + template + LogicalResult routeForLoop(scf::ForOp op, Layout& base, Statistics& stats, + IRRewriter* rewriter) { + using Traits = WireTraversalTraits; + + assert(llvm::all_of(op.getInitArgs(), + [](Value v) { return isa(v.getType()); })); + + // In the forward direction we start the block arguments of the loop body, + // whereas in the backward direction we start at the yielded values. + + SmallVector wires; + if constexpr (Direction == WireDirection::Forward) { + llvm::for_each(op.getRegionIterArgs(), + [&](Value v) { wires.emplace_back(v); }); + } else { + auto yield = cast(op.getBody()->getTerminator()); + assert(yield != nullptr); + llvm::for_each(yield.getResults(), + [&](Value v) { wires.emplace_back(v); }); + } + + Layout remote(base); + if (route(wires, remote, stats, rewriter).failed()) { + return failure(); + } + + const auto swaps = restore(remote, base); + + if constexpr (Mode == RoutingMode::Hot) { + assert(llvm::all_of(wires, [](const WireIterator& it) { + return it == std::default_sentinel; + })); + + llvm::for_each(wires, [](auto& it) { + std::ranges::advance(it, -2 * Traits::stride()); + }); + } + + insertSWAPs(swaps, wires, remote, rewriter); + stats.nswaps += swaps.size(); + + if constexpr (Mode == RoutingMode::Hot) { + sortTopologically(op.getBody()); } - return nswaps; + return success(); } AugmentedDevice device; From 4d014015240ea580878dab3dd4ec7084692a5564 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 19 Jun 2026 10:32:01 +0200 Subject: [PATCH 13/25] Minor cleanup --- .../QCO/Transforms/Mapping/Mapping.cpp | 38 ++++++------------- 1 file changed, 12 insertions(+), 26 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 1ea807e404..1cb58e0360 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -988,8 +988,7 @@ struct MappingPass : impl::MappingPassBase { [&](BarrierOp op) { released.emplace_back(op); }) .template Case([&](scf::ForOp op) { // TODO: Don't ignore result here. - std::ignore = - routeForLoop(op, layout, stats, rewriter); + std::ignore = route(op, layout, stats, rewriter); released.emplace_back(op); }); @@ -1022,10 +1021,7 @@ struct MappingPass : impl::MappingPassBase { const auto in0 = w0.qubit(); const auto in1 = w1.qubit(); - // Hot routing is only called on the forward direction. Hence, - // setInsertionPoint*After*Value is valid here. - - rewriter->setInsertionPointAfterValue(in0); + rewriter->setInsertionPointAfterValue(in0); // Valid bc. Hot => Forward. auto swapOp = SWAPOp::create(*rewriter, in0.getLoc(), in0, in1); const auto out0 = swapOp.getQubit0Out(); @@ -1055,10 +1051,9 @@ struct MappingPass : impl::MappingPassBase { * @returns failure() if A* search isn't able to find a solution. */ template + requires(Mode != RoutingMode::Hot || Direction == WireDirection::Forward) LogicalResult route(SmallVector& wires, Layout& layout, Statistics& stats, IRRewriter* rewriter = nullptr) { - using Traits = WireTraversalTraits; - while (true) { advanceAndRecurse(wires, layout, stats, rewriter); @@ -1082,9 +1077,7 @@ struct MappingPass : impl::MappingPassBase { // gates. for (auto& it : wires) { - std::ranges::advance(it, it == std::default_sentinel - ? -2 * Traits::stride() - : -Traits::stride()); + std::ranges::advance(it, it == std::default_sentinel ? -2 : -1); } } @@ -1100,9 +1093,7 @@ struct MappingPass : impl::MappingPassBase { // multi-qubit op of the current or subsequent layer or to a sink (and // thus std::default_sentinel). - llvm::for_each(wires, [](auto& it) { - std::ranges::advance(it, Traits::stride()); - }); + llvm::for_each(wires, [](auto& it) { std::ranges::advance(it, 1); }); } stats.nswaps += swaps->size(); @@ -1112,9 +1103,9 @@ struct MappingPass : impl::MappingPassBase { } template - LogicalResult routeForLoop(scf::ForOp op, Layout& base, Statistics& stats, - IRRewriter* rewriter) { - using Traits = WireTraversalTraits; + requires(Mode != RoutingMode::Hot || Direction == WireDirection::Forward) + LogicalResult route(scf::ForOp op, Layout& base, Statistics& stats, + IRRewriter* rewriter) { assert(llvm::all_of(op.getInitArgs(), [](Value v) { return isa(v.getType()); })); @@ -1138,18 +1129,13 @@ struct MappingPass : impl::MappingPassBase { return failure(); } - const auto swaps = restore(remote, base); - if constexpr (Mode == RoutingMode::Hot) { - assert(llvm::all_of(wires, [](const WireIterator& it) { - return it == std::default_sentinel; - })); - - llvm::for_each(wires, [](auto& it) { - std::ranges::advance(it, -2 * Traits::stride()); - }); + assert(llvm::all_of( + wires, [](const auto& it) { return it == std::default_sentinel; })); + llvm::for_each(wires, [](auto& it) { std::ranges::advance(it, -2); }); } + const auto swaps = restore(remote, base); insertSWAPs(swaps, wires, remote, rewriter); stats.nswaps += swaps.size(); From f0931f3ae27529cd2f50ca65d8fa6ba1f915a315 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 19 Jun 2026 10:56:17 +0200 Subject: [PATCH 14/25] Extract layout to seperate module --- mlir/include/mlir/Dialect/QCO/Utils/Layout.h | 68 ++++++++ .../QCO/Transforms/Mapping/Mapping.cpp | 153 +----------------- mlir/lib/Dialect/QCO/Utils/Layout.cpp | 54 +++++++ 3 files changed, 127 insertions(+), 148 deletions(-) create mode 100644 mlir/include/mlir/Dialect/QCO/Utils/Layout.h create mode 100644 mlir/lib/Dialect/QCO/Utils/Layout.cpp diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Layout.h b/mlir/include/mlir/Dialect/QCO/Utils/Layout.h new file mode 100644 index 0000000000..ca52076eef --- /dev/null +++ b/mlir/include/mlir/Dialect/QCO/Utils/Layout.h @@ -0,0 +1,68 @@ +#pragma once + +#include +#include + +#include +#include + +namespace mlir::qco { +/** + * @brief A qubit layout that maps program and hardware indices without + * storing Values. Used for efficient memory usage when Value tracking isn't + * needed. + * + * Note that we use the terminology "hardware" and "program" qubits here, + * because "virtual" (opposed to physical) and "static" (opposed to dynamic) + * are C++ keywords. + */ +class Layout { +public: + /// Construct and return a random layout with size `nqubits`. + static Layout random(size_t nqubits, size_t seed); + + /// Insert program:hardware index mapping. + void add(size_t prog, size_t hw); + + /// Lookup and return program index for a hardware index. + [[nodiscard]] size_t getProgramIndex(size_t hw) const; + + /// Lookup and return hardware index for a program index. + [[nodiscard]] size_t getHardwareIndex(size_t prog) const; + + /// Lookup and return multiple hardware indices at once. + template + requires(sizeof...(ProgIndices) > 0) && + ((std::is_convertible_v) && ...) + [[nodiscard]] auto getHardwareIndices(ProgIndices... progs) const { + return std::tuple{getHardwareIndex(static_cast(progs))...}; + } + + /// Lookup and return multiple program indices at once. + template + requires(sizeof...(HwIndices) > 0) && + ((std::is_convertible_v) && ...) + [[nodiscard]] auto getProgramIndices(HwIndices... hws) const { + return std::tuple{getProgramIndex(static_cast(hws))...}; + } + + /// Swap the mapping to program indices of two hardware indices. + void swap(size_t hwA, size_t hwB); + + /// Return the number of qubits managed by the layout. + [[nodiscard]] size_t nqubits() const; + + /// Return the program to hardware mapping. + [[nodiscard]] ArrayRef getProgramToHardware() const; + +protected: + /// Maps a program qubit index to its hardware index. + SmallVector programToHardware_; + /// Maps a hardware qubit index to its program index. + SmallVector hardwareToProgram_; + +private: + explicit Layout(const size_t nqubits) + : programToHardware_(nqubits), hardwareToProgram_(nqubits) {} +}; +} // namespace mlir::qco \ No newline at end of file diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 1cb58e0360..8fbc9fb776 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -15,11 +15,11 @@ #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/Utils/Drivers.h" #include "mlir/Dialect/QCO/Utils/Graph.h" +#include "mlir/Dialect/QCO/Utils/Layout.h" #include "mlir/Dialect/QCO/Utils/WireIterator.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include "mlir/Dialect/QTensor/Utils/TensorIterator.h" -#include #include #include #include @@ -51,10 +51,8 @@ #include #include #include -#include #include #include -#include #include #include #include @@ -78,148 +76,6 @@ struct MappingPass : impl::MappingPassBase { enum class RoutingMode : std::uint8_t { Cold, Hot }; - /** - * @brief A qubit layout that maps program and hardware indices without - * storing Values. Used for efficient memory usage when Value tracking isn't - * needed. - * - * Note that we use the terminology "hardware" and "program" qubits here, - * because "virtual" (opposed to physical) and "static" (opposed to dynamic) - * are C++ keywords. - */ - class [[nodiscard]] Layout { - public: - /** - * @brief Constructs the identity (i->i) layout. - * @param nqubits The number of qubits. - * @return The identity layout. - */ - static Layout identity(const size_t nqubits) { - Layout layout(nqubits); - for (size_t i = 0; i < nqubits; ++i) { - layout.add(i, i); - } - return layout; - } - - /** - * @brief Constructs a random layout. - * @param nqubits The number of qubits. - * @param seed A seed for randomization. - * @return The random layout. - */ - static Layout random(const size_t nqubits, const size_t seed) { - SmallVector mapping(nqubits); - std::iota(mapping.begin(), mapping.end(), IndexType{0}); - std::ranges::shuffle(mapping, std::mt19937_64{seed}); - - Layout layout(nqubits); - for (const auto [prog, hw] : enumerate(mapping)) { - layout.add(prog, hw); - } - - return layout; - } - - /** - * @brief Insert program:hardware index mapping. - * @param prog The program index. - * @param hw The hardware index. - */ - void add(IndexType prog, IndexType hw) { - assert(prog < programToHardware_.size() && - "add: program index out of bounds"); - assert(hw < hardwareToProgram_.size() && - "add: hardware index out of bounds"); - programToHardware_[prog] = hw; - hardwareToProgram_[hw] = prog; - } - - /** - * @brief Look up program index for a hardware index. - * @param hw The hardware index. - * @return The program index of the respective hardware index. - */ - [[nodiscard]] IndexType getProgramIndex(const IndexType hw) const { - assert(hw < hardwareToProgram_.size() && - "getProgramIndex: hardware index out of bounds"); - return hardwareToProgram_[hw]; - } - - /** - * @brief Look up hardware index for a program index. - * @param prog The program index. - * @return The hardware index of the respective program index. - */ - [[nodiscard]] IndexType getHardwareIndex(const IndexType prog) const { - assert(prog < programToHardware_.size() && - "getHardwareIndex: program index out of bounds"); - return programToHardware_[prog]; - } - - /** - * @brief Convenience function to lookup multiple hardware indices at once. - * @param progs The program indices. - * @return A tuple of hardware indices. - */ - template - requires(sizeof...(ProgIndices) > 0) && - ((std::is_convertible_v) && ...) - [[nodiscard]] auto getHardwareIndices(ProgIndices... progs) const { - return std::tuple{getHardwareIndex(static_cast(progs))...}; - } - - /** - * @brief Convenience function to lookup multiple program indices at once. - * @param hws The hardware indices. - * @return A tuple of program indices. - */ - template - requires(sizeof...(HwIndices) > 0) && - ((std::is_convertible_v) && ...) - [[nodiscard]] auto getProgramIndices(HwIndices... hws) const { - return std::tuple{getProgramIndex(static_cast(hws))...}; - } - - /** - * @brief Swap the mapping to program indices of two hardware indices. - */ - void swap(const IndexType hw0, const IndexType hw1) { - const auto prog0 = hardwareToProgram_[hw0]; - const auto prog1 = hardwareToProgram_[hw1]; - - std::swap(hardwareToProgram_[hw0], hardwareToProgram_[hw1]); - std::swap(programToHardware_[prog0], programToHardware_[prog1]); - } - - /** - * @returns the number of qubits managed by the layout. - */ - [[nodiscard]] size_t nqubits() const { return programToHardware_.size(); } - - /** - * @returns the program to hardware mapping. - */ - [[nodiscard]] ArrayRef getProgramToHardware() const { - return programToHardware_; - } - - protected: - /** - * @brief Maps a program qubit index to its hardware index. - */ - SmallVector programToHardware_; - - /** - * @brief Maps a hardware qubit index to its program index. - */ - SmallVector hardwareToProgram_; - - private: - explicit Layout(const size_t nqubits) - : programToHardware_(nqubits), hardwareToProgram_(nqubits) {} - }; - class [[nodiscard]] AugmentedDevice { public: AugmentedDevice() = default; @@ -411,6 +267,8 @@ struct MappingPass : impl::MappingPassBase { return; } + // TODO: This really should be checked with an query to an previously run + // analysis pass... if (wires->size() > device.nqubits()) { func.emitError() << "requires " + Twine(wires.value().size()) + @@ -902,6 +760,8 @@ struct MappingPass : impl::MappingPassBase { return window; } + /// Return the sequence of SWAPs to move from one layout to another. + /// Implements the 4-Approximation algorithm described in arXiv:1602.05150v3. SmallVector restore(const Layout& from, const Layout& to) { SmallVector swaps; @@ -1033,9 +893,6 @@ struct MappingPass : impl::MappingPassBase { // Preserve program-indexed wire semantics. wires[prog0] = WireIterator(out1); wires[prog1] = WireIterator(out0); - - assert(isa(w0.operation())); - assert(isa(w1.operation())); } layout.swap(hw0, hw1); diff --git a/mlir/lib/Dialect/QCO/Utils/Layout.cpp b/mlir/lib/Dialect/QCO/Utils/Layout.cpp new file mode 100644 index 0000000000..5407df54b1 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Utils/Layout.cpp @@ -0,0 +1,54 @@ +#include "mlir/Dialect/QCO/Utils/Layout.h" + +#include +#include + +#include +#include + +namespace mlir::qco { +Layout Layout::random(const size_t nqubits, const size_t seed) { + SmallVector mapping(nqubits); + std::iota(mapping.begin(), mapping.end(), size_t{0}); + std::ranges::shuffle(mapping, std::mt19937_64{seed}); + + Layout layout(nqubits); + for (const auto [prog, hw] : enumerate(mapping)) { + layout.add(prog, hw); + } + + return layout; +} + +void Layout::add(const size_t prog, const size_t hw) { + assert(prog < programToHardware_.size() && "program index out of bounds"); + assert(hw < hardwareToProgram_.size() && "hardware index out of bounds"); + programToHardware_[prog] = hw; + hardwareToProgram_[hw] = prog; +} + +size_t Layout::getProgramIndex(const size_t hw) const { + assert(hw < hardwareToProgram_.size() && "hardware index out of bounds"); + return hardwareToProgram_[hw]; +} + +size_t Layout::getHardwareIndex(const size_t prog) const { + assert(prog < programToHardware_.size() && "program index out of bounds"); + return programToHardware_[prog]; +} + +void Layout::swap(const size_t hwA, const size_t hwB) { + const auto progA = hardwareToProgram_[hwA]; + const auto progB = hardwareToProgram_[hwB]; + + std::swap(hardwareToProgram_[hwA], hardwareToProgram_[hwB]); + std::swap(programToHardware_[progA], programToHardware_[progB]); +} + +size_t Layout::nqubits() const { return programToHardware_.size(); } + +ArrayRef Layout::getProgramToHardware() const { + return programToHardware_; +} + +} // namespace mlir::qco \ No newline at end of file From cc4b43c4e3810a001aa2fab81bfb5c2c1d169305 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 19 Jun 2026 10:56:17 +0200 Subject: [PATCH 15/25] Extract layout to seperate module --- mlir/include/mlir/Dialect/QCO/Utils/Layout.h | 67 +++++++ .../QCO/Transforms/Mapping/Mapping.cpp | 185 ++---------------- mlir/lib/Dialect/QCO/Utils/Layout.cpp | 54 +++++ 3 files changed, 135 insertions(+), 171 deletions(-) create mode 100644 mlir/include/mlir/Dialect/QCO/Utils/Layout.h create mode 100644 mlir/lib/Dialect/QCO/Utils/Layout.cpp diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Layout.h b/mlir/include/mlir/Dialect/QCO/Utils/Layout.h new file mode 100644 index 0000000000..ceb4a8b413 --- /dev/null +++ b/mlir/include/mlir/Dialect/QCO/Utils/Layout.h @@ -0,0 +1,67 @@ +#pragma once + +#include +#include + +#include +#include + +namespace mlir::qco { + +/// A qubit layout that maps program and hardware indices without +/// storing Values. Used for efficient memory usage when Value tracking isn't +/// needed. +/// +/// Note that we use the terminology "hardware" and "program" qubits +/// here, because "virtual" (opposed to physical) and "static" (opposed to +/// dynamic) are C++ keywords. +class Layout { +public: + /// Construct and return a random layout with size `nqubits`. + static Layout random(size_t nqubits, size_t seed); + + /// Insert program:hardware index mapping. + void add(size_t prog, size_t hw); + + /// Lookup and return program index for a hardware index. + [[nodiscard]] size_t getProgramIndex(size_t hw) const; + + /// Lookup and return hardware index for a program index. + [[nodiscard]] size_t getHardwareIndex(size_t prog) const; + + /// Lookup and return multiple hardware indices at once. + template + requires(sizeof...(ProgIndices) > 0) && + ((std::is_convertible_v) && ...) + [[nodiscard]] auto getHardwareIndices(ProgIndices... progs) const { + return std::tuple{getHardwareIndex(static_cast(progs))...}; + } + + /// Lookup and return multiple program indices at once. + template + requires(sizeof...(HwIndices) > 0) && + ((std::is_convertible_v) && ...) + [[nodiscard]] auto getProgramIndices(HwIndices... hws) const { + return std::tuple{getProgramIndex(static_cast(hws))...}; + } + + /// Swap the mapping to program indices of two hardware indices. + void swap(size_t hwA, size_t hwB); + + /// Return the number of qubits managed by the layout. + [[nodiscard]] size_t nqubits() const; + + /// Return the program to hardware mapping. + [[nodiscard]] ArrayRef getProgramToHardware() const; + +protected: + /// Maps a program qubit index to its hardware index. + SmallVector programToHardware_; + /// Maps a hardware qubit index to its program index. + SmallVector hardwareToProgram_; + +private: + explicit Layout(const size_t nqubits) + : programToHardware_(nqubits), hardwareToProgram_(nqubits) {} +}; +} // namespace mlir::qco \ No newline at end of file diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 1cb58e0360..0a9dc7bb53 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -15,11 +15,11 @@ #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/Utils/Drivers.h" #include "mlir/Dialect/QCO/Utils/Graph.h" +#include "mlir/Dialect/QCO/Utils/Layout.h" #include "mlir/Dialect/QCO/Utils/WireIterator.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include "mlir/Dialect/QTensor/Utils/TensorIterator.h" -#include #include #include #include @@ -51,10 +51,8 @@ #include #include #include -#include #include #include -#include #include #include #include @@ -78,148 +76,6 @@ struct MappingPass : impl::MappingPassBase { enum class RoutingMode : std::uint8_t { Cold, Hot }; - /** - * @brief A qubit layout that maps program and hardware indices without - * storing Values. Used for efficient memory usage when Value tracking isn't - * needed. - * - * Note that we use the terminology "hardware" and "program" qubits here, - * because "virtual" (opposed to physical) and "static" (opposed to dynamic) - * are C++ keywords. - */ - class [[nodiscard]] Layout { - public: - /** - * @brief Constructs the identity (i->i) layout. - * @param nqubits The number of qubits. - * @return The identity layout. - */ - static Layout identity(const size_t nqubits) { - Layout layout(nqubits); - for (size_t i = 0; i < nqubits; ++i) { - layout.add(i, i); - } - return layout; - } - - /** - * @brief Constructs a random layout. - * @param nqubits The number of qubits. - * @param seed A seed for randomization. - * @return The random layout. - */ - static Layout random(const size_t nqubits, const size_t seed) { - SmallVector mapping(nqubits); - std::iota(mapping.begin(), mapping.end(), IndexType{0}); - std::ranges::shuffle(mapping, std::mt19937_64{seed}); - - Layout layout(nqubits); - for (const auto [prog, hw] : enumerate(mapping)) { - layout.add(prog, hw); - } - - return layout; - } - - /** - * @brief Insert program:hardware index mapping. - * @param prog The program index. - * @param hw The hardware index. - */ - void add(IndexType prog, IndexType hw) { - assert(prog < programToHardware_.size() && - "add: program index out of bounds"); - assert(hw < hardwareToProgram_.size() && - "add: hardware index out of bounds"); - programToHardware_[prog] = hw; - hardwareToProgram_[hw] = prog; - } - - /** - * @brief Look up program index for a hardware index. - * @param hw The hardware index. - * @return The program index of the respective hardware index. - */ - [[nodiscard]] IndexType getProgramIndex(const IndexType hw) const { - assert(hw < hardwareToProgram_.size() && - "getProgramIndex: hardware index out of bounds"); - return hardwareToProgram_[hw]; - } - - /** - * @brief Look up hardware index for a program index. - * @param prog The program index. - * @return The hardware index of the respective program index. - */ - [[nodiscard]] IndexType getHardwareIndex(const IndexType prog) const { - assert(prog < programToHardware_.size() && - "getHardwareIndex: program index out of bounds"); - return programToHardware_[prog]; - } - - /** - * @brief Convenience function to lookup multiple hardware indices at once. - * @param progs The program indices. - * @return A tuple of hardware indices. - */ - template - requires(sizeof...(ProgIndices) > 0) && - ((std::is_convertible_v) && ...) - [[nodiscard]] auto getHardwareIndices(ProgIndices... progs) const { - return std::tuple{getHardwareIndex(static_cast(progs))...}; - } - - /** - * @brief Convenience function to lookup multiple program indices at once. - * @param hws The hardware indices. - * @return A tuple of program indices. - */ - template - requires(sizeof...(HwIndices) > 0) && - ((std::is_convertible_v) && ...) - [[nodiscard]] auto getProgramIndices(HwIndices... hws) const { - return std::tuple{getProgramIndex(static_cast(hws))...}; - } - - /** - * @brief Swap the mapping to program indices of two hardware indices. - */ - void swap(const IndexType hw0, const IndexType hw1) { - const auto prog0 = hardwareToProgram_[hw0]; - const auto prog1 = hardwareToProgram_[hw1]; - - std::swap(hardwareToProgram_[hw0], hardwareToProgram_[hw1]); - std::swap(programToHardware_[prog0], programToHardware_[prog1]); - } - - /** - * @returns the number of qubits managed by the layout. - */ - [[nodiscard]] size_t nqubits() const { return programToHardware_.size(); } - - /** - * @returns the program to hardware mapping. - */ - [[nodiscard]] ArrayRef getProgramToHardware() const { - return programToHardware_; - } - - protected: - /** - * @brief Maps a program qubit index to its hardware index. - */ - SmallVector programToHardware_; - - /** - * @brief Maps a hardware qubit index to its program index. - */ - SmallVector hardwareToProgram_; - - private: - explicit Layout(const size_t nqubits) - : programToHardware_(nqubits), hardwareToProgram_(nqubits) {} - }; - class [[nodiscard]] AugmentedDevice { public: AugmentedDevice() = default; @@ -228,21 +84,15 @@ struct MappingPass : impl::MappingPassBase { const llvm::DenseSet>& couplingSet) : coupling_(couplingSet), dist_(coupling_.getDistMatrix()) {} - /** - * @returns the device's number of qubits. - */ + /// Return the device's number of qubits. [[nodiscard]] size_t nqubits() const { return coupling_.getNumNodes(); } - /** - * @returns true if @p u and @p v are adjacent. - */ + /// Return true if two qubits are adjacent. [[nodiscard]] bool areAdjacent(size_t u, size_t v) const { return dist_[u][v] == 1UL; } - /** - * @returns the length of the shortest path between @p u and @p v. - */ + /// Return the length of the shortest path between two qubits. [[nodiscard]] size_t distanceBetween(size_t u, size_t v) const { if (dist_[u][v] == UINT64_MAX) { report_fatal_error("Failed to compute the distance between qubits " + @@ -251,30 +101,22 @@ struct MappingPass : impl::MappingPassBase { return dist_[u][v]; } - /** - * @returns all neighbours of @p u. - */ + /// Return all neighbours of a qubit. [[nodiscard]] ArrayRef neighboursOf(size_t u) const { return coupling_.getEdges(u); } - /** - * @returns the qubit identifiers. - */ + /// Return the qubit identifiers. [[nodiscard]] ArrayRef qubits() const { return coupling_.getNodes(); } - /** - * @returns the links of the device. - */ - [[nodiscard]] llvm::DenseSet> links() const { + /// Return the couplings of the device. + [[nodiscard]] llvm::DenseSet> couplings() const { return coupling_.getEdges(); } - /** - * @returns the max degree (connectivity) of any qubit of the device. - */ + /// Return the max degree (connectivity) of any qubit of the device. [[nodiscard]] size_t maxDegree() const { return coupling_.getMaxDegree(); } private: @@ -411,6 +253,8 @@ struct MappingPass : impl::MappingPassBase { return; } + // TODO: This really should be checked with an query to an previously run + // analysis pass... if (wires->size() > device.nqubits()) { func.emitError() << "requires " + Twine(wires.value().size()) + @@ -902,6 +746,8 @@ struct MappingPass : impl::MappingPassBase { return window; } + /// Return the sequence of SWAPs to move from one layout to another. + /// Implements the 4-Approximation algorithm described in arXiv:1602.05150v3. SmallVector restore(const Layout& from, const Layout& to) { SmallVector swaps; @@ -927,7 +773,7 @@ struct MappingPass : impl::MappingPassBase { do { // Construct 'F' graph. g.clear(); - for (const auto& [hwA, hwB] : device.links()) { + for (const auto& [hwA, hwB] : device.couplings()) { constructEdge(hwA, hwB); constructEdge(hwB, hwA); } @@ -1033,9 +879,6 @@ struct MappingPass : impl::MappingPassBase { // Preserve program-indexed wire semantics. wires[prog0] = WireIterator(out1); wires[prog1] = WireIterator(out0); - - assert(isa(w0.operation())); - assert(isa(w1.operation())); } layout.swap(hw0, hw1); diff --git a/mlir/lib/Dialect/QCO/Utils/Layout.cpp b/mlir/lib/Dialect/QCO/Utils/Layout.cpp new file mode 100644 index 0000000000..5407df54b1 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Utils/Layout.cpp @@ -0,0 +1,54 @@ +#include "mlir/Dialect/QCO/Utils/Layout.h" + +#include +#include + +#include +#include + +namespace mlir::qco { +Layout Layout::random(const size_t nqubits, const size_t seed) { + SmallVector mapping(nqubits); + std::iota(mapping.begin(), mapping.end(), size_t{0}); + std::ranges::shuffle(mapping, std::mt19937_64{seed}); + + Layout layout(nqubits); + for (const auto [prog, hw] : enumerate(mapping)) { + layout.add(prog, hw); + } + + return layout; +} + +void Layout::add(const size_t prog, const size_t hw) { + assert(prog < programToHardware_.size() && "program index out of bounds"); + assert(hw < hardwareToProgram_.size() && "hardware index out of bounds"); + programToHardware_[prog] = hw; + hardwareToProgram_[hw] = prog; +} + +size_t Layout::getProgramIndex(const size_t hw) const { + assert(hw < hardwareToProgram_.size() && "hardware index out of bounds"); + return hardwareToProgram_[hw]; +} + +size_t Layout::getHardwareIndex(const size_t prog) const { + assert(prog < programToHardware_.size() && "program index out of bounds"); + return programToHardware_[prog]; +} + +void Layout::swap(const size_t hwA, const size_t hwB) { + const auto progA = hardwareToProgram_[hwA]; + const auto progB = hardwareToProgram_[hwB]; + + std::swap(hardwareToProgram_[hwA], hardwareToProgram_[hwB]); + std::swap(programToHardware_[progA], programToHardware_[progB]); +} + +size_t Layout::nqubits() const { return programToHardware_.size(); } + +ArrayRef Layout::getProgramToHardware() const { + return programToHardware_; +} + +} // namespace mlir::qco \ No newline at end of file From fd818ffdb71e077593216ee85af8540fe4097019 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 19 Jun 2026 14:58:32 +0200 Subject: [PATCH 16/25] Update isExecutable to handle nested regions --- mlir/include/mlir/Dialect/QCO/Utils/Drivers.h | 67 -------- mlir/include/mlir/Dialect/QCO/Utils/Qubits.h | 52 ------ .../QCO/Transforms/Mapping/Mapping.cpp | 148 +++++++++--------- mlir/lib/Dialect/QCO/Utils/Qubits.cpp | 113 ------------- .../QCO/Transforms/Mapping/test_mapping.cpp | 130 ++++++++++----- .../Dialect/QCO/Utils/test_drivers.cpp | 71 --------- 6 files changed, 161 insertions(+), 420 deletions(-) delete mode 100644 mlir/include/mlir/Dialect/QCO/Utils/Qubits.h delete mode 100644 mlir/lib/Dialect/QCO/Utils/Qubits.cpp diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h index e8663567b0..9188ece1b6 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h @@ -10,10 +10,8 @@ #pragma once -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/IR/QCOInterfaces.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" -#include "mlir/Dialect/QCO/Utils/Qubits.h" #include "mlir/Dialect/QCO/Utils/WireIterator.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" @@ -34,71 +32,6 @@ namespace mlir::qco { -using WalkProgramFn = function_ref; - -/** - * @brief Perform top-down non-recursive walk of all operations within a - * region of a quantum program and apply a callback function. - * @details The signature of the callback function is: - * - * (Operation*, Qubits& q) -> WalkResult - * - * where the Qubits object tracks the front of qubit SSA values. - * Depending on the template parameter, the callback is executed before or after - * updating the Qubits state. - * @param region The targeted region. - * @param qubits An empty or filled qubits object. - * @param fn The callback function. - * @returns success(), if all operations have been visited. - */ -template -LogicalResult walkProgram(Region& region, Qubits& qubits, - const WalkProgramFn& fn) { - for (Operation& curr : region.getOps()) { - if constexpr (Order == WalkOrder::PreOrder) { - if (fn(&curr, qubits).wasInterrupted()) { - return failure(); - } - } - - TypeSwitch(&curr) - .template Case( - [&](StaticOp op) { qubits.add(op.getQubit(), op.getIndex()); }) - .template Case([&](AllocOp op) { qubits.add(op.getResult()); }) - .template Case([&](UnitaryOpInterface& op) { - for (const auto& [prevV, nextV] : - llvm::zip(op.getInputQubits(), op.getOutputQubits())) { - const auto prevQ = cast>(prevV); - const auto nextQ = cast>(nextV); - qubits.remap(prevQ, nextQ); - } - }) - .template Case([&](scf::ForOp op) { - for (OpOperand& operand : op.getInitsMutable()) { - auto result = op.getTiedLoopResult(&operand); - const auto prevQ = cast>(operand.get()); - const auto nextQ = cast>(result); - qubits.remap(prevQ, nextQ); - } - }) - .template Case([&](ResetOp op) { - qubits.remap(op.getQubitIn(), op.getQubitOut()); - }) - .template Case([&](MeasureOp op) { - qubits.remap(op.getQubitIn(), op.getQubitOut()); - }) - .template Case( - [&](SinkOp op) { qubits.remove(op.getQubit()); }); - - if constexpr (Order == WalkOrder::PostOrder) { - if (fn(&curr, qubits).wasInterrupted()) { - return failure(); - } - } - } - return success(); -} - using ReleasedOps = SmallVector; using PendingWiresMap = DenseMap>; diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Qubits.h b/mlir/include/mlir/Dialect/QCO/Utils/Qubits.h deleted file mode 100644 index ccd9d5a577..0000000000 --- a/mlir/include/mlir/Dialect/QCO/Utils/Qubits.h +++ /dev/null @@ -1,52 +0,0 @@ -/* - * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM - * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH - * All rights reserved. - * - * SPDX-License-Identifier: MIT - * - * Licensed under the MIT License - */ - -#pragma once - -#include "mlir/Dialect/QCO/IR/QCODialect.h" - -#include -#include - -#include -#include -#include - -namespace mlir::qco { - -class Qubits { - /// Specifies the qubit "location" (hardware or program). - enum class QubitLocation : std::uint8_t { Hardware, Program }; - -public: - /// Update the qubits map based on the given op. - void update(Operation*); - /// Add qubit with automatically assigned dynamic index. - void add(TypedValue q); - /// Add qubit with static index. - void add(TypedValue q, std::size_t hw); - /// Remap the qubit value from prev to next. - void remap(TypedValue prev, TypedValue next); - /// Remove the qubit value. - void remove(TypedValue q); - /// Return the qubit value assigned to a program index. - [[nodiscard]] TypedValue getProgramQubit(std::size_t index) const; - /// Return the qubit value assigned to a hardware index. - [[nodiscard]] TypedValue getHardwareQubit(std::size_t index) const; - /// Return the index assigned to the qubit value. - [[nodiscard]] std::size_t getIndex(TypedValue q) const; - -private: - DenseMap> programToValue_; - DenseMap> hardwareToValue_; - DenseMap, std::pair> - valueToIndex_; -}; -} // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 0a9dc7bb53..4dc4d4dcd8 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -282,8 +282,6 @@ struct MappingPass : impl::MappingPassBase { return; } - func->dumpPretty(); - // Collect statistics. numSwaps += s.nswaps; @@ -459,28 +457,28 @@ struct MappingPass : impl::MappingPassBase { stack.emplace_back(body, DenseSet{}); while (!stack.empty()) { - auto [region, m] = stack.pop_back_val(); + auto [region, qubits] = stack.pop_back_val(); for (Operation& op : llvm::make_early_inc_range(region.getOps())) { TypeSwitch(&op) - .Case([&](StaticOp op) { m.insert(op.getQubit()); }) + .Case([&](StaticOp op) { qubits.insert(op.getQubit()); }) .Case([&](UnitaryOpInterface& op) { for (const auto [pred, succ] : llvm::zip_equal(op.getInputQubits(), op.getOutputQubits())) { - m.insert(succ); - m.erase(pred); + qubits.insert(succ); + qubits.erase(pred); } }) .Case([&](scf::ForOp loop) { - assert(m.size() == layout.nqubits()); + assert(qubits.size() == layout.nqubits()); - DenseSet addons(m); + DenseSet addons(qubits); llvm::for_each(loop.getInits(), [&](auto v) { addons.erase(v); }); auto newLoop = extend(loop, to_vector(addons), rewriter); for (OpOperand& operand : newLoop.getInitsMutable()) { - m.insert(newLoop.getTiedLoopResult(&operand)); - m.erase(operand.get()); + qubits.insert(newLoop.getTiedLoopResult(&operand)); + qubits.erase(operand.get()); } stack.emplace_back( @@ -489,8 +487,8 @@ struct MappingPass : impl::MappingPassBase { newLoop.getRegionIterArgs().end())); }) .Case([&](auto op) { - m.insert(op.getQubitOut()); - m.erase(op.getQubitIn()); + qubits.insert(op.getQubitOut()); + qubits.erase(op.getQubitIn()); }) .Case([&](auto) { llvm::reportFatalInternalError("unexpected dynamic qubit alloc"); @@ -654,6 +652,65 @@ struct MappingPass : impl::MappingPassBase { return failure(); } + /// Return the sequence of SWAPs to move from one layout to another. + /// Implements the 4-Approximation algorithm described in arXiv:1602.05150v3. + SmallVector restore(const Layout& from, const Layout& to) { + SmallVector swaps; + + Layout curr(from); + + Graph g; + const auto constructEdge = [&](size_t hwX, size_t hwY) { + const auto prog = curr.getProgramIndex(hwX); + + const auto hwGoal = to.getHardwareIndex(prog); + const auto distPre = device.distanceBetween(hwX, hwGoal); + const auto distPost = device.distanceBetween(hwY, hwGoal); + + if (distPost < distPre) { + llvm::dbgs() << "prog=" << prog << " edge=(" << hwX << ", " << hwY + << ")" + << " dist(pre)=" << distPre << " dist(post)=" << distPost + << '\n'; + g.addEdge(hwX, hwY); + } + }; + + do { + // Construct 'F' graph. + g.clear(); + for (const auto& [hwA, hwB] : device.couplings()) { + constructEdge(hwA, hwB); + constructEdge(hwB, hwA); + } + + // Find happy swap chain or unhappy swap. + if (const auto cycle = g.findCycle(); cycle) { + // Apply happy SWAP chain. + for (size_t i = 0; i < cycle->size() - 1; ++i) { + llvm::dbgs() << "happySWAP=(" << (*cycle)[i] << ", " + << (*cycle)[i + 1] << ")\n"; + curr.swap((*cycle)[i], (*cycle)[i + 1]); + swaps.emplace_back((*cycle)[i], (*cycle)[i + 1]); + } + + continue; + } + + for (const auto e : g.getEdges()) { + if (g.getDegree(e.second) == 0) { + llvm::dbgs() << "unhappySWAP=(" << e.first << ", " << e.second + << ")\n"; + curr.swap(e.first, e.second); + swaps.emplace_back(e.first, e.second); + break; + } + } + } while (!g.empty()); + + return swaps; + } + /** * @brief Skip a qubit-pair block. * @details Traverses the pair of wire iterators in tandem until a two-qubit @@ -746,65 +803,6 @@ struct MappingPass : impl::MappingPassBase { return window; } - /// Return the sequence of SWAPs to move from one layout to another. - /// Implements the 4-Approximation algorithm described in arXiv:1602.05150v3. - SmallVector restore(const Layout& from, const Layout& to) { - SmallVector swaps; - - Layout curr(from); - - Graph g; - const auto constructEdge = [&](size_t hwX, size_t hwY) { - const auto prog = curr.getProgramIndex(hwX); - - const auto hwGoal = to.getHardwareIndex(prog); - const auto distPre = device.distanceBetween(hwX, hwGoal); - const auto distPost = device.distanceBetween(hwY, hwGoal); - - if (distPost < distPre) { - llvm::dbgs() << "prog=" << prog << " edge=(" << hwX << ", " << hwY - << ")" - << " dist(pre)=" << distPre << " dist(post)=" << distPost - << '\n'; - g.addEdge(hwX, hwY); - } - }; - - do { - // Construct 'F' graph. - g.clear(); - for (const auto& [hwA, hwB] : device.couplings()) { - constructEdge(hwA, hwB); - constructEdge(hwB, hwA); - } - - // Find happy swap chain or unhappy swap. - if (const auto cycle = g.findCycle(); cycle) { - // Apply happy SWAP chain. - for (size_t i = 0; i < cycle->size() - 1; ++i) { - llvm::dbgs() << "happySWAP=(" << (*cycle)[i] << ", " - << (*cycle)[i + 1] << ")\n"; - curr.swap((*cycle)[i], (*cycle)[i + 1]); - swaps.emplace_back((*cycle)[i], (*cycle)[i + 1]); - } - - continue; - } - - for (const auto e : g.getEdges()) { - if (g.getDegree(e.second) == 0) { - llvm::dbgs() << "unhappySWAP=(" << e.first << ", " << e.second - << ")\n"; - curr.swap(e.first, e.second); - swaps.emplace_back(e.first, e.second); - break; - } - } - } while (!g.empty()); - - return swaps; - } - /** * @brief Advance past all executable gates and recurse into nested regions, * if necessary. @@ -822,7 +820,9 @@ struct MappingPass : impl::MappingPassBase { for (const auto& [readyOp, progs] : ready) { TypeSwitch(readyOp) - .Case([&](UnitaryOpInterface op) { + .template Case( + [&](BarrierOp op) { released.emplace_back(op); }) + .template Case([&](UnitaryOpInterface op) { const auto [hw0, hw1] = layout.getHardwareIndices(progs[0], progs[1]); @@ -830,8 +830,6 @@ struct MappingPass : impl::MappingPassBase { released.emplace_back(op); } }) - .template Case( - [&](BarrierOp op) { released.emplace_back(op); }) .template Case([&](scf::ForOp op) { // TODO: Don't ignore result here. std::ignore = route(op, layout, stats, rewriter); @@ -891,6 +889,8 @@ struct MappingPass : impl::MappingPassBase { * A* search to find a sequence of SWAPs that makes that layer executable. * Depending on the template parameter, this function only updates * (and hence modifies) the layout or also inserts the SWAPs into the IR. + * + * Assumes that wires[i] = i-th program qubit. * @returns failure() if A* search isn't able to find a solution. */ template diff --git a/mlir/lib/Dialect/QCO/Utils/Qubits.cpp b/mlir/lib/Dialect/QCO/Utils/Qubits.cpp deleted file mode 100644 index 12787b0edc..0000000000 --- a/mlir/lib/Dialect/QCO/Utils/Qubits.cpp +++ /dev/null @@ -1,113 +0,0 @@ -/* - * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM - * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH - * All rights reserved. - * - * SPDX-License-Identifier: MIT - * - * Licensed under the MIT License - */ - -#include "mlir/Dialect/QCO/Utils/Qubits.h" - -#include "mlir/Dialect/QCO/IR/QCODialect.h" -#include "mlir/Dialect/QCO/IR/QCOInterfaces.h" -#include "mlir/Dialect/QCO/IR/QCOOps.h" - -#include -#include -#include -#include - -#include -#include -#include - -namespace mlir::qco { -void Qubits::update(Operation* op) { - TypeSwitch(op) - .Case([&](StaticOp op) { add(op.getQubit(), op.getIndex()); }) - .Case([&](AllocOp op) { add(op.getResult()); }) - .Case([&](UnitaryOpInterface& op) { - for (const auto& [prevV, nextV] : - llvm::zip(op.getInputQubits(), op.getOutputQubits())) { - const auto prevQ = cast>(prevV); - const auto nextQ = cast>(nextV); - remap(prevQ, nextQ); - } - }) - .Case([&](scf::ForOp op) { - for (OpOperand& operand : op.getInitsMutable()) { - auto result = op.getTiedLoopResult(&operand); - const auto prevQ = cast>(operand.get()); - const auto nextQ = cast>(result); - remap(prevQ, nextQ); - } - }) - .Case( - [&](ResetOp op) { remap(op.getQubitIn(), op.getQubitOut()); }) - .Case( - [&](MeasureOp op) { remap(op.getQubitIn(), op.getQubitOut()); }) - .Case([&](SinkOp op) { remove(op.getQubit()); }) - .Default([&](Operation* op) -> WalkResult { - const auto name = op->getName().getStringRef(); - reportFatalInternalError("unexpected op" + name); - }); -} - -void Qubits::add(TypedValue q) { - const auto index = programToValue_.size(); - programToValue_.try_emplace(index, q); - valueToIndex_.try_emplace(q, std::make_pair(QubitLocation::Program, index)); -} - -void Qubits::add(TypedValue q, std::size_t hw) { - hardwareToValue_.try_emplace(hw, q); - valueToIndex_.try_emplace(q, std::make_pair(QubitLocation::Hardware, hw)); -} - -void Qubits::remap(TypedValue prev, TypedValue next) { - assert(valueToIndex_.contains(prev)); - const auto& [location, index] = valueToIndex_.lookup(prev); - - valueToIndex_.erase(prev); - valueToIndex_.try_emplace(next, std::make_pair(location, index)); - - if (location == QubitLocation::Program) { - programToValue_[index] = next; - return; - } - - hardwareToValue_[index] = next; -} - -void Qubits::remove(TypedValue q) { - assert(valueToIndex_.contains(q)); - const auto& [location, index] = valueToIndex_.lookup(q); - - valueToIndex_.erase(q); - - if (location == QubitLocation::Program) { - programToValue_.erase(index); - return; - } - - hardwareToValue_.erase(index); -} - -TypedValue Qubits::getProgramQubit(std::size_t index) const { - assert(programToValue_.contains(index)); - return programToValue_.lookup(index); -} - -TypedValue Qubits::getHardwareQubit(std::size_t index) const { - assert(hardwareToValue_.contains(index)); - return hardwareToValue_.lookup(index); -} - -std::size_t Qubits::getIndex(TypedValue q) const { - assert(valueToIndex_.contains(q)); - const auto& res = valueToIndex_.lookup(q); - return res.second; -} -} // namespace mlir::qco diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index f6284443d1..33dcb48f9a 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -10,14 +10,13 @@ #include "mlir/Dialect/QCO/Builder/QCOProgramBuilder.h" #include "mlir/Dialect/QCO/IR/QCODialect.h" -#include "mlir/Dialect/QCO/IR/QCOInterfaces.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/Transforms/Mapping/Mapping.h" #include "mlir/Dialect/QCO/Transforms/Passes.h" -#include "mlir/Dialect/QCO/Utils/Drivers.h" -#include "mlir/Dialect/QCO/Utils/Qubits.h" #include +#include +#include #include #include #include @@ -47,43 +46,90 @@ struct Device { DenseSet> couplingSet; }; -/** - * @returns llvm::success() if all two-qubit gates inside @p region - * fulfill the given coupling constraints. llvm::failure(), otherwise. - */ -static LogicalResult -isExecutable(Region& region, +/// Return true, if the entry point fulfills the given coupling constraints. +static bool +isExecutable(func::FuncOp entry, const DenseSet>& couplingSet) { - Qubits qubits; - return walkProgram( - region, qubits, [&](Operation* curr, const Qubits& qubits) { - if (auto op = dyn_cast(curr)) { - if (isa(op)) { - return WalkResult::advance(); - } - - assert(op.getNumQubits() <= 2 && - "isExecutable: expected two-qubit gate decomposition"); - - if (op.getNumQubits() > 1) { - const auto q0 = cast>(op.getInputQubit(0)); - const auto q1 = cast>(op.getInputQubit(1)); - const auto i0 = qubits.getIndex(q0); - const auto i1 = qubits.getIndex(q1); - - if (!couplingSet.contains(std::make_pair(i0, i1))) { - return WalkResult::interrupt(); + DenseMap indices; + + SmallVector>> stack; + stack.emplace_back(entry.getFunctionBody(), DenseMap{}); + + while (!stack.empty()) { + auto [region, qubits] = stack.pop_back_val(); + + for (Operation& rop : region.getOps()) { + bool executable = true; + TypeSwitch(&rop) + .Case([&](StaticOp op) { + qubits.try_emplace(op.getQubit(), op.getIndex()); + }) + .Case([&](BarrierOp op) { + for (const auto [pred, succ] : + llvm::zip_equal(op.getInputQubits(), op.getOutputQubits())) { + const auto hw = qubits.at(pred); + qubits.try_emplace(succ, hw); + qubits.erase(pred); + } + }) + .Case([&](UnitaryOpInterface& op) { + assert(op.getNumQubits() <= 2 && "expected two-qubit decomp."); + + if (op.getNumQubits() > 1) { + const auto hwA = qubits.at(op.getInputQubit(0)); + const auto hwB = qubits.at(op.getInputQubit(1)); + if (!couplingSet.contains(std::make_pair(hwA, hwB))) { + llvm::dbgs() << "not executable: \n"; + op->dump(); + executable = false; + } } - } - } - return WalkResult::advance(); - }); + for (const auto [pred, succ] : + llvm::zip_equal(op.getInputQubits(), op.getOutputQubits())) { + const auto hw = qubits.at(pred); + qubits.try_emplace(succ, hw); + qubits.erase(pred); + } + }) + .Case([&](scf::ForOp op) { + SmallVector permutation; + DenseMap rqubits; + for (const auto [init, arg] : + llvm::zip_equal(op.getInits(), op.getRegionIterArgs())) { + const auto hw = qubits.at(init); + permutation.emplace_back(hw); + rqubits.try_emplace(arg, hw); + } + + for (OpOperand& operand : op.getInitsMutable()) { + const auto pred = operand.get(); + const auto succ = op.getTiedLoopResult(&operand); + const auto hw = qubits.at(pred); + qubits.try_emplace(succ, hw); + qubits.erase(pred); + } + + stack.emplace_back(op.getRegion(), rqubits); + }) + .Case([&](auto op) { + const auto pred = op.getQubitIn(); + const auto succ = op.getQubitOut(); + const auto hw = qubits.at(pred); + qubits.try_emplace(succ, hw); + qubits.erase(pred); + }); + + if (!executable) { + return false; + } + } + } + + return true; } -/** - * @returns a 9x9 square-grid coupling set. - */ +/// Return a 9x9 square-grid coupling set. static Device getNineQubitSquareGrid() { return {.nqubits = 9, .couplingSet = {{0, 3}, {3, 0}, {0, 1}, {1, 0}, {1, 4}, {4, 1}, @@ -227,8 +273,7 @@ TEST_P(MappingPassTest, GHZ) { auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE( - isExecutable(entry.getFunctionBody(), device.couplingSet).succeeded()); + EXPECT_TRUE(isExecutable(entry, device.couplingSet)); } TEST_P(MappingPassTest, GHZUnrolled) { @@ -272,8 +317,7 @@ TEST_P(MappingPassTest, GHZUnrolled) { auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE( - isExecutable(entry.getFunctionBody(), device.couplingSet).succeeded()); + EXPECT_TRUE(isExecutable(entry, device.couplingSet)); } TEST_P(MappingPassTest, GroverLike) { @@ -361,9 +405,10 @@ TEST_P(MappingPassTest, GroverLike) { auto res = pm.run(m.get()); auto entry = getEntryPoint(m.get()); + m->dump(); + ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE( - isExecutable(entry.getFunctionBody(), device.couplingSet).succeeded()); + EXPECT_TRUE(isExecutable(entry, device.couplingSet)); } TEST_P(MappingPassTest, Sabre) { @@ -457,8 +502,7 @@ TEST_P(MappingPassTest, Sabre) { auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); - EXPECT_TRUE( - isExecutable(entry.getFunctionBody(), device.couplingSet).succeeded()); + EXPECT_TRUE(isExecutable(entry, device.couplingSet)); } INSTANTIATE_TEST_SUITE_P(NineQubitSquareGrid, MappingPassTest, diff --git a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp index 9606398018..9b6c5d2c1a 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp @@ -49,77 +49,6 @@ class DriversTest : public testing::Test { }; } // namespace -TEST_F(DriversTest, ProgramWalk) { - qco::QCOProgramBuilder builder(context.get()); - builder.initialize(); - - Value q0 = builder.allocQubit(); - Value q1 = builder.allocQubit(); - Value q2 = builder.allocQubit(); - Value q3 = builder.allocQubit(); - - q0 = builder.h(q0); - std::tie(q0, q1) = builder.cx(q0, q1); - std::tie(q2, q3) = builder.cx(q2, q3); - - const auto forOut = builder.scfFor( - 1, 3, 1, {q0, q1, q2, q3}, [&builder](Value, ValueRange iterArgs) { - return SmallVector{iterArgs[0], iterArgs[1], iterArgs[2], iterArgs[3]}; - }); - - q0 = forOut[0]; - q1 = forOut[1]; - q2 = forOut[2]; - q3 = forOut[3]; - - Value c0; - Value c1; - Value c2; - Value c3; - - std::tie(q0, c0) = builder.measure(q0); - Operation* firstMeasure = q0.getDefiningOp(); - - std::tie(q1, c1) = builder.measure(q1); - std::tie(q2, c2) = builder.measure(q2); - std::tie(q3, c3) = builder.measure(q3); - - builder.sink(q0); - builder.sink(q1); - builder.sink(q2); - builder.sink(q3); - - auto m = builder.finalize(); - auto func = qco::getEntryPoint(m.get()); - - Value ex0 = nullptr; - Value ex1 = nullptr; - Value ex2 = nullptr; - Value ex3 = nullptr; - - // Walk until the first measurement operation is encountered and stop. - // Since WalkOrder::PreOrder is used here, the state of the qubits is not yet - // updated with the SSA values of the measurement op. - // Consequently, the program qubits point at the outputs of the controlled-Xs. - qco::Qubits qubits; - std::ignore = qco::walkProgram(func.getBody(), qubits, - [&](Operation* op, const qco::Qubits& qubits) { - if (op == firstMeasure) { - ex0 = qubits.getProgramQubit(0); - ex1 = qubits.getProgramQubit(1); - ex2 = qubits.getProgramQubit(2); - ex3 = qubits.getProgramQubit(3); - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - - ASSERT_EQ(ex0, forOut[0]); - ASSERT_EQ(ex1, forOut[1]); - ASSERT_EQ(ex2, forOut[2]); - ASSERT_EQ(ex3, forOut[3]); -} - TEST_F(DriversTest, ProgramGraphWalk) { qco::QCOProgramBuilder builder(context.get()); builder.initialize(); From cbab3f87d7c170be14406f099bacc75d00b62e38 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Mon, 22 Jun 2026 15:16:52 +0200 Subject: [PATCH 17/25] Cleanup --- .../QCO/Transforms/Mapping/Mapping.cpp | 464 ++++++++++-------- .../QCO/Transforms/Mapping/test_mapping.cpp | 155 +++--- 2 files changed, 333 insertions(+), 286 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 4dc4d4dcd8..603ec9fdf6 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -20,6 +20,7 @@ #include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include "mlir/Dialect/QTensor/Utils/TensorIterator.h" +#include #include #include #include @@ -70,13 +71,13 @@ namespace { struct MappingPass : impl::MappingPassBase { private: - using IndexType = size_t; - using IndexPairType = std::pair; + using IndexPairType = std::pair; using Window = SmallVector; + using Wires = SmallVector; - enum class RoutingMode : std::uint8_t { Cold, Hot }; + enum class RoutingMode : bool { Cold, Hot }; - class [[nodiscard]] AugmentedDevice { + class AugmentedDevice { public: AugmentedDevice() = default; @@ -88,8 +89,8 @@ struct MappingPass : impl::MappingPassBase { [[nodiscard]] size_t nqubits() const { return coupling_.getNumNodes(); } /// Return true if two qubits are adjacent. - [[nodiscard]] bool areAdjacent(size_t u, size_t v) const { - return dist_[u][v] == 1UL; + [[nodiscard]] bool areAdjacent(std::pair qubits) const { + return dist_[qubits.first][qubits.second] == 1UL; } /// Return the length of the shortest path between two qubits. @@ -124,23 +125,51 @@ struct MappingPass : impl::MappingPassBase { Matrix dist_; }; + struct WireInfos { + /// Return the mapped wire index of a program index. + [[nodiscard]] size_t lookupIndex(size_t prog) const { + return programToIndex.at(prog); + } + + /// Return the mapped program index of a wire index. + [[nodiscard]] size_t lookupProgram(size_t index) const { + return indexToProgram.at(index); + } + + /// Bidirectionally map a wire index to a program index. + /// Overwrites existing mappings. + void map(size_t index, size_t prog) { + indexToProgram[index] = prog; + programToIndex[prog] = index; + } + + /// Swap two program indices. + void swap(size_t prog0, size_t prog1) { + const auto i0 = lookupIndex(prog0); + const auto i1 = lookupIndex(prog1); + std::swap(programToIndex[prog0], programToIndex[prog1]); + std::swap(indexToProgram[i0], indexToProgram[i1]); + } + + /// Maps the i-th wire index to a program index. + DenseMap indexToProgram; + /// Maps a program index to the i-th wire index. + DenseMap programToIndex; + }; + /// Statistics collected while routing. struct Statistics { size_t nswaps{0}; }; - /** - * @brief Parameters influencing the behavior of the A* search algorithm. - */ - struct [[nodiscard]] Parameters { + /// Parameters influencing the behavior of the A* search algorithm. + struct Parameters { float alpha; float lambda; }; - /** - * @brief Describes a node in the A* search graph. - */ - struct [[nodiscard]] Node { + /// Describes a node in the A* search graph. + struct Node { struct ComparePointer { bool operator()(const Node* lhs, const Node* rhs) const { return lhs->f > rhs->f; @@ -178,8 +207,8 @@ struct MappingPass : impl::MappingPassBase { */ [[nodiscard]] bool isGoal(const IndexPairType& front, const AugmentedDevice& device) const { - return device.areAdjacent(layout.getHardwareIndex(front.first), - layout.getHardwareIndex(front.second)); + return device.areAdjacent( + layout.getHardwareIndices(front.first, front.second)); } private: @@ -221,8 +250,10 @@ struct MappingPass : impl::MappingPassBase { public: /// Construct default mapping pass. MappingPass() = default; + /// Construct default mapping pass with options. explicit MappingPass(MappingPassOptions options) : MappingPassBase(options) {} + /// Construct mapping from coupling set. explicit MappingPass( const llvm::DenseSet>& couplingSet, @@ -231,9 +262,9 @@ struct MappingPass : impl::MappingPassBase { protected: void runOnOperation() override { - assert(alpha > 0 && "runOnOperation: expected alpha > 0"); - assert(niterations > 0 && "runOnOperation: expected niterations > 0"); - assert(ntrials > 0 && "runOnOperation: expected ntrials > 0"); + assert(alpha > 0 && "expected alpha > 0"); + assert(niterations > 0 && "expected niterations > 0"); + assert(ntrials > 0 && "expected ntrials > 0"); IRRewriter rewriter(&getContext()); @@ -245,37 +276,36 @@ struct MappingPass : impl::MappingPassBase { return; } - auto& body = func.getFunctionBody(); - - auto wires = getComputation(func); - if (failed(wires)) { + auto comp = getComputation(func); + if (failed(comp)) { signalPassFailure(); return; } - // TODO: This really should be checked with an query to an previously run - // analysis pass... - if (wires->size() > device.nqubits()) { + auto& body = func.getFunctionBody(); + auto& [wires, infos] = *comp; + + if (wires.size() > device.nqubits()) { func.emitError() - << "requires " + Twine(wires.value().size()) + + << "requires " + Twine(wires.size()) + " qubits. However, the architecture only supports " + Twine(device.nqubits()) + "qubits."; signalPassFailure(); return; } - auto layout = generateLayout(*wires); + auto layout = generateLayout(wires, infos); if (failed(layout)) { func->emitError() << "failed to refine random initial layouts."; signalPassFailure(); } - wires = std::move(place(body, *layout, rewriter)); + std::tie(wires, infos) = std::move(place(body, *layout, rewriter)); Statistics s; const auto res = route( - *wires, *layout, s, &rewriter); + wires, infos, *layout, s, &rewriter); if (res.failed()) { func.emitError() << "failed to map the function"; signalPassFailure(); @@ -286,8 +316,7 @@ struct MappingPass : impl::MappingPassBase { numSwaps += s.nswaps; // Fix SSA Dominance issues. - llvm::for_each(func.getFunctionBody().getBlocks(), - [](Block& b) { sortTopologically(&b); }); + llvm::for_each(body.getBlocks(), [](Block& b) { sortTopologically(&b); }); } private: @@ -332,45 +361,49 @@ struct MappingPass : impl::MappingPassBase { rewriter.replaceOpWithNewOp(yield, results); rewriter.eraseOp(loop); - return newLoop; } - /** - * @brief Collect wires of the quantum computation before placement. - * @details - * The mapping pass currently assumes that the quantum computation allocates - * all tensors at the start of the function. The required qubits are extracted - * from these tensors and used for the computation. Finally, the qubits are - * inserted back into the tensors at the end of the function. - * Thus, a valid program has the following structure: - * - * T ⨉ [qtensor::AllocOp] - * → N ⨉ [qtensor::ExtractOp] - * → (Computation) - * → N ⨉ [qtensor::InsertOp] - * → T ⨉ [qtensor::DeallocOp] - * - * @returns a vector of wire iterator, or failure() if any of the above - * assumptions are violated. - */ - static FailureOr> + /// Return the wires of a dynamic computation. + /// The mapping pass currently assumes that + /// - there are no `qco.alloc` operation + /// - there is an "extraction" and "insertion" phase, where the i-th extract + /// defines the i-th program qubit + /// Thus, supported programs have the following structure: + /// + /// T ⨉ [qtensor::AllocOp] + /// → N ⨉ [qtensor::ExtractOp] + /// → (Computation) + /// → N ⨉ [qtensor::InsertOp] + /// → T ⨉ [qtensor::DeallocOp] + /// + /// If any of the above assumptions are violated, the function returns + /// failure. + static FailureOr> getComputation(func::FuncOp func) { if (!func.getOps().empty()) { return func.emitError() << "must not contain qco.alloc operations"; } - SmallVector wires; - for (auto tensor : func.getOps()) { + Wires wires; + WireInfos infos; + + for (auto alloc : func.getOps()) { bool isInitPhase = true; - TensorIterator it(tensor.getResult()); + TensorIterator it(alloc.getResult()); for (; it != std::default_sentinel; ++it) { if (auto extract = dyn_cast(it.operation())) { if (!isInitPhase) { return func.emitError() << "must extract and insert all qubits at once."; } - wires.emplace_back(extract.getResult()); + + const auto qubit = extract.getResult(); + const auto index = wires.size(); + + wires.emplace_back(qubit); + infos.map(index, index); + continue; } @@ -380,20 +413,19 @@ struct MappingPass : impl::MappingPassBase { } } } - return wires; + + return std::make_pair(wires, infos); } - /** - * @brief Perform placement by replacing dynamic with static qubits. - * @details - * Creates static qubits and replaces the extracted qubits with it. - * Moreover, the function extends the computation with as many static qubits - * as the architecture supports. - * @returns a vector of wire iterators, where the i-th wire points at the i-th - * static program qubit. - */ - static SmallVector place(Region& body, const Layout& layout, - IRRewriter& rewriter) { + /// Perform placement by + /// - initializing as many hardware qubits as the architecture supports + /// - replacing dynamic with static qubits + /// - extending the inputs of `scf::ForOp` to all hardware qubits. + /// + /// Analogously to the getComputation function, the i-th extract + /// operation defines the i-th program qubit. + static std::pair place(Region& body, const Layout& layout, + IRRewriter& rewriter) { SmallVector staticOps; staticOps.reserve(layout.nqubits()); @@ -406,10 +438,12 @@ struct MappingPass : impl::MappingPassBase { } // Replace extract ops and collect in program-qubit order. - SmallVector wires(layout.nqubits()); - size_t prog = 0UL; - for (auto alloc : make_early_inc_range(body.getOps())) { + Wires wires; + WireInfos infos; + + for (auto alloc : + llvm::make_early_inc_range(body.getOps())) { TensorIterator it(alloc.getResult()); while (it != std::default_sentinel) { // Get the operation and early increment to avoid issues after erasure. @@ -418,6 +452,7 @@ struct MappingPass : impl::MappingPassBase { TypeSwitch(curr) .Case([&](auto op) { + const auto prog = wires.size(); const auto hw = layout.getHardwareIndex(prog); const auto qubit = staticOps[hw].getQubit(); @@ -425,8 +460,8 @@ struct MappingPass : impl::MappingPassBase { rewriter.replaceAllUsesWith(op.getOutTensor(), op.getTensor()); rewriter.eraseOp(op); - wires[prog] = WireIterator(qubit); - ++prog; + wires.emplace_back(qubit); + infos.map(prog, prog); }) .Case([&](auto op) { rewriter.setInsertionPointAfter(op); @@ -441,11 +476,15 @@ struct MappingPass : impl::MappingPassBase { } // Create sinks for remaining, unused, static qubits. + rewriter.setInsertionPoint(body.back().getTerminator()); - for (; prog < layout.nqubits(); ++prog) { + for (size_t prog = wires.size(); prog < layout.nqubits(); ++prog) { const auto hw = layout.getHardwareIndex(prog); const auto qubit = staticOps[hw].getQubit(); - wires[prog] = WireIterator(qubit); + + wires.emplace_back(qubit); + infos.map(prog, prog); + SinkOp::create(rewriter, body.getLoc(), qubit); } @@ -496,7 +535,7 @@ struct MappingPass : impl::MappingPassBase { } } - return wires; + return std::make_pair(wires, infos); } /// Execute `ntrials` many (parallel) initial layout refinement trials and @@ -507,7 +546,7 @@ struct MappingPass : impl::MappingPassBase { /// cold-route along the way. Repeat this procedure "niterations" times and /// finally find the trial with the fewest SWAPs on the final backwards pass /// and return the respective layout. - FailureOr generateLayout(ArrayRef wires) { + FailureOr generateLayout(const Wires& wires, const WireInfos& infos) { std::mt19937_64 rng{seed}; struct Trial { @@ -523,15 +562,21 @@ struct MappingPass : impl::MappingPassBase { } parallelForEach(&getContext(), trials, [&, this](Trial& t) { - SmallVector local(wires); + Wires locWires(wires); + WireInfos locInfos(infos); + for (size_t i = 0; i < niterations; ++i) { - if (route(local, t.layout, t.stats).failed()) { + if (route(locWires, locInfos, t.layout, t.stats) + .failed()) { return; } - if (route(local, t.layout, t.stats).failed()) { + if (route(locWires, locInfos, t.layout, + t.stats) + .failed()) { return; } } + t.success = true; }); @@ -588,7 +633,7 @@ struct MappingPass : impl::MappingPassBase { frontier.emplace(root); - DenseMap, size_t> bestDepth; + DenseMap, size_t> bestDepth; DenseSet expansionSet; size_t i = 0; @@ -711,16 +756,16 @@ struct MappingPass : impl::MappingPassBase { return swaps; } - /** - * @brief Skip a qubit-pair block. - * @details Traverses the pair of wire iterators in tandem until a two-qubit - * operation is found. If the two-qubit operation is equivalent, continue. - * Otherwise stop. - */ + /// Skip to the end of the two-qubit block for both wire iterators, where + /// initially both must point at the same two-qubit operation. template static void skipQubitPairBlock(WireIterator& w0, WireIterator& w1) { using Traits = WireTraversalTraits; + // Traverses the pair of wire iterators in tandem until a two-qubit + // operation is found. If the two-qubit operation is equivalent, continue. + // Otherwise stop. + WireIterator curr0(w0); WireIterator curr1(w1); while (true) { @@ -760,35 +805,32 @@ struct MappingPass : impl::MappingPassBase { } } - /** - * @brief Build and return window of layers. - * @details Traverses the circuit-layers until the desired window sizes is - * reached. Assumes that wires[i] = i-th program qubit. The size of the - * window is 1 + nlookahead. - * @returns window of layers. - */ + /// Return a window of layers with a maximum size of `1 + nlookahead`. template - Window getWindow(ArrayRef wires) { + Window getWindow(Wires wires, const WireInfos& infos) { Window window; window.reserve(1 + nlookahead); - SmallVector local(wires); walkProgramGraph( - local, [&](const ReadyRange& ready, ReleasedOps& released) { + wires, [&](const ReadyRange& ready, ReleasedOps& released) { if (ready.empty()) { return WalkResult::advance(); } - for (const auto& [op, progs] : ready) { + for (const auto& [op, indices] : ready) { if (auto u = dyn_cast(op)) { - const auto p0 = progs[0]; - const auto p1 = progs[1]; - window.emplace_back(p0, p1); + const auto i0 = indices[0]; + const auto i1 = indices[1]; + + const auto prog0 = infos.lookupProgram(i0); + const auto prog1 = infos.lookupProgram(i1); + + window.emplace_back(prog0, prog1); if (window.size() == 1 + nlookahead) { return WalkResult::interrupt(); } - skipQubitPairBlock(local[p0], local[p1]); + skipQubitPairBlock(wires[i0], wires[i1]); released.emplace_back(u); return WalkResult::advance(); } @@ -803,6 +845,44 @@ struct MappingPass : impl::MappingPassBase { return window; } + /// Insert SWAP operations, exchanging two qubits, virtually (Mode::Cold) or + /// into the IR (Mode::Hot). The function expects that each wire points at the + /// correct insertion point. + template + static void insertSWAPs(ArrayRef swaps, const Wires& wires, + WireInfos& infos, Layout& layout, Statistics& stats, + IRRewriter* rewriter) { + for (const auto& [hw0, hw1] : swaps) { + if constexpr (Mode == RoutingMode::Hot) { + const auto [prog0, prog1] = layout.getProgramIndices(hw0, hw1); + + const auto i0 = infos.lookupIndex(prog0); + const auto i1 = infos.lookupIndex(prog1); + + const auto& w0 = wires[i0]; + const auto& w1 = wires[i1]; + + const auto in0 = w0.qubit(); + const auto in1 = w1.qubit(); + + rewriter->setInsertionPointAfterValue(in0); // Valid bc. Hot => Forward. + auto swapOp = SWAPOp::create(*rewriter, in0.getLoc(), in0, in1); + + const auto out0 = swapOp.getQubit0Out(); + const auto out1 = swapOp.getQubit1Out(); + + rewriter->replaceAllUsesExcept(in0, out1, swapOp); + rewriter->replaceAllUsesExcept(in1, out0, swapOp); + + infos.swap(prog0, prog1); + } + + layout.swap(hw0, hw1); + } + + stats.nswaps += swaps.size(); + } + /** * @brief Advance past all executable gates and recurse into nested regions, * if necessary. @@ -810,29 +890,75 @@ struct MappingPass : impl::MappingPassBase { * executable gates are found. */ template - void advanceAndRecurse(MutableArrayRef wires, Layout& layout, - Statistics& stats, IRRewriter* rewriter) { + void advanceAndRecurse(Wires& wires, const WireInfos& infos, + const Layout& layout, Statistics& stats, + IRRewriter* rewriter) { walkProgramGraph(wires, [&](const ReadyRange& ready, ReleasedOps& released) { if (ready.empty()) { return WalkResult::advance(); } - for (const auto& [readyOp, progs] : ready) { + for (const auto& [readyOp, indices] : ready) { TypeSwitch(readyOp) .template Case( [&](BarrierOp op) { released.emplace_back(op); }) .template Case([&](UnitaryOpInterface op) { - const auto [hw0, hw1] = - layout.getHardwareIndices(progs[0], progs[1]); - - if (device.areAdjacent(hw0, hw1)) { + const auto prog0 = infos.lookupProgram(indices[0]); + const auto prog1 = infos.lookupProgram(indices[1]); + if (device.areAdjacent(layout.getHardwareIndices(prog0, prog1))) { released.emplace_back(op); } }) .template Case([&](scf::ForOp op) { - // TODO: Don't ignore result here. - std::ignore = route(op, layout, stats, rewriter); + Wires childWires; + WireInfos childInfos; + + // First map parent (results) to child values (iteration + // arguments). Going forwards, the recursive routing starts at + // block arguments, while the backwards go starts at the yielded + // values. + + for (auto i : indices) { + const auto prog = infos.lookupProgram(i); + const auto arg = op.getTiedLoopRegionIterArg( + cast(wires[i].qubit())); + const auto index = childWires.size(); + + if constexpr (Direction == WireDirection::Forward) { + childWires.emplace_back(arg); + childInfos.map(index, prog); + } else { + const auto yield = op.getTiedLoopYieldedValue(arg)->get(); + childWires.emplace_back(yield); + childInfos.map(index, prog); + } + } + + Layout childLayout(layout); + const auto res = route( + childWires, childInfos, childLayout, stats, rewriter); + if (failed(res)) { + // TODO: How to propagate errors here? + } + + if constexpr (Mode == RoutingMode::Hot) { + + // After routing the loop body, all iterators point to + // std::default_sentinel. To move the iterators to the correct + // qubit SSA values for the epilogue SWAPs, decrement each + // twice: (sentinel → yield → unitary/block arg). + + llvm::for_each(childWires, + [](auto& it) { std::advance(it, -2); }); + } + + insertSWAPs(restore(childLayout, layout), childWires, + childInfos, childLayout, stats, rewriter); + + if constexpr (Mode == RoutingMode::Hot) { + sortTopologically(op.getBody()); + } released.emplace_back(op); }); @@ -847,60 +973,19 @@ struct MappingPass : impl::MappingPassBase { }); } - /// Insert SWAP operations, exchanging two qubits, virtually (Mode::Cold) or - /// into the IR (Mode::Hot). The function expects that the i-th wire points - /// at the i-th program qubit and that each wire also points at the correct - /// insertion point. The function preserves the program qubit order, by - /// exchanging wires after a SWAP. - template - static void insertSWAPs(ArrayRef swaps, - MutableArrayRef wires, Layout& layout, - IRRewriter* rewriter) { - for (const auto& [hw0, hw1] : swaps) { - if constexpr (Mode == RoutingMode::Hot) { - const auto& [prog0, prog1] = layout.getProgramIndices(hw0, hw1); - const auto& w0 = wires[prog0]; - const auto& w1 = wires[prog1]; - - const auto in0 = w0.qubit(); - const auto in1 = w1.qubit(); - - rewriter->setInsertionPointAfterValue(in0); // Valid bc. Hot => Forward. - auto swapOp = SWAPOp::create(*rewriter, in0.getLoc(), in0, in1); - - const auto out0 = swapOp.getQubit0Out(); - const auto out1 = swapOp.getQubit1Out(); - - rewriter->replaceAllUsesExcept(in0, out1, swapOp); - rewriter->replaceAllUsesExcept(in1, out0, swapOp); - - // Preserve program-indexed wire semantics. - wires[prog0] = WireIterator(out1); - wires[prog1] = WireIterator(out0); - } - - layout.swap(hw0, hw1); - } - } - - /** - * @brief Route via SWAP insertion. - * @details Iterates over a dynamically computed window of layers and uses - * A* search to find a sequence of SWAPs that makes that layer executable. - * Depending on the template parameter, this function only updates - * (and hence modifies) the layout or also inserts the SWAPs into the IR. - * - * Assumes that wires[i] = i-th program qubit. - * @returns failure() if A* search isn't able to find a solution. - */ + /// Iterates over a dynamically computed window of layers and uses A* search + /// to find a SWAP sequence that makes each layer executable. Depending on + /// the template parameter, this function only updates the layout or also + /// inserts the SWAPs into the IR. The function returns `failure` if A* is + /// unable to find a solution. template requires(Mode != RoutingMode::Hot || Direction == WireDirection::Forward) - LogicalResult route(SmallVector& wires, Layout& layout, + LogicalResult route(Wires& wires, WireInfos& infos, Layout& layout, Statistics& stats, IRRewriter* rewriter = nullptr) { while (true) { - advanceAndRecurse(wires, layout, stats, rewriter); + advanceAndRecurse(wires, infos, layout, stats, rewriter); - const auto window = getWindow(wires); + const auto window = getWindow(wires, infos); if (window.empty()) { break; } @@ -913,18 +998,17 @@ struct MappingPass : impl::MappingPassBase { if constexpr (Mode == RoutingMode::Hot) { // At this point the wire iterators either point to - // std::default_sentinel or a multi-qubit gate (including barriers) of + // std::default_sentinel or a multi-qubit gate (incl. barriers) of // the current or subsequent layers. The former must be decremented - // twice (sentinel -> sink -> unitary/static). For the latter we - // simply must ensure the insertion point is before the multi-qubit - // gates. + // twice (sentinel → sink → unitary/static). For the latter, we + // must ensure the insertion point is before the multi-qubit gates. for (auto& it : wires) { - std::ranges::advance(it, it == std::default_sentinel ? -2 : -1); + std::advance(it, it == std::default_sentinel ? -2 : -1); } } - insertSWAPs(*swaps, wires, layout, rewriter); + insertSWAPs(*swaps, wires, infos, layout, stats, rewriter); if constexpr (Mode == RoutingMode::Hot) { @@ -936,54 +1020,8 @@ struct MappingPass : impl::MappingPassBase { // multi-qubit op of the current or subsequent layer or to a sink (and // thus std::default_sentinel). - llvm::for_each(wires, [](auto& it) { std::ranges::advance(it, 1); }); + llvm::for_each(wires, [](auto& it) { std::advance(it, 1); }); } - - stats.nswaps += swaps->size(); - } - - return success(); - } - - template - requires(Mode != RoutingMode::Hot || Direction == WireDirection::Forward) - LogicalResult route(scf::ForOp op, Layout& base, Statistics& stats, - IRRewriter* rewriter) { - - assert(llvm::all_of(op.getInitArgs(), - [](Value v) { return isa(v.getType()); })); - - // In the forward direction we start the block arguments of the loop body, - // whereas in the backward direction we start at the yielded values. - - SmallVector wires; - if constexpr (Direction == WireDirection::Forward) { - llvm::for_each(op.getRegionIterArgs(), - [&](Value v) { wires.emplace_back(v); }); - } else { - auto yield = cast(op.getBody()->getTerminator()); - assert(yield != nullptr); - llvm::for_each(yield.getResults(), - [&](Value v) { wires.emplace_back(v); }); - } - - Layout remote(base); - if (route(wires, remote, stats, rewriter).failed()) { - return failure(); - } - - if constexpr (Mode == RoutingMode::Hot) { - assert(llvm::all_of( - wires, [](const auto& it) { return it == std::default_sentinel; })); - llvm::for_each(wires, [](auto& it) { std::ranges::advance(it, -2); }); - } - - const auto swaps = restore(remote, base); - insertSWAPs(swaps, wires, remote, rewriter); - stats.nswaps += swaps.size(); - - if constexpr (Mode == RoutingMode::Hot) { - sortTopologically(op.getBody()); } return success(); diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 33dcb48f9a..f1b6a7e88e 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -46,89 +46,98 @@ struct Device { DenseSet> couplingSet; }; -/// Return true, if the entry point fulfills the given coupling constraints. +/// Return true, if the operations within a region fulfill the given coupling +/// constraints. static bool -isExecutable(func::FuncOp entry, +isExecutable(Region& body, DenseMap& m, const DenseSet>& couplingSet) { - DenseMap indices; - - SmallVector>> stack; - stack.emplace_back(entry.getFunctionBody(), DenseMap{}); - - while (!stack.empty()) { - auto [region, qubits] = stack.pop_back_val(); - - for (Operation& rop : region.getOps()) { - bool executable = true; - TypeSwitch(&rop) - .Case([&](StaticOp op) { - qubits.try_emplace(op.getQubit(), op.getIndex()); - }) - .Case([&](BarrierOp op) { - for (const auto [pred, succ] : - llvm::zip_equal(op.getInputQubits(), op.getOutputQubits())) { - const auto hw = qubits.at(pred); - qubits.try_emplace(succ, hw); - qubits.erase(pred); + for (Operation& rop : body.getOps()) { + bool executable = true; + TypeSwitch(&rop) + .Case( + [&](StaticOp op) { m.try_emplace(op.getQubit(), op.getIndex()); }) + .Case([&](BarrierOp op) { + for (const auto [pred, succ] : + llvm::zip_equal(op.getInputQubits(), op.getOutputQubits())) { + const auto hw = m.at(pred); + m.try_emplace(succ, hw); + } + }) + .Case([&](UnitaryOpInterface& op) { + assert(op.getNumQubits() <= 2 && "expected two-qubit decomp."); + + if (op.getNumQubits() > 1) { + const auto hwA = m.at(op.getInputQubit(0)); + const auto hwB = m.at(op.getInputQubit(1)); + if (!couplingSet.contains(std::make_pair(hwA, hwB))) { + llvm::dbgs() << "not executable: \n"; + op->dump(); + executable = false; } - }) - .Case([&](UnitaryOpInterface& op) { - assert(op.getNumQubits() <= 2 && "expected two-qubit decomp."); - - if (op.getNumQubits() > 1) { - const auto hwA = qubits.at(op.getInputQubit(0)); - const auto hwB = qubits.at(op.getInputQubit(1)); - if (!couplingSet.contains(std::make_pair(hwA, hwB))) { - llvm::dbgs() << "not executable: \n"; - op->dump(); - executable = false; - } - } - - for (const auto [pred, succ] : - llvm::zip_equal(op.getInputQubits(), op.getOutputQubits())) { - const auto hw = qubits.at(pred); - qubits.try_emplace(succ, hw); - qubits.erase(pred); - } - }) - .Case([&](scf::ForOp op) { - SmallVector permutation; - DenseMap rqubits; - for (const auto [init, arg] : - llvm::zip_equal(op.getInits(), op.getRegionIterArgs())) { - const auto hw = qubits.at(init); - permutation.emplace_back(hw); - rqubits.try_emplace(arg, hw); - } - - for (OpOperand& operand : op.getInitsMutable()) { - const auto pred = operand.get(); - const auto succ = op.getTiedLoopResult(&operand); - const auto hw = qubits.at(pred); - qubits.try_emplace(succ, hw); - qubits.erase(pred); + } + + for (const auto [pred, succ] : + llvm::zip_equal(op.getInputQubits(), op.getOutputQubits())) { + const auto hw = m.at(pred); + m.try_emplace(succ, hw); + } + }) + .Case([&](scf::ForOp op) { + DenseMap loopM; + for (const auto [init, arg] : + llvm::zip_equal(op.getInits(), op.getRegionIterArgs())) { + const auto hw = m.at(init); + loopM.try_emplace(arg, hw); + } + + for (OpOperand& operand : op.getInitsMutable()) { + const auto pred = operand.get(); + const auto succ = op.getTiedLoopResult(&operand); + const auto hw = m.at(pred); + m.try_emplace(succ, hw); + } + + if (!isExecutable(op.getRegion(), loopM, couplingSet)) { + executable = false; + return; + } + + for (const auto& [arg, yielded] : + llvm::zip_equal(op.getRegionIterArgs(), op.getYieldedValues())) { + if (loopM.at(arg) != loopM.at(yielded)) { + llvm::dbgs() << "for loop layout not restored!\n"; + executable = false; + return; } - - stack.emplace_back(op.getRegion(), rqubits); - }) - .Case([&](auto op) { - const auto pred = op.getQubitIn(); - const auto succ = op.getQubitOut(); - const auto hw = qubits.at(pred); - qubits.try_emplace(succ, hw); - qubits.erase(pred); - }); - - if (!executable) { - return false; - } + } + }) + .Case([&](scf::YieldOp op) { + assert(isa(op->getParentOp())); + auto forOp = cast(op->getParentOp()); + }) + .Case([&](auto op) { + const auto pred = op.getQubitIn(); + const auto succ = op.getQubitOut(); + const auto hw = m.at(pred); + m.try_emplace(succ, hw); + }); + + if (!executable) { + return false; } } return true; } +/// Return true, if the entry point fulfills the given coupling constraints. +static bool +isExecutable(func::FuncOp entry, + const DenseSet>& couplingSet) { + DenseMap m; + return isExecutable(entry.getFunctionBody(), m, couplingSet); +} + /// Return a 9x9 square-grid coupling set. static Device getNineQubitSquareGrid() { return {.nqubits = 9, From a7bd6f1d03aa15d7728d964b173a5b9d58e4aa43 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Tue, 23 Jun 2026 08:32:19 +0200 Subject: [PATCH 18/25] Add parallelLoop unit test --- mlir/include/mlir/Dialect/QCO/Utils/Drivers.h | 2 +- .../QCO/Transforms/Mapping/Mapping.cpp | 263 ++++++++++-------- .../QCO/Transforms/Mapping/test_mapping.cpp | 88 +++++- 3 files changed, 227 insertions(+), 126 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h index 9188ece1b6..7c07279ac1 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h @@ -33,7 +33,7 @@ namespace mlir::qco { using ReleasedOps = SmallVector; -using PendingWiresMap = DenseMap>; +using PendingWiresMap = DenseMap>; struct IsReady { bool operator()(PendingWiresMap::value_type& kv) const { diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 603ec9fdf6..4ea101f8b3 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -595,23 +595,16 @@ struct MappingPass : impl::MappingPassBase { return best->layout; } - /** - * @brief Perform A* search to find a sequence of SWAPs that makes the - * two-qubit operations inside the first layer (the front) executable. - * @details - * The iteration budget is b^{3} node expansions, i.e. roughly a depth-3 - * search in a tree with branching factor b. A hard cap prevents impractical - * runtimes on larger architectures. - * - * The branching factor b of the A* search is the product of the - * architecture's maximum qubit degree and the maximum number of two-qubit - * gates in any layer: - * - * b = maxDegree × ⌈N/2⌉ - * - * @returns a vector of hardware-index pairs (each denoting a SWAP) or - * failure() if A* fails. - */ + /// Perform A* search to find a sequence of SWAPs that makes all two-qubit ops + /// inside the first layer executable. + /// + /// The iteration budget is b^{3} node expansions, i.e. roughly a depth-3 + /// search in a tree with branching factor b, where b is the product of the + /// architecture's maximum qubit degree and the maximum number of two-qubit + /// gates in any layer: `b = maxDegree × ⌈N/2⌉`. A hard cap prevents + /// impractical runtimes on larger architectures. + /// + /// Returns `failure`, if the A* search fails. FailureOr> search(const Window& window, const Layout& layout) { constexpr size_t cap = 25'000'000UL; @@ -700,41 +693,36 @@ struct MappingPass : impl::MappingPassBase { /// Return the sequence of SWAPs to move from one layout to another. /// Implements the 4-Approximation algorithm described in arXiv:1602.05150v3. SmallVector restore(const Layout& from, const Layout& to) { - SmallVector swaps; Layout curr(from); - Graph g; - const auto constructEdge = [&](size_t hwX, size_t hwY) { - const auto prog = curr.getProgramIndex(hwX); - + Graph f; + const auto constructEdge = [&](size_t hwA, size_t hwB) { + const auto prog = curr.getProgramIndex(hwA); const auto hwGoal = to.getHardwareIndex(prog); - const auto distPre = device.distanceBetween(hwX, hwGoal); - const auto distPost = device.distanceBetween(hwY, hwGoal); + const auto distPre = device.distanceBetween(hwA, hwGoal); + const auto distPost = device.distanceBetween(hwB, hwGoal); if (distPost < distPre) { - llvm::dbgs() << "prog=" << prog << " edge=(" << hwX << ", " << hwY - << ")" - << " dist(pre)=" << distPre << " dist(post)=" << distPost - << '\n'; - g.addEdge(hwX, hwY); + f.addEdge(hwA, hwB); } }; + SmallVector swaps; do { - // Construct 'F' graph. - g.clear(); + + f.clear(); for (const auto& [hwA, hwB] : device.couplings()) { constructEdge(hwA, hwB); constructEdge(hwB, hwA); } // Find happy swap chain or unhappy swap. - if (const auto cycle = g.findCycle(); cycle) { + + if (const auto cycle = f.findCycle(); cycle) { // Apply happy SWAP chain. + for (size_t i = 0; i < cycle->size() - 1; ++i) { - llvm::dbgs() << "happySWAP=(" << (*cycle)[i] << ", " - << (*cycle)[i + 1] << ")\n"; curr.swap((*cycle)[i], (*cycle)[i + 1]); swaps.emplace_back((*cycle)[i], (*cycle)[i + 1]); } @@ -742,16 +730,14 @@ struct MappingPass : impl::MappingPassBase { continue; } - for (const auto e : g.getEdges()) { - if (g.getDegree(e.second) == 0) { - llvm::dbgs() << "unhappySWAP=(" << e.first << ", " << e.second - << ")\n"; + for (const auto e : f.getEdges()) { + if (f.getDegree(e.second) == 0) { curr.swap(e.first, e.second); swaps.emplace_back(e.first, e.second); break; } } - } while (!g.empty()); + } while (!f.empty()); return swaps; } @@ -845,11 +831,11 @@ struct MappingPass : impl::MappingPassBase { return window; } - /// Insert SWAP operations, exchanging two qubits, virtually (Mode::Cold) or - /// into the IR (Mode::Hot). The function expects that each wire points at the - /// correct insertion point. + /// Insert SWAP operations, exchanging two qubits, virtually + /// (`RoutingMode::Cold`) or into the IR (`RoutingMode::Hot`). The function + /// expects that each wire points at the correct insertion point. template - static void insertSWAPs(ArrayRef swaps, const Wires& wires, + static void insertSWAPs(ArrayRef swaps, Wires& wires, WireInfos& infos, Layout& layout, Statistics& stats, IRRewriter* rewriter) { for (const auto& [hw0, hw1] : swaps) { @@ -859,8 +845,8 @@ struct MappingPass : impl::MappingPassBase { const auto i0 = infos.lookupIndex(prog0); const auto i1 = infos.lookupIndex(prog1); - const auto& w0 = wires[i0]; - const auto& w1 = wires[i1]; + auto& w0 = wires[i0]; + auto& w1 = wires[i1]; const auto in0 = w0.qubit(); const auto in1 = w1.qubit(); @@ -875,6 +861,9 @@ struct MappingPass : impl::MappingPassBase { rewriter->replaceAllUsesExcept(in1, out0, swapOp); infos.swap(prog0, prog1); + + std::advance(w0, 1); // Move to SWAP. + std::advance(w1, 1); } layout.swap(hw0, hw1); @@ -883,94 +872,124 @@ struct MappingPass : impl::MappingPassBase { stats.nswaps += swaps.size(); } - /** - * @brief Advance past all executable gates and recurse into nested regions, - * if necessary. - * @details Traverses the multi-qubit gates of the circuit until no more - * executable gates are found. - */ + /// Advance past all executable gates and recurse into nested regions. + /// Stops when no more executable gates are found. template - void advanceAndRecurse(Wires& wires, const WireInfos& infos, - const Layout& layout, Statistics& stats, - IRRewriter* rewriter) { - walkProgramGraph(wires, [&](const ReadyRange& ready, - ReleasedOps& released) { - if (ready.empty()) { + LogicalResult advanceAndRecurse(Wires& wires, const WireInfos& infos, + const Layout& layout, Statistics& stats, + IRRewriter* rewriter) { + using Traits = WireTraversalTraits; + using StackFrame = std::pair>; + SmallVector stack; + + while (true) { + + // Advance wires past all executable gates and push operations with + // nested regions and the respective wire indices of their inputs onto the + // stack. + + walkProgramGraph(wires, [&](const ReadyRange& ready, + ReleasedOps& released) { + if (ready.empty()) { + return WalkResult::advance(); + } + + for (const auto& [readyOp, indices] : ready) { + TypeSwitch(readyOp) + .template Case( + [&](BarrierOp op) { released.emplace_back(op); }) + .template Case([&](UnitaryOpInterface op) { + const auto prog0 = infos.lookupProgram(indices[0]); + const auto prog1 = infos.lookupProgram(indices[1]); + if (device.areAdjacent( + layout.getHardwareIndices(prog0, prog1))) { + released.emplace_back(op); + } + }) + .template Case( + [&](scf::ForOp op) { stack.emplace_back(op, indices); }); + } + + if (released.empty()) { + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + + if (stack.empty()) { + break; } - for (const auto& [readyOp, indices] : ready) { - TypeSwitch(readyOp) - .template Case( - [&](BarrierOp op) { released.emplace_back(op); }) - .template Case([&](UnitaryOpInterface op) { - const auto prog0 = infos.lookupProgram(indices[0]); - const auto prog1 = infos.lookupProgram(indices[1]); - if (device.areAdjacent(layout.getHardwareIndices(prog0, prog1))) { - released.emplace_back(op); - } - }) - .template Case([&](scf::ForOp op) { - Wires childWires; - WireInfos childInfos; - - // First map parent (results) to child values (iteration - // arguments). Going forwards, the recursive routing starts at - // block arguments, while the backwards go starts at the yielded - // values. - - for (auto i : indices) { - const auto prog = infos.lookupProgram(i); - const auto arg = op.getTiedLoopRegionIterArg( - cast(wires[i].qubit())); - const auto index = childWires.size(); - - if constexpr (Direction == WireDirection::Forward) { - childWires.emplace_back(arg); - childInfos.map(index, prog); - } else { - const auto yield = op.getTiedLoopYieldedValue(arg)->get(); - childWires.emplace_back(yield); - childInfos.map(index, prog); - } - } + // The wires now point at the results of non-executable gates or + // operations with nested regions. Continue with processing the nested + // regions recursively. + + for (const auto& [op, indices] : stack) { + assert(isa(op)); + auto forOp = cast(op); + + Wires childWires; + WireInfos childInfos; + + // Map parent (results) to child values (iter args). Going forwards, the + // recursive routing starts at block arguments, while the backwards go + // starts at the yielded values. + + for (size_t i : indices) { + const auto prog = infos.lookupProgram(i); + const auto res = cast(wires[i].qubit()); + const auto arg = forOp.getTiedLoopRegionIterArg(res); + const auto index = childWires.size(); + + if constexpr (Direction == WireDirection::Forward) { + childWires.emplace_back(arg); + childInfos.map(index, prog); + } else { + const auto yield = forOp.getTiedLoopYieldedValue(arg)->get(); + childWires.emplace_back(yield); + childInfos.map(index, prog); + } + } - Layout childLayout(layout); - const auto res = route( - childWires, childInfos, childLayout, stats, rewriter); - if (failed(res)) { - // TODO: How to propagate errors here? - } + Layout childLayout(layout); + const auto res = route(childWires, childInfos, + childLayout, stats, rewriter); + if (failed(res)) { + return failure(); + } - if constexpr (Mode == RoutingMode::Hot) { + const auto swaps = restore(childLayout, layout); - // After routing the loop body, all iterators point to - // std::default_sentinel. To move the iterators to the correct - // qubit SSA values for the epilogue SWAPs, decrement each - // twice: (sentinel → yield → unitary/block arg). + if constexpr (Mode == RoutingMode::Hot) { - llvm::for_each(childWires, - [](auto& it) { std::advance(it, -2); }); - } + // After routing the loop body, all iterators point to + // std::default_sentinel. To move the iterators to the correct + // qubit SSA values for the epilogue SWAPs, decrement each + // twice: (sentinel → yield → unitary/block arg). - insertSWAPs(restore(childLayout, layout), childWires, - childInfos, childLayout, stats, rewriter); + llvm::for_each(childWires, [](auto& it) { std::advance(it, -2); }); + } - if constexpr (Mode == RoutingMode::Hot) { - sortTopologically(op.getBody()); - } + insertSWAPs(swaps, childWires, childInfos, childLayout, stats, + rewriter); - released.emplace_back(op); - }); - } + if constexpr (Mode == RoutingMode::Hot) { + sortTopologically(forOp.getBody()); + } - // Stop, if there are no more ready AND executable gates. - if (released.empty()) { - return WalkResult::interrupt(); + // Finally, move past the operation with nested regions by incrementing + // the respective global wires. + + llvm::for_each(indices, [&](size_t i) { + std::advance(wires[i], Traits::stride()); + }); } - return WalkResult::advance(); - }); + stack.clear(); + } + + return success(); } /// Iterates over a dynamically computed window of layers and uses A* search @@ -983,7 +1002,11 @@ struct MappingPass : impl::MappingPassBase { LogicalResult route(Wires& wires, WireInfos& infos, Layout& layout, Statistics& stats, IRRewriter* rewriter = nullptr) { while (true) { - advanceAndRecurse(wires, infos, layout, stats, rewriter); + const auto res = advanceAndRecurse(wires, infos, layout, + stats, rewriter); + if (failed(res)) { + return failure(); + } const auto window = getWindow(wires, infos); if (window.empty()) { diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index f1b6a7e88e..4724f3e5f4 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -70,7 +70,8 @@ isExecutable(Region& body, DenseMap& m, const auto hwA = m.at(op.getInputQubit(0)); const auto hwB = m.at(op.getInputQubit(1)); if (!couplingSet.contains(std::make_pair(hwA, hwB))) { - llvm::dbgs() << "not executable: \n"; + llvm::dbgs() << "(" << hwA << ", " << hwB << ") " + << "not executable: \n"; op->dump(); executable = false; } @@ -333,9 +334,7 @@ TEST_P(MappingPassTest, GroverLike) { const auto& device = GetParam(); PassManager pm(context.get()); - pm.addPass(createMappingPass( - device.couplingSet, - MappingPassOptions{.ntrials = 1, .niterations = 1, .nlookahead = 2})); + pm.addPass(createMappingPass(device.couplingSet, MappingPassOptions{})); QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -414,7 +413,86 @@ TEST_P(MappingPassTest, GroverLike) { auto res = pm.run(m.get()); auto entry = getEntryPoint(m.get()); - m->dump(); + ASSERT_TRUE(res.succeeded()); + EXPECT_TRUE(isExecutable(entry, device.couplingSet)); +} + +TEST_P(MappingPassTest, ParallelLoops) { + constexpr int64_t nqubits = 6; + const auto& device = GetParam(); + + PassManager pm(context.get()); + pm.addPass( + createMappingPass(device.couplingSet, MappingPassOptions{.ntrials = 1, .niterations=1})); + + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + Value tensor = builder.qtensorAlloc(nqubits); + SmallVector creg(nqubits); + SmallVector qreg(nqubits); + + for (int64_t i = 0; i < nqubits; ++i) { + std::tie(tensor, qreg[i]) = builder.qtensorExtract(tensor, i); + qreg[i] = builder.h(qreg[i]); + } + + const auto upForResults = + builder.scfFor(1, 3, 1, {qreg[0], qreg[1], qreg[2]}, + [&builder](Value, ValueRange iterArgs) { + Value iterQ0 = iterArgs[0]; + Value iterQ1 = iterArgs[1]; + Value iterQ2 = iterArgs[2]; + + std::tie(iterQ0, iterQ1) = builder.cx(iterQ0, iterQ1); + iterQ0 = builder.h(iterQ0); + std::tie(iterQ0, iterQ1) = builder.cz(iterQ0, iterQ1); + std::tie(iterQ1, iterQ2) = builder.cz(iterQ1, iterQ2); + std::tie(iterQ0, iterQ2) = builder.cx(iterQ0, iterQ2); + + return SmallVector{iterQ0, iterQ1, iterQ2}; + }); + + qreg[0] = upForResults[0]; + qreg[1] = upForResults[1]; + qreg[2] = upForResults[2]; + + const auto downForResults = + builder.scfFor(1, 3, 1, {qreg[3], qreg[4], qreg[5]}, + [&builder](Value, ValueRange iterArgs) { + Value iterQ0 = iterArgs[0]; + Value iterQ1 = iterArgs[1]; + Value iterQ2 = iterArgs[2]; + + std::tie(iterQ0, iterQ1) = builder.cx(iterQ0, iterQ1); + iterQ0 = builder.h(iterQ0); + std::tie(iterQ1, iterQ2) = builder.cz(iterQ1, iterQ2); + std::tie(iterQ0, iterQ1) = builder.cz(iterQ0, iterQ1); + std::tie(iterQ0, iterQ2) = builder.cx(iterQ0, iterQ2); + + return SmallVector{iterQ0, iterQ1, iterQ2}; + }); + + qreg[3] = downForResults[0]; + qreg[4] = downForResults[1]; + qreg[5] = downForResults[2]; + + qreg = builder.barrier(qreg); + + for (int64_t i = 0; i < nqubits; ++i) { + std::tie(qreg[i], creg[i]) = builder.measure(qreg[i]); + qreg[i] = builder.h(qreg[i]); + } + + for (int64_t i = 0; i < nqubits; ++i) { + tensor = builder.qtensorInsert(qreg[i], tensor, i); + } + + builder.qtensorDealloc(tensor); + + auto m = builder.finalize(); + auto res = pm.run(m.get()); + auto entry = getEntryPoint(m.get()); ASSERT_TRUE(res.succeeded()); EXPECT_TRUE(isExecutable(entry, device.couplingSet)); From 5bb508b1bef087f3ec96b76cca8442346fb8898b Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Tue, 23 Jun 2026 13:47:47 +0200 Subject: [PATCH 19/25] Minor cleanup --- mlir/include/mlir/Dialect/QCO/Utils/Graph.h | 5 +- .../QCO/Transforms/Mapping/Mapping.cpp | 76 +++++++++++-------- 2 files changed, 47 insertions(+), 34 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Graph.h b/mlir/include/mlir/Dialect/QCO/Utils/Graph.h index 1409158f79..a09c7cb934 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Graph.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Graph.h @@ -128,9 +128,9 @@ template class Graph { return dist; } - /// Return reverse cycle in graph or `std::nullopt` if none exists. + /// Return cycle in graph or `std::nullopt` if none exists. /// Implements an iterative depth-first search inspired by LLVM's SCC - /// utilities. + /// utilities. For a cycle [A, B, C, A], the function returns [A, B, C]. [[nodiscard]] std::optional> findCycle() const { enum struct State : uint8_t { Unseen, Seen, Finished }; @@ -189,6 +189,7 @@ template class Graph { path.emplace_back(curr); } path.emplace_back(nbrId); + std::ranges::reverse(path); return path; } } diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 4ea101f8b3..b39c3246a2 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -47,6 +47,7 @@ #include #include +#include #include #include #include @@ -128,33 +129,34 @@ struct MappingPass : impl::MappingPassBase { struct WireInfos { /// Return the mapped wire index of a program index. [[nodiscard]] size_t lookupIndex(size_t prog) const { - return programToIndex.at(prog); + return programToIndex_.at(prog); } /// Return the mapped program index of a wire index. [[nodiscard]] size_t lookupProgram(size_t index) const { - return indexToProgram.at(index); + return indexToProgram_.at(index); } /// Bidirectionally map a wire index to a program index. /// Overwrites existing mappings. void map(size_t index, size_t prog) { - indexToProgram[index] = prog; - programToIndex[prog] = index; + indexToProgram_[index] = prog; + programToIndex_[prog] = index; } /// Swap two program indices. void swap(size_t prog0, size_t prog1) { const auto i0 = lookupIndex(prog0); const auto i1 = lookupIndex(prog1); - std::swap(programToIndex[prog0], programToIndex[prog1]); - std::swap(indexToProgram[i0], indexToProgram[i1]); + std::swap(programToIndex_[prog0], programToIndex_[prog1]); + std::swap(indexToProgram_[i0], indexToProgram_[i1]); } + private: /// Maps the i-th wire index to a program index. - DenseMap indexToProgram; + DenseMap indexToProgram_; /// Maps a program index to the i-th wire index. - DenseMap programToIndex; + DenseMap programToIndex_; }; /// Statistics collected while routing. @@ -693,47 +695,57 @@ struct MappingPass : impl::MappingPassBase { /// Return the sequence of SWAPs to move from one layout to another. /// Implements the 4-Approximation algorithm described in arXiv:1602.05150v3. SmallVector restore(const Layout& from, const Layout& to) { + static constexpr size_t MIN_CYCLE_LENGTH = 2; Layout curr(from); - + SmallVector swaps; Graph f; - const auto constructEdge = [&](size_t hwA, size_t hwB) { + + const auto shouldAddEdge = [&](size_t hwA, size_t hwB) { const auto prog = curr.getProgramIndex(hwA); const auto hwGoal = to.getHardwareIndex(prog); - const auto distPre = device.distanceBetween(hwA, hwGoal); - const auto distPost = device.distanceBetween(hwB, hwGoal); - - if (distPost < distPre) { - f.addEdge(hwA, hwB); - } + return device.distanceBetween(hwB, hwGoal) < + device.distanceBetween(hwA, hwGoal); }; - SmallVector swaps; do { - f.clear(); - for (const auto& [hwA, hwB] : device.couplings()) { - constructEdge(hwA, hwB); - constructEdge(hwB, hwA); - } - // Find happy swap chain or unhappy swap. + // Build F-graph: add edges for both directions of each coupling + for (const auto& coupling : device.couplings()) { + const std::array, 2> directedEdges = { + {{coupling.first, coupling.second}, + {coupling.second, coupling.first}}}; - if (const auto cycle = f.findCycle(); cycle) { - // Apply happy SWAP chain. + for (const auto& [hwA, hwB] : directedEdges) { + if (shouldAddEdge(hwA, hwB)) { + f.addEdge(hwA, hwB); + } + } + } + + // Try to find a directed cycle in the F graph. If there is one, + // we can apply a happy swap chain. Note that this happy swap chain + // does not include the final back edge closing the cycle because the + // first SWAP changes the token (the qubit) on the target, invalidating + // the edge in F. - for (size_t i = 0; i < cycle->size() - 1; ++i) { + if (const auto cycle = f.findCycle(); + cycle && cycle->size() >= MIN_CYCLE_LENGTH) { + for (size_t i = 0; i + 1 < cycle->size(); ++i) { curr.swap((*cycle)[i], (*cycle)[i + 1]); swaps.emplace_back((*cycle)[i], (*cycle)[i + 1]); } - continue; } - for (const auto e : f.getEdges()) { - if (f.getDegree(e.second) == 0) { - curr.swap(e.first, e.second); - swaps.emplace_back(e.first, e.second); + // To be benchmarked: f.getEdges() is potentially to expensive because it builds a + // llvm::DenseMap. + + for (const auto [u, v] : f.getEdges()) { + if (f.getDegree(v) == 0) { + curr.swap(u, v); + swaps.emplace_back(u, v); break; } } @@ -880,8 +892,8 @@ struct MappingPass : impl::MappingPassBase { IRRewriter* rewriter) { using Traits = WireTraversalTraits; using StackFrame = std::pair>; - SmallVector stack; + SmallVector stack; while (true) { // Advance wires past all executable gates and push operations with From d092079e95f354f8e4d0059bf603f50be79501e5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 25 Jun 2026 06:25:21 +0000 Subject: [PATCH 20/25] =?UTF-8?q?=F0=9F=8E=A8=20pre-commit=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mlir/include/mlir/Dialect/QCO/Utils/Layout.h | 12 +++++++++++- mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h | 6 ++++-- mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp | 4 ++-- mlir/lib/Dialect/QCO/Utils/Layout.cpp | 12 +++++++++++- .../Dialect/QCO/Transforms/Mapping/test_mapping.cpp | 4 ++-- 5 files changed, 30 insertions(+), 8 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Layout.h b/mlir/include/mlir/Dialect/QCO/Utils/Layout.h index 7b8948e835..45d0a77ee0 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Layout.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Layout.h @@ -1,3 +1,13 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + #pragma once #include @@ -65,4 +75,4 @@ class Layout { explicit Layout(const size_t nqubits) : programToHardware_(nqubits), hardwareToProgram_(nqubits) {} }; -} // namespace mlir::qco \ No newline at end of file +} // namespace mlir::qco diff --git a/mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h b/mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h index cc5bef551b..5076d93643 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h @@ -32,9 +32,11 @@ class [[nodiscard]] WireIterator { using difference_type = std::ptrdiff_t; using value_type = Operation*; - WireIterator() : op_(nullptr), qubit_(nullptr), isFinal_(false), isSentinel_(false) {} + WireIterator() + : op_(nullptr), qubit_(nullptr), isFinal_(false), isSentinel_(false) {} explicit WireIterator(Value qubit) - : op_(qubit.getDefiningOp()), qubit_(qubit), isFinal_(false), isSentinel_(false) {} + : op_(qubit.getDefiningOp()), qubit_(qubit), isFinal_(false), + isSentinel_(false) {} /// @returns the operation the iterator points to. [[nodiscard]] Operation* operation() const { return op_; } diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index b39c3246a2..40a9b6fa28 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -739,8 +739,8 @@ struct MappingPass : impl::MappingPassBase { continue; } - // To be benchmarked: f.getEdges() is potentially to expensive because it builds a - // llvm::DenseMap. + // To be benchmarked: f.getEdges() is potentially to expensive because it + // builds a llvm::DenseMap. for (const auto [u, v] : f.getEdges()) { if (f.getDegree(v) == 0) { diff --git a/mlir/lib/Dialect/QCO/Utils/Layout.cpp b/mlir/lib/Dialect/QCO/Utils/Layout.cpp index e993d97a9b..8e505f0223 100644 --- a/mlir/lib/Dialect/QCO/Utils/Layout.cpp +++ b/mlir/lib/Dialect/QCO/Utils/Layout.cpp @@ -1,3 +1,13 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + #include "mlir/Dialect/QCO/Utils/Layout.h" #include @@ -53,4 +63,4 @@ ArrayRef Layout::getProgramToHardware() const { return programToHardware_; } -} // namespace mlir::qco \ No newline at end of file +} // namespace mlir::qco diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 4724f3e5f4..7cc7742068 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -422,8 +422,8 @@ TEST_P(MappingPassTest, ParallelLoops) { const auto& device = GetParam(); PassManager pm(context.get()); - pm.addPass( - createMappingPass(device.couplingSet, MappingPassOptions{.ntrials = 1, .niterations=1})); + pm.addPass(createMappingPass( + device.couplingSet, MappingPassOptions{.ntrials = 1, .niterations = 1})); QCOProgramBuilder builder(context.get()); builder.initialize(); From 995789a45e89e56eae36d2154edc771826698a1d Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Thu, 25 Jun 2026 09:04:31 +0200 Subject: [PATCH 21/25] Move recursive logic from advance to route --- .../QCO/Transforms/Mapping/Mapping.cpp | 266 +++++++++--------- 1 file changed, 132 insertions(+), 134 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index b39c3246a2..75e6a93b61 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -170,6 +170,13 @@ struct MappingPass : impl::MappingPassBase { float lambda; }; + /// Utility-struct for routing functions. + struct RoutingBundle { + Wires wires; + WireInfos infos; + Layout layout; + }; + /// Describes a node in the A* search graph. struct Node { struct ComparePointer { @@ -304,10 +311,13 @@ struct MappingPass : impl::MappingPassBase { std::tie(wires, infos) = std::move(place(body, *layout, rewriter)); - Statistics s; + Statistics stats; + RoutingBundle bundle{.infos = std::move(infos), + .wires = std::move(wires), + .layout = std::move(*layout)}; const auto res = route( - wires, infos, *layout, s, &rewriter); + bundle, stats, &rewriter); if (res.failed()) { func.emitError() << "failed to map the function"; signalPassFailure(); @@ -315,7 +325,7 @@ struct MappingPass : impl::MappingPassBase { } // Collect statistics. - numSwaps += s.nswaps; + numSwaps += stats.nswaps; // Fix SSA Dominance issues. llvm::for_each(body.getBlocks(), [](Block& b) { sortTopologically(&b); }); @@ -552,29 +562,26 @@ struct MappingPass : impl::MappingPassBase { std::mt19937_64 rng{seed}; struct Trial { - Layout layout; + RoutingBundle bundle; Statistics stats{}; bool success{false}; }; - SmallVector trials; + SmallVector trials; trials.reserve(ntrials); for (size_t i = 0; i < ntrials; ++i) { - trials.emplace_back(Layout::random(device.nqubits(), rng())); + trials.emplace_back( + RoutingBundle{.wires = wires, + .infos = infos, + .layout = Layout::random(device.nqubits(), rng())}); } parallelForEach(&getContext(), trials, [&, this](Trial& t) { - Wires locWires(wires); - WireInfos locInfos(infos); - for (size_t i = 0; i < niterations; ++i) { - if (route(locWires, locInfos, t.layout, t.stats) - .failed()) { + if (route(t.bundle, t.stats).failed()) { return; } - if (route(locWires, locInfos, t.layout, - t.stats) - .failed()) { + if (route(t.bundle, t.stats).failed()) { return; } } @@ -594,7 +601,7 @@ struct MappingPass : impl::MappingPassBase { return failure(); } - return best->layout; + return best->bundle.layout; } /// Perform A* search to find a sequence of SWAPs that makes all two-qubit ops @@ -739,8 +746,8 @@ struct MappingPass : impl::MappingPassBase { continue; } - // To be benchmarked: f.getEdges() is potentially to expensive because it builds a - // llvm::DenseMap. + // To be benchmarked: f.getEdges() is potentially to expensive because it + // builds a llvm::DenseMap. for (const auto [u, v] : f.getEdges()) { if (f.getDegree(v) == 0) { @@ -847,9 +854,9 @@ struct MappingPass : impl::MappingPassBase { /// (`RoutingMode::Cold`) or into the IR (`RoutingMode::Hot`). The function /// expects that each wire points at the correct insertion point. template - static void insertSWAPs(ArrayRef swaps, Wires& wires, - WireInfos& infos, Layout& layout, Statistics& stats, - IRRewriter* rewriter) { + static void insertSWAPs(ArrayRef swaps, RoutingBundle& bundle, + Statistics& stats, IRRewriter* rewriter) { + auto& [wires, infos, layout] = bundle; for (const auto& [hw0, hw1] : swaps) { if constexpr (Mode == RoutingMode::Hot) { const auto [prog0, prog1] = layout.getProgramIndices(hw0, hw1); @@ -884,140 +891,131 @@ struct MappingPass : impl::MappingPassBase { stats.nswaps += swaps.size(); } - /// Advance past all executable gates and recurse into nested regions. - /// Stops when no more executable gates are found. - template - LogicalResult advanceAndRecurse(Wires& wires, const WireInfos& infos, - const Layout& layout, Statistics& stats, - IRRewriter* rewriter) { - using Traits = WireTraversalTraits; - using StackFrame = std::pair>; + /// Advance past all executable gates and return operations with nested + /// regions and the respecitve wire indices. Stops when no more executable + /// gates are found. After the function returns, the wires point at the + /// results of non-executable gates or operations with nested regions. + template + SmallVector>> + advance(Wires& wires, const WireInfos& infos, const Layout& layout) { + SmallVector>> stack; - SmallVector stack; - while (true) { + // Advance wires past all executable gates and push operations with + // nested regions and the respective wire indices of their inputs onto the + // result stack. - // Advance wires past all executable gates and push operations with - // nested regions and the respective wire indices of their inputs onto the - // stack. + walkProgramGraph(wires, [&](const ReadyRange& ready, + ReleasedOps& released) { + if (ready.empty()) { + return WalkResult::advance(); + } - walkProgramGraph(wires, [&](const ReadyRange& ready, - ReleasedOps& released) { - if (ready.empty()) { - return WalkResult::advance(); - } + for (const auto& [readyOp, indices] : ready) { + TypeSwitch(readyOp) + .template Case( + [&](BarrierOp op) { released.emplace_back(op); }) + .template Case([&](UnitaryOpInterface op) { + const auto prog0 = infos.lookupProgram(indices[0]); + const auto prog1 = infos.lookupProgram(indices[1]); + if (device.areAdjacent(layout.getHardwareIndices(prog0, prog1))) { + released.emplace_back(op); + } + }) + .template Case( + [&](scf::ForOp op) { stack.emplace_back(op, indices); }); + } - for (const auto& [readyOp, indices] : ready) { - TypeSwitch(readyOp) - .template Case( - [&](BarrierOp op) { released.emplace_back(op); }) - .template Case([&](UnitaryOpInterface op) { - const auto prog0 = infos.lookupProgram(indices[0]); - const auto prog1 = infos.lookupProgram(indices[1]); - if (device.areAdjacent( - layout.getHardwareIndices(prog0, prog1))) { - released.emplace_back(op); - } - }) - .template Case( - [&](scf::ForOp op) { stack.emplace_back(op, indices); }); - } + if (released.empty()) { + return WalkResult::interrupt(); + } - if (released.empty()) { - return WalkResult::interrupt(); - } + return WalkResult::advance(); + }); - return WalkResult::advance(); - }); + return stack; + } - if (stack.empty()) { - break; - } + /// Iterates over a dynamically computed window of layers and uses A* search + /// to find a SWAP sequence that makes each layer executable. Depending on + /// the template parameter, this function only updates the layout or also + /// inserts the SWAPs into the IR. The function returns `failure` if A* is + /// unable to find a solution. + template + requires(Mode != RoutingMode::Hot || Direction == WireDirection::Forward) + LogicalResult route(RoutingBundle& bundle, Statistics& stats, + IRRewriter* rewriter = nullptr) { + using Traits = WireTraversalTraits; - // The wires now point at the results of non-executable gates or - // operations with nested regions. Continue with processing the nested - // regions recursively. - - for (const auto& [op, indices] : stack) { - assert(isa(op)); - auto forOp = cast(op); - - Wires childWires; - WireInfos childInfos; - - // Map parent (results) to child values (iter args). Going forwards, the - // recursive routing starts at block arguments, while the backwards go - // starts at the yielded values. - - for (size_t i : indices) { - const auto prog = infos.lookupProgram(i); - const auto res = cast(wires[i].qubit()); - const auto arg = forOp.getTiedLoopRegionIterArg(res); - const auto index = childWires.size(); - - if constexpr (Direction == WireDirection::Forward) { - childWires.emplace_back(arg); - childInfos.map(index, prog); - } else { - const auto yield = forOp.getTiedLoopYieldedValue(arg)->get(); - childWires.emplace_back(yield); - childInfos.map(index, prog); - } - } + auto& [wires, infos, layout] = bundle; + + while (true) { - Layout childLayout(layout); - const auto res = route(childWires, childInfos, - childLayout, stats, rewriter); - if (failed(res)) { - return failure(); + while (true) { + const auto stack = advance(wires, infos, layout); + + if (stack.empty()) { + break; } - const auto swaps = restore(childLayout, layout); + // Continue with processing the nested regions recursively. - if constexpr (Mode == RoutingMode::Hot) { + for (const auto& [op, indices] : stack) { + assert(isa(op)); + auto forOp = cast(op); - // After routing the loop body, all iterators point to - // std::default_sentinel. To move the iterators to the correct - // qubit SSA values for the epilogue SWAPs, decrement each - // twice: (sentinel → yield → unitary/block arg). + RoutingBundle child{.layout = layout}; - llvm::for_each(childWires, [](auto& it) { std::advance(it, -2); }); - } + // Map parent (results) to child values (iter args). Going forwards, + // the recursive routing starts at block arguments, while the + // backwards go starts at the yielded values. - insertSWAPs(swaps, childWires, childInfos, childLayout, stats, - rewriter); + for (size_t i : indices) { + const auto prog = infos.lookupProgram(i); + const auto res = cast(wires[i].qubit()); + const auto arg = forOp.getTiedLoopRegionIterArg(res); + const auto index = child.wires.size(); - if constexpr (Mode == RoutingMode::Hot) { - sortTopologically(forOp.getBody()); - } + if constexpr (Direction == WireDirection::Forward) { + child.wires.emplace_back(arg); + child.infos.map(index, prog); + } else { + const auto yield = forOp.getTiedLoopYieldedValue(arg)->get(); + child.wires.emplace_back(yield); + child.infos.map(index, prog); + } + } - // Finally, move past the operation with nested regions by incrementing - // the respective global wires. + Layout childLayout(layout); + const auto res = route(child, stats, rewriter); + if (failed(res)) { + return failure(); + } - llvm::for_each(indices, [&](size_t i) { - std::advance(wires[i], Traits::stride()); - }); - } + const auto swaps = restore(childLayout, layout); - stack.clear(); - } + if constexpr (Mode == RoutingMode::Hot) { - return success(); - } + // After routing the loop body, all iterators point to + // std::default_sentinel. To move the iterators to the correct + // qubit SSA values for the epilogue SWAPs, decrement each + // twice: (sentinel → yield → unitary/block arg). - /// Iterates over a dynamically computed window of layers and uses A* search - /// to find a SWAP sequence that makes each layer executable. Depending on - /// the template parameter, this function only updates the layout or also - /// inserts the SWAPs into the IR. The function returns `failure` if A* is - /// unable to find a solution. - template - requires(Mode != RoutingMode::Hot || Direction == WireDirection::Forward) - LogicalResult route(Wires& wires, WireInfos& infos, Layout& layout, - Statistics& stats, IRRewriter* rewriter = nullptr) { - while (true) { - const auto res = advanceAndRecurse(wires, infos, layout, - stats, rewriter); - if (failed(res)) { - return failure(); + llvm::for_each(child.wires, [](auto& it) { std::advance(it, -2); }); + } + + insertSWAPs(swaps, child, stats, rewriter); + + if constexpr (Mode == RoutingMode::Hot) { + sortTopologically(forOp.getBody()); + } + + // Finally, move past the operation with nested regions by + // incrementing the respective global wires. + + llvm::for_each(indices, [&](size_t i) { + std::advance(wires[i], Traits::stride()); + }); + } } const auto window = getWindow(wires, infos); @@ -1043,7 +1041,7 @@ struct MappingPass : impl::MappingPassBase { } } - insertSWAPs(*swaps, wires, infos, layout, stats, rewriter); + insertSWAPs(*swaps, bundle, stats, rewriter); if constexpr (Mode == RoutingMode::Hot) { From 49efdaa3cbaa616f70fa9a0328d92b4a8672c63f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 25 Jun 2026 07:05:05 +0000 Subject: [PATCH 22/25] =?UTF-8?q?=F0=9F=8E=A8=20pre-commit=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 75e6a93b61..923c12be81 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -892,7 +892,7 @@ struct MappingPass : impl::MappingPassBase { } /// Advance past all executable gates and return operations with nested - /// regions and the respecitve wire indices. Stops when no more executable + /// regions and the respective wire indices. Stops when no more executable /// gates are found. After the function returns, the wires point at the /// results of non-executable gates or operations with nested regions. template From 162781e20b94cd8ffa7fcf45e9dc10afe6e95580 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Thu, 25 Jun 2026 09:10:03 +0200 Subject: [PATCH 23/25] Use correct layout in nested handling --- mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 75e6a93b61..b9afb5c102 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -26,7 +26,6 @@ #include #include #include -#include #include #include #include @@ -985,13 +984,12 @@ struct MappingPass : impl::MappingPassBase { } } - Layout childLayout(layout); const auto res = route(child, stats, rewriter); if (failed(res)) { return failure(); } - const auto swaps = restore(childLayout, layout); + const auto swaps = restore(child.layout, layout); if constexpr (Mode == RoutingMode::Hot) { From f46bad427b6f92e6b0b41041f38a572900bd5150 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Thu, 25 Jun 2026 09:33:41 +0200 Subject: [PATCH 24/25] Try fix builds --- mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp | 12 +++++++----- .../Dialect/QCO/Transforms/Mapping/test_mapping.cpp | 2 +- mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp | 1 - 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 63fc11d6ee..28772e3dca 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -89,8 +89,8 @@ struct MappingPass : impl::MappingPassBase { [[nodiscard]] size_t nqubits() const { return coupling_.getNumNodes(); } /// Return true if two qubits are adjacent. - [[nodiscard]] bool areAdjacent(std::pair qubits) const { - return dist_[qubits.first][qubits.second] == 1UL; + [[nodiscard]] bool areAdjacent(size_t u, size_t v) const { + return dist_[u][v] == 1UL; } /// Return the length of the shortest path between two qubits. @@ -215,8 +215,9 @@ struct MappingPass : impl::MappingPassBase { */ [[nodiscard]] bool isGoal(const IndexPairType& front, const AugmentedDevice& device) const { - return device.areAdjacent( - layout.getHardwareIndices(front.first, front.second)); + const auto [hw0, hw1] = + layout.getHardwareIndices(front.first, front.second); + return device.areAdjacent(hw0, hw1); } private: @@ -916,7 +917,8 @@ struct MappingPass : impl::MappingPassBase { .template Case([&](UnitaryOpInterface op) { const auto prog0 = infos.lookupProgram(indices[0]); const auto prog1 = infos.lookupProgram(indices[1]); - if (device.areAdjacent(layout.getHardwareIndices(prog0, prog1))) { + const auto [hw0, hw1] = layout.getHardwareIndices(prog0, prog1); + if (device.areAdjacent(hw0, hw1)) { released.emplace_back(op); } }) diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 7cc7742068..4bafb33b56 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -423,7 +423,7 @@ TEST_P(MappingPassTest, ParallelLoops) { PassManager pm(context.get()); pm.addPass(createMappingPass( - device.couplingSet, MappingPassOptions{.ntrials = 1, .niterations = 1})); + device.couplingSet, MappingPassOptions{})); QCOProgramBuilder builder(context.get()); builder.initialize(); diff --git a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp index 9b6c5d2c1a..80c40f49a1 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp @@ -12,7 +12,6 @@ #include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/Utils/Drivers.h" -#include "mlir/Dialect/QCO/Utils/Qubits.h" #include "mlir/Dialect/QCO/Utils/WireIterator.h" #include From 1af74e0477b61fb25a1c193f6f2444476d7266b8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 25 Jun 2026 07:34:15 +0000 Subject: [PATCH 25/25] =?UTF-8?q?=F0=9F=8E=A8=20pre-commit=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 4bafb33b56..f5392c4f2b 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -422,8 +422,7 @@ TEST_P(MappingPassTest, ParallelLoops) { const auto& device = GetParam(); PassManager pm(context.get()); - pm.addPass(createMappingPass( - device.couplingSet, MappingPassOptions{})); + pm.addPass(createMappingPass(device.couplingSet, MappingPassOptions{})); QCOProgramBuilder builder(context.get()); builder.initialize();