🚸 Remove One-Tensor Constraint From Mapping Pass

CHANGELOG.md

@@ -16,7 +16,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel
 - ✨ Add a `hadamard-lifting` pass for lifting Hadamard gates above Pauli gates ([#1605]) ([**@lirem101**], [**@burgholzer**])
 - ✨ Add a `merge-single-qubit-rotation-gates` pass for merging consecutive rotation gates using quaternions ([#1407], [#1674]) ([**@J4MMlE**], [**@denialhaag**], [**@MatthiasReumann**])
 - ✨ Add conversions between `jeff` and QCO ([#1479], [#1548], [#1565], [#1637], [#1676], [#1706]) ([**@denialhaag**], [**@burgholzer**])
-- ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568], [#1581], [#1583], [#1588], [#1600], [#1664], [#1709], [#1716]) ([**@MatthiasReumann**], [**@burgholzer**])
+- ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568], [#1581], [#1583], [#1588], [#1600], [#1664], [#1709], [#1716], [#1748]) ([**@MatthiasReumann**], [**@burgholzer**])
 - ✨ 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], [#1567], [#1569], [#1570], [#1572], [#1573], [#1580], [#1602], [#1620], [#1623], [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1717], [#1728], [#1730])
   ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**])
@@ -402,6 +402,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool
 
 <!-- PR links -->
 
+[#1748]: https://github.com/munich-quantum-toolkit/core/pull/1748
 [#1737]: https://github.com/munich-quantum-toolkit/core/pull/1737
 [#1730]: https://github.com/munich-quantum-toolkit/core/pull/1730
 [#1728]: https://github.com/munich-quantum-toolkit/core/pull/1728

mlir/include/mlir/Dialect/QTensor/Utils/TensorIterator.h

@@ -28,9 +28,11 @@ class [[nodiscard]] TensorIterator {
   using difference_type = std::ptrdiff_t;
   using value_type = Operation*;
 
-  TensorIterator() : op_(nullptr), tensor_(nullptr), isSentinel_(false) {}
+  TensorIterator()
+      : op_(nullptr), tensor_(nullptr), isFinal_(false), isSentinel_(false) {}
   explicit TensorIterator(TypedValue<RankedTensorType> tensor)
-      : op_(tensor.getDefiningOp()), tensor_(tensor), isSentinel_(false) {}
+      : op_(tensor.getDefiningOp()), tensor_(tensor), isFinal_(false),
+        isSentinel_(false) {}
 
   /// @returns the operation the iterator points to.
   [[nodiscard]] Operation* operation() const { return op_; }
@@ -81,6 +83,7 @@ class [[nodiscard]] TensorIterator {
 
   Operation* op_;
   TypedValue<RankedTensorType> tensor_;
+  bool isFinal_;
   bool isSentinel_;
 };
 } // namespace mlir::qtensor

mlir/lib/Dialect/QCO/Transforms/CMakeLists.txt

@@ -15,6 +15,7 @@ add_mlir_library(
   PRIVATE
   MLIRQCODialect
   MLIRQCOUtils
+  MLIRQTensorUtils
   MLIRArithDialect
   MLIRMathDialect
   MLIRSCFUtils

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

@@ -16,6 +16,7 @@
 #include "mlir/Dialect/QCO/Utils/Drivers.h"
 #include "mlir/Dialect/QCO/Utils/WireIterator.h"
 #include "mlir/Dialect/QTensor/IR/QTensorOps.h"
+#include "mlir/Dialect/QTensor/Utils/TensorIterator.h"
 
 #include <llvm/ADT/ArrayRef.h>
 #include <llvm/ADT/PriorityQueue.h>
@@ -453,9 +454,17 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
   /**
    * @brief Collect wires of the quantum computation before placement.
    * @details
-   * The mapping pass currently assumes that the quantum computation consists of
-   * a single quantum tensor. The required qubits are extracted and inserted "in
-   * one go" at the beginning and the end of the function, respectively.
+   * The mapping pass currently assumes that the quantum computation allocates
+   * all tensors at the start of the function. The required qubits are extracted
+   * from these tensors and used for the computation. Finally, the qubits are
+   * inserted back into the tensors at the end of the function.
+   * Thus, a valid program has the following structure:
+   *
+   *    T ⨉ [qtensor::AllocOp]
+   *  → N ⨉ [qtensor::ExtractOp]
+   *  → (Computation)
+   *  → N ⨉ [qtensor::InsertOp]
+   *  → T ⨉ [qtensor::DeallocOp]
    *
    * @returns a vector of wire iterator, or failure() if any of the above
    * assumptions are violated.
@@ -467,44 +476,26 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
       return failure();
     }
 
-    const auto tensors = func.getOps<qtensor::AllocOp>();
-    if (range_size(tensors) != 1) {
-      func.emitError() << "must contain a single qtensor.alloc operation";
-      return failure();
-    }
-
-    bool inExtractPhase = true;
     SmallVector<WireIterator> wires;
-    Value tensor = (*tensors.begin()).getResult();
-
-    while (true) {
-      assert(tensor.hasOneUse() && "getComputation: expected linear typing");
-      Operation* curr = *(tensor.user_begin());
-
-      if (isa<DeallocOp>(curr)) {
-        break;
-      }
-
-      if (auto extractOp = dyn_cast<ExtractOp>(curr)) {
-        if (!inExtractPhase) {
-          func.emitError() << "must extract and insert all qubits at once.";
-          return failure();
+    for (auto tensor : func.getOps<qtensor::AllocOp>()) {
+      bool isInitPhase = true;
+      TensorIterator it(tensor.getResult());
+      for (; it != std::default_sentinel; ++it) {
+        if (auto extract = dyn_cast<ExtractOp>(it.operation())) {
+          if (!isInitPhase) {
+            func.emitError() << "must extract and insert all qubits at once.";
+            return failure();
+          }
+          wires.emplace_back(extract.getResult());
+          continue;
         }
-        tensor = extractOp.getOutTensor();
-        wires.emplace_back(extractOp.getResult());
-        continue;
-      }
 
-      if (auto insertOp = dyn_cast<InsertOp>(curr)) {
-        inExtractPhase = false;
-        tensor = insertOp.getResult();
-        continue;
+        if (isa<InsertOp>(it.operation())) {
+          isInitPhase = false;
+          continue;
+        }
       }
-
-      report_fatal_error("unknown op in def-use chain: " +
-                         curr->getName().getStringRef());
     }
-
     return wires;
   }
 
@@ -533,50 +524,38 @@ struct MappingPass : impl::MappingPassBase<MappingPass> {
     // Replace extract ops and collect in program-qubit order.
     SmallVector<WireIterator> placedWires(layout.nqubits());
 
-    const auto tensors = func.getOps<qtensor::AllocOp>();
-    assert(range_size(tensors) == 1 && "place: expected exactly one tensor");
-
-    qtensor::AllocOp alloc = *(tensors.begin());
-    const Value tensor = alloc.getResult();
-    assert(tensor.hasOneUse() && "place: expected linear typing");
-
-    size_t prog = 0;
-    while (true) {
-      Operation* curr = *(tensor.user_begin());
-      if (isa<DeallocOp>(curr)) {
-        rewriter.eraseOp(curr);
-        break;
+    size_t prog = 0UL;
+    for (auto alloc : make_early_inc_range(func.getOps<qtensor::AllocOp>())) {
+      TensorIterator it(alloc.getResult());
+      while (it != std::default_sentinel) {
+        // Get the operation and early increment to avoid issues after erasure.
+        Operation* curr = it.operation();
+        ++it;
+
+        TypeSwitch<Operation*>(curr)
+            .Case<ExtractOp>([&](auto op) {
+              const auto hw = layout.getHardwareIndex(prog);
+              const auto qubit = staticOps[hw].getQubit();
+
+              rewriter.replaceAllUsesWith(op.getResult(), qubit);
+              rewriter.replaceAllUsesWith(op.getOutTensor(), op.getTensor());
+              rewriter.eraseOp(op);
+
+              placedWires[prog] = WireIterator(qubit);
+              ++prog;
+            })
+            .Case<InsertOp>([&](auto op) {
+              rewriter.setInsertionPointAfter(op);
+              SinkOp::create(rewriter, op.getLoc(), op.getScalar());
+              rewriter.replaceAllUsesWith(op.getResult(), op.getDest());
+              rewriter.eraseOp(op);
+            })
+            .Case<DeallocOp>([&](auto op) { rewriter.eraseOp(op); });
       }
 
-      TypeSwitch<Operation*>(curr)
-          .Case<ExtractOp>([&](ExtractOp op) {
-            const auto hw = layout.getHardwareIndex(prog);
-            const auto qubit = staticOps[hw].getQubit();
-
-            placedWires[prog] = WireIterator(qubit);
-
-            rewriter.replaceAllUsesWith(op.getResult(), qubit);
-            rewriter.replaceAllUsesWith(op.getOutTensor(), tensor);
-            rewriter.eraseOp(op);
-
-            ++prog;
-          })
-          .Case<InsertOp>([&](InsertOp op) {
-            rewriter.setInsertionPointAfter(op);
-
-            SinkOp::create(rewriter, op.getLoc(), op.getScalar());
-
-            rewriter.replaceAllUsesWith(op.getResult(), tensor);
-            rewriter.eraseOp(op);
-          })
-          .Default([&](Operation* op) {
-            report_fatal_error("unknown op in def-use chain: " +
-                               op->getName().getStringRef());
-          });
+      rewriter.eraseOp(alloc);
     }
 
-    rewriter.eraseOp(alloc);
-
     // Create sinks for remaining, unused, static qubits.
     rewriter.setInsertionPoint(func.getFunctionBody().back().getTerminator());
     for (; prog < layout.nqubits(); ++prog) {

mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp

@@ -44,13 +44,19 @@ void TensorIterator::forward() {
     return;
   }
 
+  // After the final operation comes the sentinel.
+  if (isFinal_) {
+    isSentinel_ = true;
+    return;
+  }
+
   // Find the user-operation of the tensor SSA value.
   assert(tensor_.hasOneUse() && "expected linear typing");
   op_ = *(tensor_.user_begin());
 
   // The following operations define the end of the tensor's life-chain.
   if (isa<DeallocOp, scf::YieldOp, qco::YieldOp>(op_)) {
-    isSentinel_ = true;
+    isFinal_ = true;
     return;
   }
 
@@ -80,6 +86,7 @@ void TensorIterator::backward() {
   // If the iterator is a sentinel, reactivate the iterator.
   if (isSentinel_) {
     isSentinel_ = false;
+    isFinal_ = true;
     return;
   }
 
@@ -92,6 +99,7 @@ void TensorIterator::backward() {
   // For these operations, tensor_ is an OpOperand. Hence, only get the def-op.
   if (isa<DeallocOp, scf::YieldOp, qco::YieldOp>(op_)) {
     op_ = tensor_.getDefiningOp();
+    isFinal_ = false;
     return;
   }
 
@@ -136,6 +144,7 @@ void TensorIterator::backward() {
   // If the current tensor SSA value is a BlockArgument (no defining op), the
   // operation will be a nullptr.
   op_ = tensor_.getDefiningOp();
+  isFinal_ = false;
 }
 
 static_assert(std::bidirectional_iterator<TensorIterator>);

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

@@ -135,37 +135,6 @@ TEST_P(MappingPassTest, NoQubitAllocations) {
   ASSERT_TRUE(res.failed());
 }
 
-TEST_P(MappingPassTest, NoTwoTensors) {
-  const auto& device = GetParam();
-
-  QCOProgramBuilder builder(context.get());
-  builder.initialize();
-
-  Value tensor0 = builder.qtensorAlloc(1);
-  Value tensor1 = builder.qtensorAlloc(1);
-
-  Value q0;
-  std::tie(tensor0, q0) = builder.qtensorExtract(tensor0, 0);
-  Value q1;
-  std::tie(tensor1, q1) = builder.qtensorExtract(tensor1, 0);
-
-  q0 = builder.h(q0);
-  q1 = builder.h(q1);
-
-  std::tie(q0, q1) = builder.cx(q0, q1);
-
-  tensor0 = builder.qtensorInsert(q0, tensor0, 0);
-  tensor1 = builder.qtensorInsert(q1, tensor1, 0);
-
-  builder.qtensorDealloc(tensor0);
-  builder.qtensorDealloc(tensor1);
-
-  auto m = builder.finalize();
-  auto res = runPass(m.get(), device, MappingPassOptions{});
-
-  ASSERT_TRUE(res.failed());
-}
-
 TEST_P(MappingPassTest, NoExtractAfterInsert) {
   const auto& device = GetParam();
 
@@ -259,25 +228,26 @@ TEST_P(MappingPassTest, Sabre) {
   QCOProgramBuilder builder(context.get());
   builder.initialize();
 
-  Value tensor = builder.qtensorAlloc(6);
+  Value tensorUp = builder.qtensorAlloc(4);
+  Value tensorDown = builder.qtensorAlloc(2);
 
   Value q0;
-  std::tie(tensor, q0) = builder.qtensorExtract(tensor, 0);
+  std::tie(tensorUp, q0) = builder.qtensorExtract(tensorUp, 0);
 
   Value q1;
-  std::tie(tensor, q1) = builder.qtensorExtract(tensor, 1);
+  std::tie(tensorUp, q1) = builder.qtensorExtract(tensorUp, 1);
 
   Value q2;
-  std::tie(tensor, q2) = builder.qtensorExtract(tensor, 2);
+  std::tie(tensorUp, q2) = builder.qtensorExtract(tensorUp, 2);
 
   Value q3;
-  std::tie(tensor, q3) = builder.qtensorExtract(tensor, 3);
+  std::tie(tensorUp, q3) = builder.qtensorExtract(tensorUp, 3);
 
   Value q4;
-  std::tie(tensor, q4) = builder.qtensorExtract(tensor, 4);
+  std::tie(tensorDown, q4) = builder.qtensorExtract(tensorDown, 0);
 
   Value q5;
-  std::tie(tensor, q5) = builder.qtensorExtract(tensor, 5);
+  std::tie(tensorDown, q5) = builder.qtensorExtract(tensorDown, 1);
 
   q0 = builder.h(q0);
   q1 = builder.h(q1);
@@ -329,13 +299,14 @@ TEST_P(MappingPassTest, Sabre) {
   std::tie(q4, c4) = builder.measure(q4);
   std::tie(q5, c5) = builder.measure(q5);
 
-  tensor = builder.qtensorInsert(q0, tensor, 0);
-  tensor = builder.qtensorInsert(q1, tensor, 1);
-  tensor = builder.qtensorInsert(q2, tensor, 2);
-  tensor = builder.qtensorInsert(q3, tensor, 3);
-  tensor = builder.qtensorInsert(q4, tensor, 4);
-  tensor = builder.qtensorInsert(q5, tensor, 5);
-  builder.qtensorDealloc(tensor);
+  tensorUp = builder.qtensorInsert(q0, tensorUp, 0);
+  tensorUp = builder.qtensorInsert(q1, tensorUp, 1);
+  tensorUp = builder.qtensorInsert(q2, tensorUp, 2);
+  tensorUp = builder.qtensorInsert(q3, tensorUp, 3);
+  tensorDown = builder.qtensorInsert(q4, tensorDown, 0);
+  tensorDown = builder.qtensorInsert(q5, tensorDown, 1);
+  builder.qtensorDealloc(tensorUp);
+  builder.qtensorDealloc(tensorDown);
 
   auto m = builder.finalize();
   auto res = runPass(m.get(), device, MappingPassOptions{});

mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp

@@ -144,6 +144,8 @@ TEST_F(TensorIteratorTest, Traversal) {
   ++it;
   ASSERT_EQ(it.operation(), *(tensor8.user_begin())); // qtensor.dealloc
   ASSERT_EQ(it.tensor(), nullptr);
+
+  ++it;
   ASSERT_EQ(it, std::default_sentinel);
 
   ++it;
@@ -211,13 +213,15 @@ TEST_F(TensorIteratorTest, Traversal) {
   ASSERT_EQ(recIt.tensor(), tensorElse2);
 
   ++recIt;
-  ASSERT_EQ(recIt, std::default_sentinel);
   ASSERT_EQ(recIt.operation(), *(tensorElse2.user_begin())); // qco.yield
   ASSERT_EQ(recIt.tensor(), nullptr);
 
   ++recIt;
   ASSERT_EQ(recIt, std::default_sentinel);
 
+  ++recIt;
+  ASSERT_EQ(recIt, std::default_sentinel);
+
   --recIt;
   ASSERT_EQ(recIt.operation(), *(tensorElse2.user_begin())); // qco.yield
   ASSERT_EQ(recIt.tensor(), nullptr);