✨ Extend Layout API for partial injections

CHANGELOG.md

@@ -64,8 +64,8 @@ releases may include breaking changes.
   [**@simon1hofmann**])
 - ✨ Add a compiler-target-aware `place-and-route` pass ([#1537], [#1547],
   [#1568], [#1581], [#1583], [#1588], [#1600], [#1664], [#1709], [#1716],
-  [#1748], [#1805], [#1870], [#1904], [#1911], [#1951], [#1997], [#2016],
-  [#2060]) ([**@MatthiasReumann**], [**@burgholzer**])
+  [#1748], [#1805], [#1870], [#1904], [#1911], [#1951], [#1956], [#1997],
+  [#2016], [#2060]) ([**@MatthiasReumann**], [**@burgholzer**], [**@rturrado**])
 - ✨ Add modifier and global-phase normalization passes ([#1986], [#1995],
   [#2015]) ([**@burgholzer**], [**@denialhaag**])
 - ✨ Add single-qubit optimization passes for unitary fusion, Hadamard lifting,
@@ -969,6 +969,7 @@ for previous changelogs._
 [#1965]: https://github.com/munich-quantum-toolkit/core/pull/1965
 [#1961]: https://github.com/munich-quantum-toolkit/core/pull/1961
 [#1957]: https://github.com/munich-quantum-toolkit/core/pull/1957
+[#1956]: https://github.com/munich-quantum-toolkit/core/pull/1956
 [#1953]: https://github.com/munich-quantum-toolkit/core/pull/1953
 [#1952]: https://github.com/munich-quantum-toolkit/core/pull/1952
 [#1951]: https://github.com/munich-quantum-toolkit/core/pull/1951

mlir/include/mlir/Dialect/QCO/Utils/Layout.h

@@ -20,23 +20,36 @@
 
 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.
+/// A qubit layout that maps program qubit indices to hardware qubit indices
+/// without storing Values.
 ///
-/// Note that we use the terminology "hardware" and "program" qubits
-/// here, because "virtual" (opposed to physical) and "static" (opposed to
-/// dynamic) are C++ keywords.
+/// Program and hardware qubit indices form a dense range, respectively
+/// `[0, nProgramQubits)` and `[0, nHardwareQubits)`, with `nProgramQubits <=
+/// nHardwareQubits`, and every program qubit is mapped to a distinct hardware
+/// qubit. Unmapped hardware slots carry a sentinel value.
+
+/// 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);
+  /// Construct an empty layout.
+  Layout() = default;
 
-  /// Construct a layout from a program-to-hardware mapping,
+  /// Construct and return a random layout that maps every program qubit
+  /// index in `[0, nProgramQubits)` to a distinct hardware index drawn from
+  /// `[0, nHardwareQubits)`.
+  static Layout random(size_t nProgramQubits, size_t nHardwareQubits,
+                       size_t seed);
+
+  /// Construct a layout from a bijective program-to-hardware mapping,
   /// where mapping[prog] = hw.
+  /// Sets both `nProgramQubits` and `nHardwareQubits` to `mapping.size()`.
   static Layout fromMapping(ArrayRef<size_t> mapping);
 
-  /// Insert program:hardware index mapping.
+  /// Insert a program:hardware index mapping.
+  /// Requires `prog < nProgramQubits`, `hw < nHardwareQubits`, and that
+  /// neither `prog` nor `hw` has been mapped previously.
   void add(size_t prog, size_t hw);
 
   /// Lookup and return program index for a hardware index.
@@ -61,28 +74,38 @@ class Layout {
     return std::tuple{getProgramIndex(static_cast<size_t>(hws))...};
   }
 
+  /// Return true if `hw` currently has a program qubit assigned to it.
+  [[nodiscard]] bool hasProgramAt(size_t hw) const;
+
   /// Swap the mapping to program indices of two hardware indices.
+  /// Both sides must currently have a program qubit assigned.
   void swap(size_t hwA, size_t hwB);
 
-  /// Return the number of qubits managed by the layout.
-  [[nodiscard]] size_t nqubits() const;
+  /// Return the number of program qubits this layout was declared with.
+  [[nodiscard]] size_t nProgramQubits() const;
+
+  /// Return the number of hardware qubits this layout was declared with.
+  [[nodiscard]] size_t nHardwareQubits() const;
 
-  /// Return the program to hardware mapping.
+  /// Return a view of the program to hardware mapping of length
+  /// `nProgramQubits()`, where entry `prog` is the hardware index assigned to
+  /// program qubit `prog`. Requires every program qubit to be mapped.
   [[nodiscard]] ArrayRef<size_t> getProgramToHardware() const;
 
   /// Compare two layouts for equality.
   [[nodiscard]] bool operator==(const Layout& other) const {
-    return programToHardware_ == other.programToHardware_;
+    return programToHardware_ == other.programToHardware_ &&
+           hardwareToProgram_ == other.hardwareToProgram_;
   }
 
 private:
-  /// Construct a layout with `nqubits`.
-  explicit Layout(const size_t nqubits)
-      : programToHardware_(nqubits), hardwareToProgram_(nqubits) {}
+  Layout(size_t nProgramQubits, size_t nHardwareQubits);
 
   /// Maps a program qubit index to its hardware index.
+  /// Size equals `nProgramQubits()`.
   SmallVector<size_t> programToHardware_;
   /// Maps a hardware qubit index to its program index.
+  /// Size equals `nHardwareQubits()`.
   SmallVector<size_t> hardwareToProgram_;
 };
 } // namespace mlir::qco

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

@@ -661,7 +661,7 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
 
     // Create and save static qubit operations.
     rewriter.setInsertionPointToStart(&body.front());
-    for (size_t hw = 0; hw < layout.nqubits(); ++hw) {
+    for (size_t hw = 0; hw < layout.nHardwareQubits(); ++hw) {
       const auto site = target->siteForVertex(hw);
       auto op = StaticOp::create(rewriter, body.getLoc(), site);
       staticQubits.emplace_back(op.getQubit());
@@ -713,7 +713,7 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
     // Create sinks for remaining, unused, static qubits.
 
     rewriter.setInsertionPoint(body.back().getTerminator());
-    for (size_t prog = wires.size(); prog < layout.nqubits(); ++prog) {
+    for (size_t prog = wires.size(); prog < layout.nHardwareQubits(); ++prog) {
       const auto hw = layout.getHardwareIndex(prog);
       auto qubit = staticQubits[hw];
 
@@ -754,7 +754,8 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
       trials.emplace_back(
           RoutingBundle{.wires = wires,
                         .infos = infos,
-                        .layout = Layout::random(target->numQubits(), rng())});
+                        .layout = Layout::random(target->numQubits(),
+                                                 target->numQubits(), rng())});
     }
 
     parallelForEach(&getContext(), trials, [&, this](Trial& t) {

mlir/lib/Dialect/QCO/Utils/Layout.cpp

@@ -17,16 +17,35 @@
 #include <llvm/Support/ErrorHandling.h>
 #include <mlir/Support/LLVM.h>
 
-#include <algorithm>
 #include <cassert>
 #include <cstddef>
+#include <limits>
 #include <random>
 
 namespace mlir::qco {
-Layout Layout::random(const size_t nqubits, const size_t seed) {
-  auto mapping = llvm::to_vector(llvm::seq(nqubits));
-  llvm::shuffle(mapping.begin(), mapping.end(), std::mt19937_64{seed});
-  return fromMapping(mapping);
+
+namespace {
+/// Sentinel stored in `programToHardware_` and `hardwareToProgram_` entries
+/// that do not currently hold a valid index.
+constexpr size_t UNMAPPED = std::numeric_limits<size_t>::max();
+} // namespace
+
+Layout::Layout(const size_t nProgramQubits, const size_t nHardwareQubits)
+    : programToHardware_(nProgramQubits, UNMAPPED),
+      hardwareToProgram_(nHardwareQubits, UNMAPPED) {}
+
+Layout Layout::random(const size_t nProgramQubits, const size_t nHardwareQubits,
+                      const size_t seed) {
+  assert(nProgramQubits <= nHardwareQubits &&
+         "cannot map more program qubits than hardware qubits");
+  auto hwIndices = llvm::to_vector(llvm::seq(nHardwareQubits));
+  llvm::shuffle(hwIndices.begin(), hwIndices.end(), std::mt19937_64{seed});
+
+  Layout layout(nProgramQubits, nHardwareQubits);
+  for (size_t prog = 0; prog < nProgramQubits; ++prog) {
+    layout.add(prog, hwIndices[prog]);
+  }
+  return layout;
 }
 
 Layout Layout::fromMapping(ArrayRef<size_t> mapping) {
@@ -38,7 +57,7 @@ Layout Layout::fromMapping(ArrayRef<size_t> mapping) {
     seen.set(hw);
   }
 
-  Layout layout(mapping.size());
+  Layout layout(mapping.size(), mapping.size());
   for (const auto [prog, hw] : enumerate(mapping)) {
     layout.add(prog, hw);
   }
@@ -48,33 +67,58 @@ Layout Layout::fromMapping(ArrayRef<size_t> mapping) {
 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");
+  assert(programToHardware_[prog] == UNMAPPED &&
+         "program index already mapped");
+  assert(hardwareToProgram_[hw] == UNMAPPED && "hardware index already mapped");
   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];
+  const auto prog = hardwareToProgram_[hw];
+  assert(prog != UNMAPPED && "hardware index not mapped");
+  return prog;
 }
 
 size_t Layout::getHardwareIndex(const size_t prog) const {
   assert(prog < programToHardware_.size() && "program index out of bounds");
-  return programToHardware_[prog];
+  const auto hw = programToHardware_[prog];
+  assert(hw != UNMAPPED && "program index not mapped");
+  return hw;
+}
+
+bool Layout::hasProgramAt(const size_t hw) const {
+  assert(hw < hardwareToProgram_.size() && "hardware index out of bounds");
+  return hardwareToProgram_[hw] != UNMAPPED;
 }
 
 void Layout::swap(const size_t hwA, const size_t hwB) {
   assert(hwA < hardwareToProgram_.size() && "hardware index out of bounds");
   assert(hwB < hardwareToProgram_.size() && "hardware index out of bounds");
-  const auto progA = hardwareToProgram_[hwA];
-  const auto progB = hardwareToProgram_[hwB];
-
-  std::swap(hardwareToProgram_[hwA], hardwareToProgram_[hwB]);
-  std::swap(programToHardware_[progA], programToHardware_[progB]);
+  if (hwA == hwB) {
+    return;
+  }
+  const size_t progA = hardwareToProgram_[hwA];
+  const size_t progB = hardwareToProgram_[hwB];
+  assert(progA != UNMAPPED && "hardware index not mapped");
+  assert(progB != UNMAPPED && "hardware index not mapped");
+  hardwareToProgram_[hwA] = progB;
+  hardwareToProgram_[hwB] = progA;
+  programToHardware_[progA] = hwB;
+  programToHardware_[progB] = hwA;
 }
 
-size_t Layout::nqubits() const { return programToHardware_.size(); }
+size_t Layout::nProgramQubits() const { return programToHardware_.size(); }
+
+size_t Layout::nHardwareQubits() const { return hardwareToProgram_.size(); }
 
 ArrayRef<size_t> Layout::getProgramToHardware() const {
+#ifndef NDEBUG
+  for (const size_t hw : programToHardware_) {
+    assert(hw != UNMAPPED && "program qubit not mapped");
+  }
+#endif
   return programToHardware_;
 }