✨ Implement Trials For Mapping Pass

CHANGELOG.md

@@ -13,7 +13,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel
 
 - ✨ Add Sampler and Estimator Primitives to the QDMI-Qiskit Interface ([#1507]) ([**@marcelwa**])
 - ✨ Add conversions between Jeff and QCO ([#1479], [#1548]) ([**@denialhaag**])
-- ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547]) ([**@MatthiasReumann**])
+- ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568]) ([**@MatthiasReumann**])
 - ✨ 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], [#1513], [#1521], [#1548], [#1550], [#1554])
   ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**])
@@ -331,6 +331,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool
 
 <!-- PR links -->
 
+[#1568]: https://github.com/munich-quantum-toolkit/core/pull/1568
 [#1554]: https://github.com/munich-quantum-toolkit/core/pull/1554
 [#1550]: https://github.com/munich-quantum-toolkit/core/pull/1550
 [#1549]: https://github.com/munich-quantum-toolkit/core/pull/1549

mlir/include/mlir/Dialect/QCO/Transforms/Passes.td

@@ -49,6 +49,9 @@ def MappingPass : Pass<"place-and-route", "mlir::ModuleOp"> {
     each traversal, the pass routes the circuit and updates the dynamic-to-static mapping based on the routing decisions
     made during that traversal. By performing multiple traversals, the pass can iteratively refine the mapping and
     potentially find a more optimal solution. This is behavior is controlled by the `niterations` parameter.
+
+    The pass option `ntrials` determines how many random initial layouts the pass explores. Per default, the pass always
+    explores the identity layout. If compiled with multi-threading on, these trials will be executed in parallel.
   }];
   let options = [
     Option<"nlookahead", "nlookahead", "std::size_t", "1",
@@ -58,7 +61,11 @@ def MappingPass : Pass<"place-and-route", "mlir::ModuleOp"> {
     Option<"lambda", "lambda", "float", "0.5F",
            "The lambda factor in the cost function.">,
     Option<"niterations", "niterations", "std::size_t", "2",
-           "The number of forwards and backwards traversal to improve the initial layout.">
+           "The number of forwards and backwards traversal to improve the initial layout.">,
+    Option<"ntrials", "ntrials", "std::size_t", "4",
+           "The number of (possibly parallel) random trials of the forwards and backwards mechanism.">,
+    Option<"seed", "seed", "std::size_t", "42",
+           "A seed used for randomization.">
   ];
 }
 

mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp

@@ -27,6 +27,7 @@
 #include <mlir/IR/Location.h>
 #include <mlir/IR/Operation.h>
 #include <mlir/IR/PatternMatch.h>
+#include <mlir/IR/Threading.h>
 #include <mlir/IR/Value.h>
 #include <mlir/Support/LLVM.h>
 #include <mlir/Support/WalkResult.h>
@@ -37,7 +38,10 @@
 #include <cstdint>
 #include <functional>
 #include <iterator>
+#include <numeric>
+#include <optional>
 #include <queue>
+#include <random>
 #include <string>
 #include <string_view>
 #include <tuple>
@@ -88,6 +92,25 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
       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 std::size_t nqubits, const std::size_t seed) {
+      SmallVector<IndexType> 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.
@@ -300,6 +323,9 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
   using MappingPassBase::MappingPassBase;
 
   void runOnOperation() override {
+    std::mt19937_64 rng{this->seed};
+    IRRewriter rewriter(&getContext());
+
     Parameters params(this->alpha, this->lambda, this->nlookahead);
     // TODO: Hardcoded architecture.
     Architecture arch("RigettiNovera", 9,
@@ -308,7 +334,6 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
                        {3, 4}, {4, 3}, {4, 7}, {7, 4}, {4, 5}, {5, 4},
                        {5, 8}, {8, 5}, {6, 7}, {7, 6}, {7, 8}, {8, 7}});
 
-    IRRewriter rewriter(&getContext());
     for (auto func : getOperation().getOps<func::FuncOp>()) {
       const auto dyn = collectDynamicQubits(func.getFunctionBody());
       if (dyn.size() > arch.nqubits()) {
@@ -320,34 +345,114 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
 
       const auto [ltr, rtl] = computeBidirectionalLayers(dyn);
 
-      // Use the SABRE Approach to improve the initial layout choice (here:
-      // identity): Traverse the layers from left-to-right-to-left and
-      // cold-route along the way. Repeat this procedure "niterations" times.
+      // Create trials. Currently this includes the identity layout and
+      // `ntrials` many random layouts.
 
-      Layout layout = Layout::identity(arch.nqubits());
-      for (std::size_t r = 0; r < this->niterations; ++r) {
-        if (failed(routeCold(ltr, layout, arch, params))) {
-          signalPassFailure();
-          return;
-        }
-        if (failed(routeCold(rtl, layout, arch, params))) {
-          signalPassFailure();
-          return;
-        }
+      SmallVector<Layout> trials;
+      trials.reserve(1 + this->ntrials);
+      trials.emplace_back(Layout::identity(arch.nqubits()));
+      for (std::size_t i = 0; i < this->ntrials; ++i) {
+        trials.emplace_back(Layout::random(arch.nqubits(), rng()));
+      }
+
+      // Execute each of the trials (possibly in parallel). Collect the results
+      // and find the one with the fewest SWAPs.
+
+      SmallVector<std::optional<TrialResult>> results(trials.size());
+      parallelForEach(
+          &getContext(), enumerate(trials), [&, this](auto indexedTrial) {
+            auto [idx, layout] = indexedTrial;
+            auto res = runMappingTrial(ltr, rtl, arch, params, layout);
+            if (succeeded(res)) {
+              results[idx] = std::move(*res);
+            }
+          });
+
+      TrialResult* best = findBestTrial(results);
+      if (best == nullptr) {
+        signalPassFailure();
+        return;
       }
 
       // Once the initial layout is found, replace the dynamic with static
       // qubits ("placement") and hot-route the circuit layer-by-layer.
 
-      const auto stat = place(dyn, layout, func.getFunctionBody(), rewriter);
-      if (failed(routeHot(ltr, layout, stat, arch, params, rewriter))) {
+      const auto stat =
+          place(dyn, best->layout, func.getFunctionBody(), rewriter);
+      if (failed(commitTrial(*best, stat, rewriter))) {
         signalPassFailure();
         return;
       };
     }
   }
 
 private:
+  struct [[nodiscard]] TrialResult {
+    explicit TrialResult(Layout layout) : layout(std::move(layout)) {}
+
+    /// @brief The computed initial layout.
+    Layout layout;
+    /// @brief A vector of SWAPs for each layer.
+    SmallVector<SmallVector<IndexGate>> swaps;
+    /// @brief The number of inserted SWAPs.
+    std::size_t nswaps{};
+  };
+
+  /**
+   * @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 TrialResult*
+  findBestTrial(MutableArrayRef<std::optional<TrialResult>> results) {
+    TrialResult* best = nullptr;
+    for (auto& opt : results) {
+      if (opt.has_value()) {
+        if (best == nullptr || best->nswaps > opt->nswaps) {
+          best = &opt.value();
+        }
+      }
+    }
+    return best;
+  }
+
+  /**
+   * @brief Run a mapping trial.
+   * @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 the trial result or failure() on failure.
+   */
+  FailureOr<TrialResult> runMappingTrial(ArrayRef<Layer> ltr,
+                                         ArrayRef<Layer> rtl,
+                                         const Architecture& arch,
+                                         const Parameters& params,
+                                         Layout& layout) {
+    // Perform forwards and backwards traversals.
+    for (std::size_t i = 0; i < this->niterations; ++i) {
+      if (failed(route(ltr, arch, params, layout, [](const auto&) {}))) {
+        return failure();
+      }
+      if (failed(route(rtl, arch, params, layout, [](const auto&) {}))) {
+        return failure();
+      }
+    }
+
+    TrialResult result(layout); // Copies the final initial layout.
+
+    // Helper function that adds the SWAPs to the trial result.
+    const auto collectSwaps = [&](ArrayRef<IndexGate> swaps) {
+      result.nswaps += swaps.size();
+      result.swaps.emplace_back(swaps);
+    };
+
+    // Perform final left-to-right traversal whilst collecting SWAPs.
+    if (failed(route(ltr, arch, params, layout, collectSwaps))) {
+      return failure();
+    }
+
+    return result;
+  }
+
   /**
    * @brief Collect dynamic qubits contained in the given function body.
    * @returns a vector of SSA values produced by qco.alloc operations.
@@ -406,12 +511,14 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
   /**
    * @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 parameter @p maxIterations determines how many nodes will
+   * be explored until the current search is considered a failure.
    * @returns a vector of hardware-index pairs (each denoting a SWAP) or
    * failure() if A* fails.
    */
   [[nodiscard]] static FailureOr<SmallVector<IndexGate>>
   search(ArrayRef<Layer> layers, const Layout& layout, const Architecture& arch,
-         const Parameters& params) {
+         const Parameters& params, const std::size_t maxIterations = 50'000) {
 
     Node root(layout);
     if (root.isGoal(layers.front(), arch)) {
@@ -422,7 +529,8 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
     frontier.emplace(root);
     DenseSet<IndexGate> expansionSet;
 
-    while (!frontier.empty()) {
+    std::size_t i = 0;
+    while (!frontier.empty() && i < maxIterations) {
       Node curr = frontier.top();
       frontier.pop();
 
@@ -452,6 +560,8 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
           }
         }
       }
+
+      ++i;
     }
 
     return failure();
@@ -553,11 +663,16 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
    * @details Iterates over a sliding window of layers and uses the A* search
    * engine to find a sequence of SWAPs that makes that layer executable.
    * Instead of inserting these SWAPs into the IR, this function only updates
-   * the layout.
+   * (and hence modifies) the layout. The function calls the callback @p onSwaps
+   * for each layer with the found sequence of SWAPs.
    * @returns failure() if A* search isn't able to find a solution.
    */
-  LogicalResult routeCold(ArrayRef<Layer> layers, Layout& layout,
-                          const Architecture& arch, const Parameters& params) {
+  template <typename OnSwaps>
+  LogicalResult route(ArrayRef<Layer> layers, const Architecture& arch,
+                      const Parameters& params, Layout& layout,
+                      OnSwaps&& onSwaps) {
+    auto&& callback = std::forward<OnSwaps>(onSwaps);
+
     for (std::size_t i = 0; i < layers.size(); ++i) {
       const std::size_t len = std::min(1 + nlookahead, layers.size() - i);
       const auto window = layers.slice(i, len);
@@ -569,6 +684,8 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
       for (const auto& [hw0, hw1] : *swaps) {
         layout.swap(hw0, hw1);
       }
+
+      std::invoke(callback, *swaps);
     }
 
     return success();
@@ -581,9 +698,9 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
    * This function inserts SWAP ops.
    * @returns failure() if A* search isn't able to find a solution.
    */
-  LogicalResult routeHot(ArrayRef<Layer> ltr, Layout& layout,
-                         ArrayRef<QubitValue> statics, const Architecture& arch,
-                         const Parameters& params, IRRewriter& rewriter) {
+  static LogicalResult commitTrial(const TrialResult& result,
+                                   ArrayRef<QubitValue> statics,
+                                   IRRewriter& rewriter) {
     // Helper function that advances the iterator to the input qubit (the
     // operation producing it) of a deallocation or two-qubit op.
     const auto advFront = [](WireIterator& it) {
@@ -602,21 +719,15 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
       }
     };
 
+    DenseMap<Operation*, WireIterator*> seen;
     auto wires = toWires(statics);
-    for (const auto [i, layer] : enumerate(ltr)) {
-      // Advance all wires to the next front of one-qubit outputs (the SSA
-      // values).
+    for (const auto [i, swaps] : enumerate(result.swaps)) {
+      // Advance all wires to the next front of one-qubit outputs
+      // (the SSA values).
       for_each(wires, advFront);
 
-      // Collect window and use A* to find and insert a sequence of swaps.
-      const auto len = std::min(1 + this->nlookahead, ltr.size() - i);
-      const auto window = ltr.slice(i, len);
-      const auto swaps = search(window, layout, arch, params);
-      if (failed(swaps)) {
-        return failure();
-      }
-
-      for (const auto& [hw0, hw1] : *swaps) {
+      // Apply the sequence of SWAPs and rewire the qubit SSA values.
+      for (const auto& [hw0, hw1] : swaps) {
         Operation* op0 = wires[hw0].operation();
         Operation* op1 = wires[hw1].operation();
         const auto in0 = wires[hw0].qubit();
@@ -641,15 +752,23 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
         // Jump over the SWAPOp.
         std::ranges::advance(wires[hw0], 1);
         std::ranges::advance(wires[hw1], 1);
-
-        layout.swap(hw0, hw1);
       }
 
-      // Jump over two-qubit gates contained in the layer.
-      for (const auto& [prog0, prog1] : layer) {
-        std::ranges::advance(wires[layout.getHardwareIndex(prog0)], 1);
-        std::ranges::advance(wires[layout.getHardwareIndex(prog1)], 1);
+      // Jump over "ready" two-qubit gates.
+      for (auto& it : wires) {
+        auto op = dyn_cast<UnitaryOpInterface>(std::next(it).operation());
+        if (op && op.getNumQubits() > 1) {
+          if (seen.contains(op)) {
+            std::ranges::advance(it, 1);
+            std::ranges::advance(*seen[op], 1);
+            continue;
+          }
+
+          seen.try_emplace(op, &it);
+        }
       }
+
+      seen.clear(); // Prepare for next iteration.
     }
 
     return success();

mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp

@@ -104,8 +104,12 @@ class MappingPassTest : public testing::Test,
   static void runHeuristicMapping(OwningOpRef<ModuleOp>& moduleOp) {
     PassManager pm(moduleOp->getContext());
     pm.addPass(createQCToQCO());
-    pm.addPass(qco::createMappingPass(qco::MappingPassOptions{
-        .nlookahead = 5, .alpha = 1, .lambda = 0.85, .niterations = 2}));
+    pm.addPass(qco::createMappingPass(qco::MappingPassOptions{.nlookahead = 5,
+                                                              .alpha = 1,
+                                                              .lambda = 0.85,
+                                                              .niterations = 2,
+                                                              .ntrials = 8,
+                                                              .seed = 1337}));
     pm.addPass(createQCOToQC());
     auto res = pm.run(*moduleOp);
     ASSERT_TRUE(succeeded(res));