Prepare Layout API for partial injections

`Layout` can now describe partial program-to-hardware mappings.
`nqubits()` is split into `nProgramQubits()` and `nHardwareQubits()`.
`hasProgramAt(hw)` reports whether a hardware qubit currently carries a program qubit.
`swap(hwA, hwB)` accepts either side as unplaced and treats empty as a legitimate value that also gets swapped.
`random(nProg, nHw, seed)` places `nProg` program qubits on distinct hardware qubits drawn from `[0, nHw)`, so `nProg < nHw` leaves the remaining hardware slots unplaced.

The mapping pass keeps calling `random(n, n, rng())` and `nHardwareQubits()`, so behavior is unchanged.
Actually using `nProg < nHw` and the query API from inside the pass is left for a follow-up.

Part of #1867

Assisted-by: Claude Opus 4.7 via Claude Code
Signed-off-by: rturrado <rturrado@gmail.com>

Author: rturrado

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

@@ -21,25 +21,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.
+/// The mapping is a partial injection: at most `nProgramQubits` program qubits
+/// are placed on distinct hardware qubits drawn from `[0, nHardwareQubits)`.
+/// The two bounds may differ (usually `nProgramQubits <= nHardwareQubits`), so
+/// some hardware qubits may remain unplaced.
+///
+/// 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 an empty layout.
+  Layout() = default;
+
   /// Construct and return a random layout that places every program qubit
-  /// index in `[0, nqubits)` on a distinct hardware index in the same range.
-  static Layout random(size_t nqubits, size_t seed);
+  /// index in `[0, nProgramQubits)` on 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 program-to-hardware mapping,
   /// where mapping[prog] = hw.
+  /// Sets both `nProgramQubits` and `nHardwareQubits` to `mapping.size()`.
   static Layout fromMapping(ArrayRef<size_t> mapping);
 
   /// Insert a program:hardware index mapping.
-  /// Requires that neither `prog` nor `hw` has been mapped previously.
+  /// 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.
@@ -64,15 +75,24 @@ 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.
+  /// Either side may be currently unplaced; the empty state is treated as a
+  /// legitimate value that also gets swapped.
   void swap(size_t hwA, size_t hwB);
 
-  /// Return the number of qubits this layout was declared with.
-  [[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 as a materialized vector of
-  /// length `nqubits()`, where entry `prog` is the hardware index assigned
-  /// to program qubit `prog`.
+  /// length `nProgramQubits()`, where entry `prog` is the hardware index
+  /// assigned to program qubit `prog`. Requires every program qubit to be
+  /// placed.
   [[nodiscard]] SmallVector<size_t> getProgramToHardware() const;
 
   /// Compare two layouts for equality.
@@ -81,14 +101,16 @@ class Layout {
   }
 
 private:
-  /// Construct a layout with `nqubits`.
-  explicit Layout(const size_t nqubits) : nqubits_(nqubits) {}
+  Layout(const size_t nProgramQubits, const size_t nHardwareQubits)
+      : nProgramQubits_(nProgramQubits), nHardwareQubits_(nHardwareQubits) {}
 
   /// Maps a program qubit index to its hardware index.
   DenseMap<size_t, size_t> programToHardware_;
   /// Maps a hardware qubit index to its program index.
   DenseMap<size_t, size_t> hardwareToProgram_;
-  /// Number of qubits this layout was declared with.
-  size_t nqubits_ = 0;
+  /// Number of program qubits this layout was declared with.
+  size_t nProgramQubits_ = 0;
+  /// Number of hardware qubits this layout was declared with.
+  size_t nHardwareQubits_ = 0;
 };
 } // namespace mlir::qco

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

@@ -636,7 +636,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());
@@ -688,7 +688,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);
       const auto site = target->siteForVertex(hw);
       const auto qubit = staticQubits[site];
@@ -720,7 +720,7 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
               }
             })
             .Case<scf::ForOp>([&](scf::ForOp forOp) {
-              assert(qubits.size() == layout.nqubits());
+              assert(qubits.size() == layout.nHardwareQubits());
 
               llvm::for_each(getQubitValues(forOp.getInits()),
                              [&](Value v) { qubits.erase(v); });
@@ -741,7 +741,7 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
                   DenseSet<Value>(regionQubits.begin(), regionQubits.end()));
             })
             .Case<scf::WhileOp>([&](scf::WhileOp whileOp) {
-              assert(qubits.size() == layout.nqubits());
+              assert(qubits.size() == layout.nHardwareQubits());
 
               llvm::for_each(getQubitValues(whileOp.getInits()),
                              [&](Value v) { qubits.erase(v); });
@@ -767,7 +767,7 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
                   DenseSet<Value>(afterArgs.begin(), afterArgs.end()));
             })
             .Case<IfOp>([&](IfOp ifOp) {
-              assert(qubits.size() == layout.nqubits());
+              assert(qubits.size() == layout.nHardwareQubits());
 
               llvm::for_each(ifOp.getQubits(),
                              [&](Value v) { qubits.erase(v); });
@@ -790,7 +790,7 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
                   DenseSet<Value>(elseArgs.begin(), elseArgs.end()));
             })
             .Case<IndexSwitchOp>([&](IndexSwitchOp switchOp) {
-              assert(qubits.size() == layout.nqubits());
+              assert(qubits.size() == layout.nHardwareQubits());
 
               llvm::for_each(switchOp.getTargets(),
                              [&](Value value) { qubits.erase(value); });
@@ -850,7 +850,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) {
@@ -921,8 +922,8 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
     SmallVector<IndexPairType, 6> expansionSet;
 
     const auto materializeKey = [](const Layout& layout) {
-      SmallVector<size_t> key(layout.nqubits());
-      for (size_t prog = 0; prog < layout.nqubits(); ++prog) {
+      SmallVector<size_t> key(layout.nHardwareQubits());
+      for (size_t prog = 0; prog < layout.nHardwareQubits(); ++prog) {
         key[prog] = layout.getHardwareIndex(prog);
       }
       return key;
@@ -1088,7 +1089,7 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
   /// is the order (the permutation) of program-to-hardware indices.
   template <typename Range> static Layout vote(Range layouts) {
     assert(!layouts.empty() && "expected at least one layout");
-    const auto ncandidates = (*layouts.begin()).nqubits();
+    const auto ncandidates = (*layouts.begin()).nHardwareQubits();
 
     SmallVector<size_t> scores(ncandidates, 0);
     for (const Layout& layout : layouts) {

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

@@ -18,16 +18,25 @@
 #include <llvm/Support/ErrorHandling.h>
 #include <mlir/Support/LLVM.h>
 
+#include <algorithm>
 #include <cassert>
 #include <cstddef>
 #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);
+Layout Layout::random(const size_t nProgramQubits, const size_t nHardwareQubits,
+                      const size_t seed) {
+  assert(nProgramQubits <= nHardwareQubits &&
+         "cannot place 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) {
@@ -39,16 +48,16 @@ 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);
   }
   return layout;
 }
 
 void Layout::add(const size_t prog, const size_t hw) {
-  assert(prog < nqubits_ && "program index out of bounds");
-  assert(hw < nqubits_ && "hardware index out of bounds");
+  assert(prog < nProgramQubits_ && "program index out of bounds");
+  assert(hw < nHardwareQubits_ && "hardware index out of bounds");
   assert(!programToHardware_.contains(prog) && "program index already mapped");
   assert(!hardwareToProgram_.contains(hw) && "hardware index already mapped");
   programToHardware_[prog] = hw;
@@ -67,24 +76,43 @@ size_t Layout::getHardwareIndex(const size_t prog) const {
   return it->second;
 }
 
+bool Layout::hasProgramAt(const size_t hw) const {
+  return hardwareToProgram_.contains(hw);
+}
+
 void Layout::swap(const size_t hwA, const size_t hwB) {
+  assert(hwA < nHardwareQubits_ && "hardware index out of bounds");
+  assert(hwB < nHardwareQubits_ && "hardware index out of bounds");
+  if (hwA == hwB) {
+    return;
+  }
+  // Read the current value on each side (may be empty), then write each to the
+  // other side. Empty is treated as a legitimate value.
   const auto itA = hardwareToProgram_.find(hwA);
   const auto itB = hardwareToProgram_.find(hwB);
-  assert(itA != hardwareToProgram_.end() && "hardware index not mapped");
-  assert(itB != hardwareToProgram_.end() && "hardware index not mapped");
-  const auto progA = itA->second;
-  const auto progB = itB->second;
-  itA->second = progB;
-  itB->second = progA;
-  programToHardware_[progA] = hwB;
-  programToHardware_[progB] = hwA;
+  const bool hasA = itA != hardwareToProgram_.end();
+  const bool hasB = itB != hardwareToProgram_.end();
+  const size_t progA = hasA ? itA->second : 0;
+  const size_t progB = hasB ? itB->second : 0;
+  hardwareToProgram_.erase(hwA);
+  hardwareToProgram_.erase(hwB);
+  if (hasA) {
+    hardwareToProgram_[hwB] = progA;
+    programToHardware_[progA] = hwB;
+  }
+  if (hasB) {
+    hardwareToProgram_[hwA] = progB;
+    programToHardware_[progB] = hwA;
+  }
 }
 
-size_t Layout::nqubits() const { return nqubits_; }
+size_t Layout::nProgramQubits() const { return nProgramQubits_; }
+
+size_t Layout::nHardwareQubits() const { return nHardwareQubits_; }
 
 SmallVector<size_t> Layout::getProgramToHardware() const {
-  SmallVector<size_t> result(nqubits_);
-  for (size_t prog = 0; prog < nqubits_; ++prog) {
+  SmallVector<size_t> result(nProgramQubits_);
+  for (size_t prog = 0; prog < nProgramQubits_; ++prog) {
     result[prog] = getHardwareIndex(prog);
   }
   return result;