✨Map scf.for operations by restoration using the 4-approximation token swap algorithm

mlir/include/mlir/Dialect/QCO/Transforms/Mapping/Mapping.h

@@ -11,7 +11,6 @@
 #pragma once
 
 #include "mlir/Dialect/QCO/Transforms/Passes.h"
-#include "mlir/Dialect/QCO/Utils/Algorithms.h"
 
 #include <llvm/Support/LogicalResult.h>
 #include <mlir/IR/Region.h>
@@ -25,7 +24,8 @@ namespace mlir::qco {
  * @brief Create a mapping pass instance with the given target architecture.
  * @returns a pass object.
  */
-std::unique_ptr<Pass> createMappingPass(size_t nqubits, const Edges& coupling,
-                                        MappingPassOptions options);
+std::unique_ptr<Pass>
+createMappingPass(const llvm::DenseSet<std::pair<size_t, size_t>>&,
+                  MappingPassOptions);
 
 } // namespace mlir::qco

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

@@ -10,19 +10,20 @@
 
 #pragma once
 
-#include "mlir/Dialect/QCO/IR/QCODialect.h"
 #include "mlir/Dialect/QCO/IR/QCOInterfaces.h"
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
-#include "mlir/Dialect/QCO/Utils/Qubits.h"
 #include "mlir/Dialect/QCO/Utils/WireIterator.h"
 #include "mlir/Dialect/QTensor/IR/QTensorOps.h"
 
 #include <llvm/ADT/TypeSwitch.h>
+#include <llvm/Support/ErrorHandling.h>
+#include <mlir/Dialect/SCF/IR/SCF.h>
 #include <mlir/IR/Region.h>
 #include <mlir/IR/Value.h>
 #include <mlir/IR/Visitors.h>
 #include <mlir/Support/LLVM.h>
 
+#include <cassert>
 #include <cstddef>
 #include <functional>
 #include <iterator>
@@ -31,70 +32,22 @@
 
 namespace mlir::qco {
 
-using WalkProgramFn = function_ref<WalkResult(Operation*, Qubits&)>;
-
-/**
- * @brief Perform top-down non-recursive walk of all operations within a
- * region of a quantum program and apply a callback function.
- * @details The signature of the callback function is:
- *
- *     (Operation*, Qubits& q) -> WalkResult
- *
- * where the Qubits object tracks the front of qubit SSA values.
- * Depending on the template parameter, the callback is executed before or after
- * updating the Qubits state.
- * @param region The targeted region.
- * @param fn The callback function.
- * @returns success(), if all operations have been visited.
- */
-template <WalkOrder Order = WalkOrder::PreOrder>
-LogicalResult walkProgram(Region& region, const WalkProgramFn& fn) {
-  Qubits qubits;
-  for (Operation& curr : region.getOps()) {
-    if constexpr (Order == WalkOrder::PreOrder) {
-      if (fn(&curr, qubits).wasInterrupted()) {
-        return failure();
-      }
-    }
-
-    TypeSwitch<Operation*>(&curr)
-        .template Case<StaticOp>(
-            [&](StaticOp op) { qubits.add(op.getQubit(), op.getIndex()); })
-        .template Case<AllocOp>([&](AllocOp op) { qubits.add(op.getResult()); })
-        .template Case<UnitaryOpInterface>([&](UnitaryOpInterface& op) {
-          for (const auto& [prevV, nextV] :
-               llvm::zip(op.getInputQubits(), op.getOutputQubits())) {
-            const auto prevQ = cast<TypedValue<QubitType>>(prevV);
-            const auto nextQ = cast<TypedValue<QubitType>>(nextV);
-            qubits.remap(prevQ, nextQ);
-          }
-        })
-        .template Case<ResetOp>([&](ResetOp op) {
-          qubits.remap(op.getQubitIn(), op.getQubitOut());
-        })
-        .template Case<MeasureOp>([&](MeasureOp op) {
-          qubits.remap(op.getQubitIn(), op.getQubitOut());
-        })
-        .template Case<SinkOp>(
-            [&](SinkOp op) { qubits.remove(op.getQubit()); });
-
-    if constexpr (Order == WalkOrder::PostOrder) {
-      if (fn(&curr, qubits).wasInterrupted()) {
-        return failure();
-      }
-    }
-  }
-
-  return success();
-}
-
-using ReleasedOps = SmallVector<UnitaryOpInterface, 8>;
-using PendingWiresMap =
-    DenseMap<UnitaryOpInterface, SmallVector<std::size_t, 2>>;
+using ReleasedOps = SmallVector<Operation*, 8>;
+using PendingWiresMap = DenseMap<Operation*, SmallVector<size_t>>;
 
 struct IsReady {
   bool operator()(PendingWiresMap::value_type& kv) const {
-    return kv.second.size() == kv.first.getNumQubits();
+    const auto npending = kv.second.size();
+    return TypeSwitch<Operation*, bool>(kv.first)
+        .Case<UnitaryOpInterface>(
+            [&](auto& op) { return op.getNumQubits() == npending; })
+        .template Case<scf::ForOp>(
+            [&](auto& op) { return op.getInits().size() == npending; })
+        .Default([&](Operation* op) {
+          const auto name = op->getName().getStringRef();
+          reportFatalInternalError("unknown pending op: " + name);
+          return false;
+        });
   }
 };
 
@@ -138,20 +91,28 @@ LogicalResult walkProgramGraph(MutableArrayRef<WireIterator> wires,
   PendingWiresMap pending;
   pending.reserve(wires.size());
 
-  SmallVector<std::size_t> curr(wires.size());
+  SmallVector<size_t> curr(wires.size());
   std::iota(curr.begin(), curr.end(), 0UL);
 
-  SmallVector<std::size_t> next;
+  SmallVector<size_t> next;
   next.reserve(wires.size());
 
   while (!curr.empty()) {
-    for (std::size_t i : curr) {
+    for (size_t i : curr) {
       auto& it = wires[i];
       while (Traits::isActive(it)) {
+
+        if (it.qubit() != nullptr && isa<BlockArgument>(it.qubit())) {
+          std::ranges::advance(it, Traits::stride());
+          continue;
+        }
+
+        assert(it.operation() != nullptr);
+
         const auto res =
             TypeSwitch<Operation*, WalkResult>(it.operation())
-                .template Case<UnitaryOpInterface>([&](UnitaryOpInterface& op) {
-                  // If there are fewer wires than the qubit requires inputs,
+                .template Case<UnitaryOpInterface>([&](auto& op) {
+                  // If there are fewer wires than the unitary requires inputs,
                   // it's impossible to release the operation. Hence, fail.
                   if (op.getNumQubits() > wires.size()) {
                     return WalkResult::interrupt();
@@ -173,19 +134,45 @@ LogicalResult walkProgramGraph(MutableArrayRef<WireIterator> wires,
 
                   indices.emplace_back(i);
 
-                  return WalkResult::skip(); // Stop at multi-qubit gate.
+                  return WalkResult::skip(); // Stop at multi-qubit unitary.
+                })
+                .template Case<scf::ForOp>([&](scf::ForOp& op) {
+                  // If there are fewer wires than the loop requires inputs,
+                  // it's impossible to release the operation. Hence, fail.
+                  if (op.getInits().size() > wires.size()) {
+                    return WalkResult::interrupt();
+                  }
+
+                  if (op.getInits().size() == 1) {
+                    std::ranges::advance(it, Traits::stride());
+                    return WalkResult::advance();
+                  }
+
+                  // Insert the loop to the pending map.
+                  // The caller decides if this op should be released.
+                  const auto [mapIt, inserted] = pending.try_emplace(op);
+                  auto& indices = mapIt->second;
+
+                  if (inserted) {
+                    indices.reserve(op.getInits().size());
+                  }
+
+                  indices.emplace_back(i);
+
+                  return WalkResult::skip(); // Stop at multi-qubit loop.
                 })
                 // AllocOp, StaticOp, and qtensor::ExtractOp are only reachable
                 // on the forward path; backward isActive() halts before
                 // reaching them (decrementing at a source op is a no-op).
-                .template Case<AllocOp, StaticOp, qtensor::ExtractOp, ResetOp,
-                               MeasureOp, SinkOp, qtensor::InsertOp>([&](auto) {
+                .template Case<AllocOp, StaticOp, ResetOp, MeasureOp, SinkOp,
+                               YieldOp, scf::YieldOp, qtensor::ExtractOp,
+                               qtensor::InsertOp>([&](auto) {
                   std::ranges::advance(it, Traits::stride());
                   return WalkResult::advance();
                 })
                 .Default([&](Operation* op) -> WalkResult {
                   const auto name = op->getName().getStringRef();
-                  report_fatal_error("unknown op encountered: " + name);
+                  reportFatalInternalError("unknown op encountered: " + name);
                 });
 
         if (res.wasSkipped()) {
@@ -212,11 +199,11 @@ LogicalResult walkProgramGraph(MutableArrayRef<WireIterator> wires,
       }
     }
 
-    for (const UnitaryOpInterface& op : released) {
+    for (Operation* op : released) {
       const auto mapIt = pending.find(op);
       assert(mapIt != pending.end());
 
-      for (std::size_t i : mapIt->second) {
+      for (size_t i : mapIt->second) {
         std::ranges::advance(wires[i], Traits::stride());
         next.emplace_back(i);
       }