🐛 Fix SSA Dominance Issues for Mapping Pass for Circuits With Barriers (#1581)

Author: MatthiasReumann

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], [#1565]) ([**@denialhaag**])
-- ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568]) ([**@MatthiasReumann**])
+- ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568], [#1581]) ([**@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], [#1542], [#1548], [#1550], [#1554], [#1570], [#1572], [#1573])
   ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**])
@@ -333,6 +333,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool
 
 <!-- PR links -->
 
+[#1581]: https://github.com/munich-quantum-toolkit/core/pull/1581
 [#1573]: https://github.com/munich-quantum-toolkit/core/pull/1573
 [#1572]: https://github.com/munich-quantum-toolkit/core/pull/1572
 [#1571]: https://github.com/munich-quantum-toolkit/core/pull/1571

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

@@ -321,6 +321,65 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
 
   using MinQueue = std::priority_queue<Node, std::vector<Node>, std::greater<>>;
 
+  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{};
+  };
+
+  struct SynchronizationMap {
+    /**
+     * @returns true if the operation is contained in the map.
+     */
+    bool contains(Operation* op) const { return onHold.contains(op); }
+
+    /**
+     * @brief Add op with respective iterator and ref count to the map.
+     */
+    void add(Operation* op, WireIterator* it, const std::size_t cnt) {
+      onHold.try_emplace(op, SmallVector<WireIterator*>{it});
+      // Decrease the cnt by one because the op was visited when adding.
+      refCount.try_emplace(op, cnt - 1);
+    }
+
+    /**
+     * @brief Decrement ref count of op and potentially release its iterators.
+     */
+    std::optional<ArrayRef<WireIterator*>> visit(Operation* op,
+                                                 WireIterator* it) {
+      assert(refCount.contains(op) && "expected sync map to contain op");
+
+      // Add iterator for later release.
+      onHold[op].push_back(it);
+
+      // Release iterators whenever the ref count reaches zero.
+      if (--refCount[op] == 0) {
+        return onHold[op];
+      }
+
+      return std::nullopt;
+    }
+
+    /**
+     * @brief Clear the contents of the map.
+     */
+    void clear() {
+      onHold.clear();
+      refCount.clear();
+    }
+
+  private:
+    /// @brief Maps operations to to-be-released iterators.
+    DenseMap<Operation*, SmallVector<WireIterator*, 2>> onHold;
+    /// @brief Maps operations to ref counts.
+    DenseMap<Operation*, std::size_t> refCount;
+  };
+
 public:
   using MappingPassBase::MappingPassBase;
 
@@ -382,17 +441,6 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
   }
 
 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.
@@ -615,28 +663,33 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
         while (shouldContinue(it)) {
           const auto res =
               TypeSwitch<Operation*, WalkResult>(it.operation())
-                  .Case<UnitaryOpInterface>([&](UnitaryOpInterface op) {
-                    assert(op.getNumQubits() > 0 && op.getNumQubits() <= 2);
+                  .Case<BarrierOp>([&](auto) {
+                    std::ranges::advance(it, step);
+                    return WalkResult::advance();
+                  })
+                  .template Case<UnitaryOpInterface>(
+                      [&](UnitaryOpInterface op) {
+                        assert(op.getNumQubits() > 0 && op.getNumQubits() <= 2);
 
-                    if (op.getNumQubits() == 1) {
-                      std::ranges::advance(it, step);
-                      return WalkResult::advance();
-                    }
+                        if (op.getNumQubits() == 1) {
+                          std::ranges::advance(it, step);
+                          return WalkResult::advance();
+                        }
 
-                    if (visited.contains(op)) {
-                      const auto otherIndex = visited[op];
-                      layer.insert(std::make_pair(index, otherIndex));
+                        if (visited.contains(op)) {
+                          const auto otherIndex = visited[op];
+                          layer.insert(std::make_pair(index, otherIndex));
 
-                      std::ranges::advance(wires[index], step);
-                      std::ranges::advance(wires[otherIndex], step);
+                          std::ranges::advance(wires[index], step);
+                          std::ranges::advance(wires[otherIndex], step);
 
-                      visited.erase(op);
-                    } else {
-                      visited.try_emplace(op, index);
-                    }
+                          visited.erase(op);
+                        } else {
+                          visited.try_emplace(op, index);
+                        }
 
-                    return WalkResult::interrupt();
-                  })
+                        return WalkResult::interrupt();
+                      })
                   .template Case<AllocOp, StaticOp, ResetOp, MeasureOp,
                                  DeallocOp>([&](auto) {
                     std::ranges::advance(it, step);
@@ -709,25 +762,30 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
     // 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) {
+      auto next = std::next(it);
       while (true) {
-        const auto next = std::next(it);
         if (isa<DeallocOp>(next.operation())) {
           break;
         }
 
+        if (isa<BarrierOp>(next.operation())) {
+          break;
+        }
+
         auto op = dyn_cast<UnitaryOpInterface>(next.operation());
         if (op && op.getNumQubits() > 1) {
           break;
         }
 
         std::ranges::advance(it, 1);
+        std::ranges::advance(next, 1);
       }
     };
 
     auto wires = toWires(place(dynQubits, result.layout, funcBody, rewriter));
 
-    DenseMap<Operation*, WireIterator*> seen;
-    for (const auto [i, swaps] : enumerate(result.swaps)) {
+    SynchronizationMap ready;
+    for (const auto& swaps : result.swaps) {
       // Advance all wires to the next front of one-qubit outputs
       // (the SSA values).
       for_each(wires, advFront);
@@ -760,21 +818,30 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
         std::ranges::advance(wires[hw1], 1);
       }
 
-      // Jump over "ready" two-qubit gates.
+      // Jump over "ready" 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;
-          }
+        if (!op) {
+          continue;
+        }
+
+        if (op.getNumQubits() < 2) {
+          continue;
+        }
 
-          seen.try_emplace(op, &it);
+        if (!ready.contains(op)) {
+          ready.add(op, &it, op.getNumQubits());
+          continue;
+        }
+
+        if (auto opt = ready.visit(op, &it)) {
+          for (WireIterator* wire : *opt) {
+            std::ranges::advance(*wire, 1);
+          }
         }
       }
 
-      seen.clear(); // Prepare for next iteration.
+      ready.clear(); // Prepare for next iteration.
     }
   }
 };

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

@@ -63,6 +63,9 @@ class MappingPassTest : public testing::Test,
 
     bool executable = true;
     std::ignore = moduleOp->walk([&](qc::UnitaryOpInterface op) {
+      if (isa<qc::BarrierOp>(op)) {
+        return WalkResult::advance();
+      }
       if (op.getNumQubits() > 1) {
         assert(op.getNumQubits() == 2 &&
                "Expected only 2-qubit gates after decomposition");
@@ -178,6 +181,14 @@ TEST_P(MappingPassTest, Sabre) {
 
   builder.cx(q3, q0);
 
+  builder.barrier({q0, q1, q2, q3, q4, q5});
+  builder.measure(q0);
+  builder.measure(q1);
+  builder.measure(q2);
+  builder.measure(q3);
+  builder.measure(q4);
+  builder.measure(q5);
+
   builder.dealloc(q0);
   builder.dealloc(q1);
   builder.dealloc(q2);