diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a6a60e02c..5a37ea0ab0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,8 +40,8 @@ releases may include breaking changes. [#1676], [#1706], [#1776]) ([**@denialhaag**], [**@burgholzer**]) - ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568], [#1581], [#1583], [#1588], - [#1600], [#1664], [#1709], [#1716], [#1748], [#1805]) ([**@MatthiasReumann**], - [**@burgholzer**]) + [#1600], [#1664], [#1709], [#1716], [#1748], [#1805], [#1870]) + ([**@MatthiasReumann**], [**@burgholzer**]) - ✨ Add initial infrastructure for new QC and QCO MLIR dialects ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], @@ -630,6 +630,7 @@ changelogs._ [#1877]: https://github.com/munich-quantum-toolkit/core/pull/1877 [#1873]: https://github.com/munich-quantum-toolkit/core/pull/1873 [#1872]: https://github.com/munich-quantum-toolkit/core/pull/1872 +[#1870]: https://github.com/munich-quantum-toolkit/core/pull/1870 [#1869]: https://github.com/munich-quantum-toolkit/core/pull/1869 [#1850]: https://github.com/munich-quantum-toolkit/core/pull/1850 [#1849]: https://github.com/munich-quantum-toolkit/core/pull/1849 diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Graph.h b/mlir/include/mlir/Dialect/QCO/Utils/Graph.h index f7e5096808..6cd25c8144 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Graph.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Graph.h @@ -75,7 +75,9 @@ class Graph { [[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(); } + [[nodiscard]] size_t getDegree(const size_t id) const { + return adj_.at(id).size(); + } /// Return the max degree of the graph. [[nodiscard]] size_t getMaxDegree() const; diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp index 1cd45a58f6..c232728250 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp @@ -313,14 +313,14 @@ BlockArgument IfOp::getTiedElseBlockArgument(OpOperand* qubit) { } OpOperand* IfOp::getTiedThenYieldedValue(BlockArgument bbArg) { - if (bbArg.getDefiningOp() != getOperation()) { + if (bbArg.getOwner()->getParentOp() != getOperation()) { return nullptr; } return &thenYield().getTargetsMutable()[bbArg.getArgNumber()]; } OpOperand* IfOp::getTiedElseYieldedValue(BlockArgument bbArg) { - if (bbArg.getDefiningOp() != getOperation()) { + if (bbArg.getOwner()->getParentOp() != getOperation()) { return nullptr; } return &elseYield().getTargetsMutable()[bbArg.getArgNumber()]; diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index dd0ab588ff..e93c14b180 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -51,6 +51,7 @@ #include #include #include +#include #include #include #include @@ -73,25 +74,27 @@ struct MappingPass : impl::MappingPassBase { using IndexPairType = std::pair; using Window = SmallVector; using Wires = SmallVector; + using RecursiveRoutingStackItem = std::pair>; + using RecursiveRoutingStack = SmallVector; enum class RoutingMode : bool { Cold, Hot }; class AugmentedDevice { public: explicit AugmentedDevice( - const llvm::DenseSet>& couplingSet) + const DenseSet>& couplingSet) : coupling_(couplingSet), dist_(coupling_.getDistMatrix()) {} /// Return the device's number of qubits. [[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 { + [[nodiscard]] bool areAdjacent(const size_t u, const size_t v) const { return dist_[u][v] == 1UL; } /// Return the length of the shortest path between two qubits. - [[nodiscard]] size_t distanceBetween(size_t u, size_t v) const { + [[nodiscard]] size_t distanceBetween(const size_t u, const size_t v) const { const auto dist = dist_[u][v]; if (dist == UINT64_MAX) { report_fatal_error("Failed to compute the distance between qubits " + @@ -106,7 +109,7 @@ struct MappingPass : impl::MappingPassBase { } /// Return all neighbours of a qubit. - [[nodiscard]] ArrayRef neighboursOf(size_t u) const { + [[nodiscard]] ArrayRef neighboursOf(const size_t u) const { return coupling_.getNeighbours(u); } @@ -120,24 +123,30 @@ 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); + [[nodiscard]] size_t lookupIndex(const size_t prog) const { + return programToIndex_[prog]; } /// Return the mapped program index of a wire index. - [[nodiscard]] size_t lookupProgram(size_t index) const { - return indexToProgram_.at(index); + [[nodiscard]] size_t lookupProgram(const size_t index) const { + return indexToProgram_[index]; } /// Bidirectionally map a wire index to a program index. /// Overwrites existing mappings. - void map(size_t index, size_t prog) { + void map(const size_t index, const size_t prog) { + if (index >= indexToProgram_.size()) { + indexToProgram_.resize(index + 1); + } + if (prog >= programToIndex_.size()) { + programToIndex_.resize(prog + 1); + } indexToProgram_[index] = prog; programToIndex_[prog] = index; } /// Swap two program indices. - void swap(size_t prog0, size_t prog1) { + void swap(const size_t prog0, const size_t prog1) { const auto i0 = lookupIndex(prog0); const auto i1 = lookupIndex(prog1); std::swap(programToIndex_[prog0], programToIndex_[prog1]); @@ -146,9 +155,9 @@ struct MappingPass : impl::MappingPassBase { private: /// Maps the i-th wire index to a program index. - DenseMap indexToProgram_; + SmallVector indexToProgram_; /// Maps a program index to the i-th wire index. - DenseMap programToIndex_; + SmallVector programToIndex_; }; /// Statistics collected while routing. @@ -237,17 +246,88 @@ struct MappingPass : impl::MappingPassBase { } }; + /// Describes the graph F of arXiv:1602.05150v3. + struct FGraph { + explicit FGraph(std::shared_ptr device) + : f_(device->qubits()), device_(std::move(device)) {}; + + /// Build F-graph: Add edges to F for each edge in the coupling graph. + /// Note that this assumes that the coupling graph is directed, but + /// symmetric (essentially: undirected). + void construct(const Layout& from, const Layout& to) { + for (const auto u : device_->qubits()) { + for (const auto v : device_->neighboursOf(u)) { + if (shouldAddEdge(u, v, from, to)) { + f_.addEdge(u, v); + } + } + } + } + + /// 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. + [[nodiscard]] std::optional> + findHappySWAPChain() const { + const auto optCycle = f_.findCycle(); + if (!optCycle) { + return std::nullopt; + } + const auto& cycle = *optCycle; + + SmallVector swaps; + for (size_t i = cycle.size() - 1; i > 0; --i) { + swaps.emplace_back(cycle[i], cycle[i - 1]); + } + return swaps; + } + + /// Find an unhappy SWAP. That is, find an edge (u, v), where exchanging u + /// and v, reduces u's distance to its target location (by one) and + /// increases v's distance from 0 (already at the correct location) to one. + [[nodiscard]] std::optional findUnhappySWAP() const { + for (const auto u : f_.getNodes()) { + for (const auto v : f_.getNeighbours(u)) { + if (f_.getDegree(v) == 0) { + return {{u, v}}; + } + } + } + + return std::nullopt; + } + + /// Reset the F graph for rebuilding. + void reset() { f_.clearEdges(); } + + private: + /// Return true, if moving the program qubit on hardware qubit u to hardware + /// qubit v brings it closer to its destination hardware qubit. + [[nodiscard]] bool shouldAddEdge(const size_t u, const size_t v, + const Layout& from, + const Layout& to) const { + const auto dest = to.getHardwareIndex(from.getProgramIndex(u)); + return device_->distanceBetween(v, dest) < + device_->distanceBetween(u, dest); + } + + Graph f_; + std::shared_ptr device_; + }; + public: /// Construct default mapping pass. MappingPass() = default; /// Construct default mapping pass with options. - explicit MappingPass(MappingPassOptions options) : MappingPassBase(options) {} + explicit MappingPass(const MappingPassOptions& options) + : MappingPassBase(options) {} /// Construct mapping from coupling set. - explicit MappingPass( - const llvm::DenseSet>& couplingSet, - MappingPassOptions options) + explicit MappingPass(const DenseSet>& couplingSet, + const MappingPassOptions& options) : MappingPassBase(options), device(std::make_shared(couplingSet)) {} @@ -315,7 +395,7 @@ struct MappingPass : impl::MappingPassBase { numSwaps += stats.nswaps; // Fix SSA Dominance issues. - llvm::for_each(body.getBlocks(), [](Block& b) { sortTopologically(&b); }); + for_each(body.getBlocks(), [](Block& b) { sortTopologically(&b); }); } private: @@ -327,19 +407,35 @@ struct MappingPass : impl::MappingPassBase { OpBuilder::InsertionGuard guard(rewriter); rewriter.setInsertionPoint(forOp); - const auto naddons = addons.size(); const auto res = forOp.replaceWithAdditionalIterOperands(rewriter, addons, true); assert(succeeded(res)); - auto newForOp = cast(*res); - for (const auto [oldUse, newResult] : - llvm::zip_equal(addons, newForOp.getResults().take_back(naddons))) { - rewriter.replaceAllUsesExcept(oldUse, newResult, newForOp); + + for (const auto [before, after] : llvm::zip_equal( + addons, newForOp.getResults().take_back(addons.size()))) { + rewriter.replaceAllUsesExcept(before, after, newForOp); } return newForOp; } + /// Extend the qubit arguments of an `IfOp` by adding a given range of + /// additional SSA values. Replaces the existing operation and returns the + /// newly created one. + static IfOp extend(IfOp ifOp, ValueRange addons, IRRewriter& rewriter) { + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(ifOp); + + auto newIfOp = ifOp.replaceWithAdditionalQubits(rewriter, addons); + + for (const auto [before, after] : llvm::zip_equal( + addons, newIfOp->getResults().take_back(addons.size()))) { + rewriter.replaceAllUsesExcept(before, after, newIfOp); + } + + return newIfOp; + } + /// Return the wires of a dynamic computation. /// The mapping pass currently assumes that /// - there are no `qco.alloc` operation @@ -418,8 +514,7 @@ struct MappingPass : impl::MappingPassBase { Wires wires; WireInfos infos; - for (auto alloc : - llvm::make_early_inc_range(body.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. @@ -472,38 +567,62 @@ struct MappingPass : impl::MappingPassBase { stack.emplace_back(body, DenseSet{}); while (!stack.empty()) { - auto [region, qubits] = stack.pop_back_val(); - - for (Operation& op : llvm::make_early_inc_range(region.getOps())) { + for (auto [region, qubits] = stack.pop_back_val(); + Operation& op : make_early_inc_range(region.getOps())) { TypeSwitch(&op) - .Case([&](StaticOp op) { qubits.insert(op.getQubit()); }) - .Case([&](UnitaryOpInterface& op) { - for (const auto [pred, succ] : - llvm::zip_equal(op.getInputQubits(), op.getOutputQubits())) { + .Case( + [&](StaticOp staticOp) { qubits.insert(staticOp.getQubit()); }) + .Case([&](UnitaryOpInterface& uOp) { + for (const auto [pred, succ] : llvm::zip_equal( + uOp.getInputQubits(), uOp.getOutputQubits())) { qubits.insert(succ); qubits.erase(pred); } }) - .Case([&](scf::ForOp loop) { + .Case([&](scf::ForOp forOp) { assert(qubits.size() == layout.nqubits()); - DenseSet addons(qubits); - llvm::for_each(loop.getInits(), [&](auto v) { addons.erase(v); }); - auto newLoop = extend(loop, to_vector(addons), rewriter); + llvm::for_each(forOp.getInits(), + [&](Value v) { qubits.erase(v); }); - for (OpOperand& operand : newLoop.getInitsMutable()) { - qubits.insert(newLoop.getTiedLoopResult(&operand)); - qubits.erase(operand.get()); + auto newForOp = extend(forOp, to_vector(qubits), rewriter); + for (const auto [init, result] : llvm::zip_equal( + newForOp.getInits(), *newForOp.getLoopResults())) { + qubits.insert(result); + qubits.erase(init); } stack.emplace_back( - newLoop.getRegion(), - DenseSet(newLoop.getRegionIterArgs().begin(), - newLoop.getRegionIterArgs().end())); + newForOp.getRegion(), + DenseSet(newForOp.getRegionIterArgs().begin(), + newForOp.getRegionIterArgs().end())); }) - .Case([&](auto op) { - qubits.insert(op.getQubitOut()); - qubits.erase(op.getQubitIn()); + .Case([&](IfOp ifOp) { + assert(qubits.size() == layout.nqubits()); + + llvm::for_each(ifOp.getQubits(), + [&](Value v) { qubits.erase(v); }); + + auto newIfOp = extend(ifOp, to_vector(qubits), rewriter); + + for (const auto [qubit, result] : + llvm::zip_equal(newIfOp.getQubits(), newIfOp.getResults())) { + qubits.insert(result); + qubits.erase(qubit); + } + + const auto thenArgs = newIfOp.getThenRegion().getArguments(); + const auto elseArgs = newIfOp.getElseRegion().getArguments(); + stack.emplace_back( + newIfOp.getThenRegion(), + DenseSet(thenArgs.begin(), thenArgs.end())); + stack.emplace_back( + newIfOp.getElseRegion(), + DenseSet(elseArgs.begin(), elseArgs.end())); + }) + .Case([&](auto resetOp) { + qubits.insert(resetOp.getQubitOut()); + qubits.erase(resetOp.getQubitIn()); }) .Case([&](auto) { llvm::reportFatalInternalError("unexpected dynamic qubit alloc"); @@ -580,7 +699,7 @@ struct MappingPass : impl::MappingPassBase { /// /// Returns `failure`, if the A* search fails. FailureOr> search(const Window& window, - const Layout& layout) { + const Layout& layout) const { constexpr size_t cap = 25'000'000UL; const size_t b = device->maxDegree() * ((device->nqubits() + 1) / 2); @@ -601,7 +720,7 @@ struct MappingPass : impl::MappingPassBase { frontier.emplace(root); DenseMap, size_t> bestDepth; - DenseSet expansionSet; + SmallVector expansionSet; size_t i = 0; while (!frontier.empty() && i < budget) { @@ -616,8 +735,8 @@ struct MappingPass : impl::MappingPassBase { const auto [it, inserted] = bestDepth.try_emplace( curr->layout.getProgramToHardware(), curr->depth); if (!inserted) { - const auto otherDepth = it->getSecond(); - if (curr->depth >= otherDepth) { + if (const auto otherDepth = it->getSecond(); + curr->depth >= otherDepth) { ++i; continue; } @@ -631,7 +750,7 @@ struct MappingPass : impl::MappingPassBase { if (curr->isGoal(window.front(), *device)) { SmallVector seq(curr->depth); size_t j = seq.size() - 1; - for (Node* n = curr; n->parent != nullptr; n = n->parent) { + for (const Node* n = curr; n->parent != nullptr; n = n->parent) { seq[j] = n->swap; --j; } @@ -640,18 +759,18 @@ struct MappingPass : impl::MappingPassBase { } // Given a layout, create child-nodes for each possible SWAP - // between two neighbouring hardware qubits. + // between two neighboring hardware qubits. expansionSet.clear(); - const auto& [q0, q1] = window.front(); - for (const auto prog : {q0, q1}) { + for (const auto& [q0, q1] = window.front(); const auto prog : {q0, q1}) { for (const auto hw0 = curr->layout.getHardwareIndex(prog); const auto hw1 : device->neighboursOf(hw0)) { // Ensure consistent hashing/comparison. const IndexPairType swap = std::minmax(hw0, hw1); - if (!expansionSet.insert(swap).second) { + if (is_contained(expansionSet, swap)) { continue; } + expansionSet.push_back(swap); frontier.emplace(std::construct_at(arena.Allocate(), curr, swap, window, *device, params)); @@ -664,80 +783,92 @@ struct MappingPass : impl::MappingPassBase { return failure(); } - /// Return the sequence of SWAPs to move from one layout to another. + /// Return the SWAP sequence 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) { - + [[nodiscard]] SmallVector restore(const Layout& from, + const Layout& to) const { Layout curr(from); - Graph f(device->qubits()); + FGraph f(device); SmallVector swaps; - const auto shouldAddEdge = [&](size_t u, size_t v) { - const auto prog = curr.getProgramIndex(u); - const auto hwGoal = to.getHardwareIndex(prog); - return device->distanceBetween(v, hwGoal) < - device->distanceBetween(u, hwGoal); - }; - while (true) { + f.reset(); + f.construct(curr, to); - // Build F-graph: Add edges to F for each edge in the coupling graph. - // Note that this assumes that the coupling graph is directed, but - // symmetric (essentially: undirected). - - f.clearEdges(); - for (const auto u : device->qubits()) { - for (const auto v : device->neighboursOf(u)) { - if (shouldAddEdge(u, v)) { - f.addEdge(u, v); - } + if (const auto happy = f.findHappySWAPChain()) { + for (const auto& swap : *happy) { + swaps.emplace_back(swap); + curr.swap(swap.first, swap.second); } + continue; } - // 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. + // If there are no happy or unhappy swaps anymore, + // the final placement of every token is reached. - if (const auto cycle = f.findCycle(); cycle) { - for (size_t i = cycle->size() - 1; i > 0; --i) { - curr.swap((*cycle)[i], (*cycle)[i - 1]); - swaps.emplace_back((*cycle)[i], (*cycle)[i - 1]); - } - continue; + const auto unhappy = f.findUnhappySWAP(); + if (!unhappy) { + break; } - // Otherwise, search for an unhappy SWAP. That is, search for an edge (u, - // v), where exchanging u and v, reduces u's distance to its target - // location (by one) and increases v's distance from 0 (already at the - // correct location) to one. - - bool found{false}; - for (const auto u : f.getNodes()) { - for (const auto v : f.getNeighbours(u)) { - if (f.getDegree(v) == 0) { - curr.swap(u, v); - swaps.emplace_back(u, v); - found = true; - break; - } - } + swaps.emplace_back(*unhappy); + curr.swap(unhappy->first, unhappy->second); + } + + return swaps; + } + + /// Return a pair of SWAP sequences to transform two layouts into each other. + /// Inspired by the 4-Approximation algorithm described in arXiv:1602.05150v3, + /// with the key difference that the goal permutation is not static. + [[nodiscard]] std::pair, + SmallVector> + converge(const Layout& lhs, const Layout& rhs) const { + std::array layouts{Layout(lhs), Layout(rhs)}; + std::array graphs{FGraph(device), FGraph(device)}; + std::array, 2> swaps{}; - if (found) { + std::mt19937 gen(seed); + std::uniform_int_distribution coin(0, 1); + + while (true) { + size_t i = 0; + for (; i < 2; ++i) { + FGraph& f = graphs[i]; + + f.reset(); + f.construct(layouts[i], layouts[(i + 1) % 2]); + + if (const auto happy = f.findHappySWAPChain()) { + for (const auto& swap : *happy) { + swaps[i].emplace_back(swap); + layouts[i].swap(swap.first, swap.second); + } break; } } - // If there are no happy or unhappy swaps anymore, - // the final placement of every token is reached. + // If we exit early from the loop, we've found a happy SWAP chain. + if (i != 2) { + continue; + } + + // Otherwise, we randomly apply an unhappy SWAP to one of the layouts. + // If there is no happy or unhappy swaps anymore, the final placement of + // every token is reached. + + i = coin(gen); - if (!found) { + const auto unhappy = graphs[i].findUnhappySWAP(); + if (!unhappy) { break; } + + swaps[i].emplace_back(*unhappy); + layouts[i].swap(unhappy->first, unhappy->second); } - return swaps; + return {std::move(swaps[0]), std::move(swaps[1])}; } /// Skip to the end of the two-qubit block for both wire iterators, where @@ -748,9 +879,9 @@ struct MappingPass : impl::MappingPassBase { // 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. + // Otherwise, stop. - std::array block{it0, it1}; + std::array block{it0, it1}; while (true) { for (auto& it : block) { while (Traits::isActive(it)) { @@ -871,111 +1002,94 @@ struct MappingPass : impl::MappingPassBase { /// 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; + RecursiveRoutingStack advance(Wires& wires, const WireInfos& infos, + const Layout& layout) { + RecursiveRoutingStack stack; // Advance wires past all executable gates and push operations with // nested regions and the respective wire indices of their inputs onto the // result 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]); - const auto [hw0, hw1] = layout.getHardwareIndices(prog0, prog1); - if (device->areAdjacent(hw0, hw1)) { - 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 (const auto [hw0, hw1] = + layout.getHardwareIndices(prog0, prog1); + device->areAdjacent(hw0, hw1)) { + released.emplace_back(op); + } + }) + .template Case( + [&](auto 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; } - /// 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. + /// Processes the recursive stack item by routing the nested operation and + /// inserting epilogue SWAPs. template requires(Mode != RoutingMode::Hot || Direction == WireDirection::Forward) - LogicalResult route(RoutingBundle& bundle, Statistics& stats, - IRRewriter* rewriter = nullptr) { - using Traits = WireTraversalTraits; - - auto& [wires, infos, layout] = bundle; - - while (true) { - - while (true) { - const auto stack = advance(wires, infos, layout); - - if (stack.empty()) { - break; - } - - // Continue with processing the nested regions recursively. - - for (const auto& [op, indices] : stack) { - assert(isa(op)); - auto forOp = cast(op); - - RoutingBundle child{.layout = layout}; - - // 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. + LogicalResult dispatch(const RecursiveRoutingStackItem& item, + RoutingBundle& parent, Statistics& stats, + IRRewriter* rewriter = nullptr) { + const auto& [op, indices] = item; + return TypeSwitch(op) + .template Case([&](scf::ForOp forOp) { + RoutingBundle child{.layout = parent.layout}; + + // 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 prog = parent.infos.lookupProgram(i); + const auto res = cast(parent.wires[i].qubit()); const auto arg = forOp.getTiedLoopRegionIterArg(res); const auto index = child.wires.size(); 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); } + child.infos.map(index, prog); } - const auto res = route(child, stats, rewriter); - if (failed(res)) { + if (failed(route(child, stats, rewriter))) { return failure(); } - const auto swaps = restore(child.layout, layout); + const auto swaps = restore(child.layout, parent.layout); 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). + // 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(child.wires, [](auto& it) { std::advance(it, -2); }); + for_each(child.wires, [](auto& it) { std::advance(it, -2); }); } insertSWAPs(swaps, child, stats, rewriter); @@ -987,9 +1101,117 @@ struct MappingPass : impl::MappingPassBase { // 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()); + for_each(indices, [&](size_t i) { + std::advance(parent.wires[i], + WireTraversalTraits::stride()); }); + + return success(); + }) + .template Case([&](IfOp ifOp) { + std::array children{RoutingBundle{.layout = parent.layout}, + RoutingBundle{.layout = parent.layout}}; + + // Map parent (results) to child values (qubits). 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 = parent.infos.lookupProgram(i); + const auto res = cast(parent.wires[i].qubit()); + const auto index = children[0].wires.size(); + + OpOperand* qubit = ifOp.getTiedQubit(res); + const std::array args{ifOp.getTiedThenBlockArgument(qubit), + ifOp.getTiedElseBlockArgument(qubit)}; + + if constexpr (Direction == WireDirection::Forward) { + for (size_t j = 0; j < children.size(); ++j) { + children[j].wires.emplace_back(args[j]); + children[j].infos.map(index, prog); + } + } else { + const std::array yields{ + ifOp.getTiedThenYieldedValue(args[0])->get(), + ifOp.getTiedElseYieldedValue(args[1])->get()}; + for (size_t j = 0; j < children.size(); ++j) { + children[j].wires.emplace_back(yields[j]); + children[j].infos.map(index, prog); + } + } + } + + for (auto& child : children) { + if (failed(route(child, stats, rewriter))) { + return failure(); + } + } + + if constexpr (Mode == RoutingMode::Hot) { + + // After routing the branch 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). + + for_each(children[0].wires, [](auto& it) { std::advance(it, -2); }); + for_each(children[1].wires, [](auto& it) { std::advance(it, -2); }); + } + + const auto [fst, snd] = + converge(children[0].layout, children[1].layout); + + insertSWAPs(fst, children[0], stats, rewriter); + insertSWAPs(snd, children[1], stats, rewriter); + + if constexpr (Mode == RoutingMode::Hot) { + + // The IfOp implements the SingleBlockImplicitTerminator trait. + assert(ifOp.getThenRegion().hasOneBlock()); + assert(ifOp.getElseRegion().hasOneBlock()); + + sortTopologically(&ifOp.getThenRegion().getBlocks().front()); + sortTopologically(&ifOp.getElseRegion().getBlocks().front()); + } + + // Finally, move past the operation with nested regions by + // incrementing the respective global wires. + + for_each(indices, [&](size_t i) { + std::advance(parent.wires[i], + WireTraversalTraits::stride()); + }); + + return success(); + }) + .Default([](Operation*) { return failure(); }); + } + + /// 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) { + auto& [wires, infos, layout] = bundle; + + while (true) { + + while (true) { + const auto stack = advance(wires, infos, layout); + if (stack.empty()) { + break; + } + for (const auto& item : stack) { + if (dispatch(item, bundle, stats, rewriter) + .failed()) { + return failure(); + } } } @@ -1028,7 +1250,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::advance(it, 1); }); + for_each(wires, [](auto& it) { std::advance(it, 1); }); } } @@ -1041,7 +1263,7 @@ struct MappingPass : impl::MappingPassBase { } // namespace std::unique_ptr -createMappingPass(const llvm::DenseSet>& couplingSet, +createMappingPass(const DenseSet>& couplingSet, MappingPassOptions options) { // Verify the assumption that the coupling set is symmetric: @@ -1050,7 +1272,6 @@ createMappingPass(const llvm::DenseSet>& couplingSet, for (const auto& [u, v] : couplingSet) { if (u == v) { llvm::reportFatalUsageError("Found an invalid (u, u) edge."); - return nullptr; } if (!couplingSet.contains({v, u})) { diff --git a/mlir/lib/Dialect/QCO/Utils/Graph.cpp b/mlir/lib/Dialect/QCO/Utils/Graph.cpp index 0f48b53281..5e3de25add 100644 --- a/mlir/lib/Dialect/QCO/Utils/Graph.cpp +++ b/mlir/lib/Dialect/QCO/Utils/Graph.cpp @@ -30,6 +30,7 @@ void Graph::addEdge(size_t u, size_t v) { } ArrayRef Graph::getNeighbours(size_t id) const { return adj_.at(id); } + SmallVector Graph::getNodes() const { return to_vector(adj_.keys()); } size_t Graph::getMaxDegree() const { diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 98ea6538d9..7038ef1b0c 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -32,6 +32,7 @@ #include #include +#include #include #include #include @@ -86,35 +87,78 @@ isExecutable(Region& body, DenseMap& m, m.try_emplace(succ, hw); } }) - .Case([&](scf::ForOp op) { + .Case([&](scf::ForOp forOp) { DenseMap loopM; for (const auto [init, arg] : - llvm::zip_equal(op.getInits(), op.getRegionIterArgs())) { + llvm::zip_equal(forOp.getInits(), forOp.getRegionIterArgs())) { const auto hw = m.at(init); loopM.try_emplace(arg, hw); } - for (OpOperand& operand : op.getInitsMutable()) { + for (OpOperand& operand : forOp.getInitsMutable()) { const auto pred = operand.get(); - const auto succ = op.getTiedLoopResult(&operand); + const auto succ = forOp.getTiedLoopResult(&operand); const auto hw = m.at(pred); m.try_emplace(succ, hw); } - if (!isExecutable(op.getRegion(), loopM, couplingSet)) { + if (!isExecutable(forOp.getRegion(), loopM, couplingSet)) { executable = false; return; } - for (const auto& [arg, yielded] : - llvm::zip_equal(op.getRegionIterArgs(), op.getYieldedValues())) { + for (const auto& [arg, yielded] : llvm::zip_equal( + forOp.getRegionIterArgs(), forOp.getYieldedValues())) { if (loopM.at(arg) != loopM.at(yielded)) { - llvm::dbgs() << "for loop layout not restored!\n"; + llvm::dbgs() << "scf::forOp: layout not restored!\n"; executable = false; return; } } }) + .Case([&](qco::IfOp ifOp) { + std::array mappings{DenseMap{}, + DenseMap{}}; + + const std::array regions{&ifOp.getThenRegion(), + &ifOp.getElseRegion()}; + + for (size_t i = 0; i < 2; ++i) { + for (const auto [init, arg] : llvm::zip_equal( + ifOp.getQubits(), regions[i]->getArguments())) { + const auto hw = m.at(init); + mappings[i].try_emplace(arg, hw); + } + } + + for (OpOperand& operand : ifOp.getQubitsMutable()) { + const auto pred = operand.get(); + const auto succ = ifOp.getTiedResult(&operand); + const auto hw = m.at(pred); + m.try_emplace(succ, hw); + } + + std::array, 2> finalPermutation{}; + + for (size_t i = 0; i < 2; ++i) { + Region* body = regions[i]; + if (!isExecutable(*body, mappings[i], couplingSet)) { + executable = false; + return; + } + + Block& block = body->getBlocks().front(); + auto yield = cast(block.getTerminator()); + for (const auto v : yield.getTargets()) { + finalPermutation[i].emplace_back(mappings[i].at(v)); + } + } + + if (finalPermutation[0] != finalPermutation[1]) { + executable = false; + return; + } + }) .Case([&](auto op) { const auto pred = op.getQubitIn(); const auto succ = op.getQubitOut(); @@ -590,5 +634,63 @@ TEST_P(MappingPassTest, Sabre) { EXPECT_TRUE(isExecutable(entry, device.couplingSet)); } +TEST_P(MappingPassTest, RandomOrderGHZ) { + const auto& device = GetParam(); + + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + Value tensor = builder.qtensorAlloc(9); + SmallVector qubits(9); + SmallVector cregs(9); + + for (int64_t i = 0; i < 9; ++i) { + std::tie(tensor, qubits[i]) = builder.qtensorExtract(tensor, i); + } + + qubits[0] = builder.h(qubits[0]); + std::tie(qubits[0], cregs[0]) = builder.measure(qubits[0]); + + qubits = builder.qcoIf( + cregs[0], qubits, + [&](ValueRange args) { + SmallVector values(args); + values[0] = builder.h(values[0]); + for (size_t i = 1; i < 9; ++i) { + std::tie(values[0], values[i]) = builder.cx(values[0], values[i]); + } + return values; + }, + [&](ValueRange args) { + SmallVector values(args); + values[8] = builder.h(values[8]); + for (size_t i = 8; i > 0; --i) { + std::tie(values[8], values[i - 1]) = + builder.cx(values[8], values[i - 1]); + } + return values; + }); + + qubits = builder.barrier(qubits); + + for (int64_t i = 0; i < 9; ++i) { + std::tie(qubits[i], cregs[i]) = builder.measure(qubits[i]); + } + + for (int64_t i = 0; i < 9; ++i) { + tensor = builder.qtensorInsert(qubits[i], tensor, i); + } + + builder.qtensorDealloc(tensor); + + auto m = builder.finalize(); + auto res = + runPass(m.get(), device.couplingSet, MappingPassOptions{.ntrials = 1}); + auto entry = getEntryPoint(m.get()); + + ASSERT_TRUE(res.succeeded()); + EXPECT_TRUE(isExecutable(entry, device.couplingSet)); +} + INSTANTIATE_TEST_SUITE_P(NineQubitSquareGrid, MappingPassTest, testing::Values(getNineQubitSquareGrid()));