✨Arena-Based A* Search For Mapping Pass

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

@@ -18,7 +18,6 @@
 #include <llvm/ADT/STLExtras.h>
 #include <llvm/ADT/SmallVector.h>
 #include <llvm/ADT/TypeSwitch.h>
-#include <llvm/Support/Debug.h>
 #include <llvm/Support/ErrorHandling.h>
 #include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/IR/Block.h>
@@ -70,6 +69,8 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
    */
   enum class Direction : std::uint8_t { Forward, Backward };
 
+  struct LayoutInfo;
+
   /**
    * @brief A qubit layout that maps program and hardware indices without
    * storing Values. Used for efficient memory usage when Value tracking isn't
@@ -191,18 +192,6 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
       return programToHardware_.size();
     }
 
-    void dump() {
-      llvm::dbgs() << "prog= ";
-      for (std::size_t i = 0; i < nqubits(); ++i) {
-        llvm::dbgs() << i << " ";
-      }
-      llvm::dbgs() << "\nhw=   ";
-      for (std::size_t i = 0; i < nqubits(); ++i) {
-        llvm::dbgs() << programToHardware_[i] << ' ';
-      }
-      llvm::dbgs() << '\n';
-    }
-
   protected:
     /**
      * @brief Maps a program qubit index to its hardware index.
@@ -215,10 +204,43 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
     SmallVector<IndexType> hardwareToProgram_;
 
   private:
+    friend struct MappingPass::LayoutInfo;
+
+    Layout() = default;
     explicit Layout(const std::size_t nqubits)
         : programToHardware_(nqubits), hardwareToProgram_(nqubits) {}
   };
 
+  /**
+   * @brief Required to use Layout as a key for LLVM maps and sets.
+   */
+  class LayoutInfo {
+    using Info = DenseMapInfo<SmallVector<IndexType>>;
+
+  public:
+    static Layout getEmptyKey() {
+      Layout l;
+      l.programToHardware_ = Info::getEmptyKey();
+      l.hardwareToProgram_ = Info::getEmptyKey();
+      return l;
+    }
+
+    static Layout getTombstoneKey() {
+      Layout l;
+      l.programToHardware_ = Info::getTombstoneKey();
+      l.hardwareToProgram_ = Info::getTombstoneKey();
+      return l;
+    }
+
+    static unsigned getHashValue(const Layout& l) {
+      return Info::getHashValue(l.programToHardware_);
+    }
+
+    static bool isEqual(const Layout& a, const Layout& b) {
+      return Info::isEqual(a.programToHardware_, b.programToHardware_);
+    }
+  };
+
   /**
    * @brief Parameters influencing the behavior of the A* search algorithm.
    */
@@ -252,19 +274,23 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
 
     /**
      * @brief Construct a non-root node from its parent node. Apply the given
-     * swap to the layout of the parent node and evaluate the cost.
+     * swap to the layout of the parent node.
      */
-    Node(const Node& parent, IndexGate swap, ArrayRef<Layer> layers,
-         const Architecture& arch, const Parameters& params)
+    Node(const Node& parent, IndexGate swap)
         : sequence(parent.sequence), layout(parent.layout), f(0) {
       // Apply node-specific swap to given layout.
       layout.swap(swap.first, swap.second);
 
       // Add swap to sequence.
       sequence.emplace_back(swap);
+    }
 
-      // Evaluate cost function.
-      f = g(params.alpha) + h(layers, arch, params); // NOLINT
+    /**
+     * @brief Evaluate the cost function.
+     */
+    void evalCost(ArrayRef<Layer> layers, const Architecture& arch,
+                  const Parameters& params) {
+      f = g(params.alpha) + h(layers, arch, params);
     }
 
     /**
@@ -507,9 +533,9 @@ 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 iteration budget is then b^{3}, which corresponds to
-   * exhausting all paths of length up to b^{2} in a search tree with branching
-   * factor b. A hard cap prevents impractical runtimes on larger architectures.
+   * 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
@@ -534,6 +560,7 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
 
     MinQueue frontier{};
     frontier.emplace(root);
+    DenseSet<Layout, LayoutInfo> discovered{root.layout};
     DenseSet<IndexGate> expansionSet;
 
     std::size_t i = 0;
@@ -549,21 +576,29 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
       // two neighbouring hardware qubits.
 
       expansionSet.clear();
-      if (!curr.sequence.empty()) {
-        expansionSet.insert(curr.sequence.back());
-      }
-
       for (const IndexGate& gate : layers.front()) {
         for (const auto prog : {gate.first, gate.second}) {
           const auto hw0 = curr.layout.getHardwareIndex(prog);
           for (const auto hw1 : arch.neighboursOf(hw0)) {
-            /// Ensure consistent hashing/comparison.
+            // Ensure consistent hashing/comparison.
             const IndexGate swap = std::minmax(hw0, hw1);
             if (!expansionSet.insert(swap).second) {
               continue;
             }
 
-            frontier.emplace(curr, swap, layers, arch, params);
+            // Multiple sequences of SWAPs may lead to the same layout.
+            // The discovered set ensures that we do not visit the same layout
+            // multiple times. Only after ensuring that the layout hasn't been
+            // discovered so far, we evaluate the "expensive" cost function.
+            // TODO: In the future, should fidelities be ever considered, the
+            // sequence of SWAPs matters - so this will become more difficult.
+
+            Node child(curr, swap);
+            if (!discovered.insert(child.layout).second) {
+              continue;
+            }
+            child.evalCost(layers, arch, params);
+            frontier.emplace(std::move(child));
           }
         }
       }