Scheduling based on mix frames (#300)

This PR introduces scheduling of a quantum circuit based on mix frames
in order to reduce the overall duration of the scheduled pulse sequence.

include/Dialect/Pulse/Transforms/Scheduling.h

@@ -1,6 +1,6 @@
 //===- scheduling.h --- quantum circuits pulse scheduling -------*- C++ -*-===//
 //
-// (C) Copyright IBM 2023.
+// (C) Copyright IBM 2023, 2024.
 //
 // This code is part of Qiskit.
 //
@@ -15,7 +15,7 @@
 //===----------------------------------------------------------------------===//
 ///
 ///  This file implements the pass for scheduling the quantum circuits at pulse
-///  level, based on the availability of involved ports
+///  level, based on the availability of involved mix frames
 ///
 //===----------------------------------------------------------------------===//
 
@@ -27,6 +27,9 @@
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Pass/Pass.h"
 
+#include <unordered_map>
+#include <unordered_set>
+
 namespace mlir::pulse {
 
 struct QuantumCircuitPulseSchedulingPass
@@ -68,13 +71,20 @@ struct QuantumCircuitPulseSchedulingPass
       llvm::cl::value_desc("delay"), llvm::cl::init(0)};
 
 private:
-  // map to keep track of next availability of ports
-  std::map<std::string, int> portNameToNextAvailabilityMap;
+  // map to keep track of next availability of mix frames
+  std::unordered_map<unsigned int, int64_t> mixFrameToNextAvailabilityMap;
 
   void scheduleAlap(mlir::pulse::CallSequenceOp quantumCircuitCallSequenceOp);
-  int getNextAvailableTimeOfPorts(mlir::ArrayAttr ports);
-  void updatePortAvailabilityMap(mlir::ArrayAttr ports,
-                                 int updatedAvailableTime);
+  // returns an unordered set of block argument numbers of the mix frames of a
+  // quantum gate; here block refers to the current block that includes the call
+  // op
+  std::unordered_set<unsigned int> getMixFramesBlockArgNums(
+      mlir::pulse::CallSequenceOp quantumGateCallSequenceOp);
+  int64_t getNextAvailableTimeOfMixFrames(
+      std::unordered_set<unsigned int> &mixFramesBlockArgNums);
+  void updateMixFrameAvailabilityMap(
+      std::unordered_set<unsigned int> &mixFramesBlockArgNums,
+      int64_t updatedAvailableTime);
   bool sequenceOpIncludeCapture(mlir::pulse::SequenceOp quantumGateSequenceOp);
   llvm::StringMap<Operation *> symbolMap;
   mlir::pulse::SequenceOp

lib/Dialect/Pulse/Transforms/Scheduling.cpp

@@ -1,6 +1,6 @@
 //===- Scheduling.cpp ---  quantum circuits pulse scheduling ----*- C++ -*-===//
 //
-// (C) Copyright IBM 2023.
+// (C) Copyright IBM 2023, 2024.
 //
 // This code is part of Qiskit.
 //
@@ -15,20 +15,23 @@
 //===----------------------------------------------------------------------===//
 ///
 ///  This file implements the pass for scheduling the quantum circuits at pulse
-///  level, based on the availability of involved ports
+///  level, based on the availability of involved mix frames
 ///
 //===----------------------------------------------------------------------===//
 
 #include "Dialect/Pulse/Transforms/Scheduling.h"
 #include "Dialect/Pulse/IR/PulseInterfaces.h"
 #include "Dialect/Pulse/IR/PulseOps.h"
+#include "Dialect/Pulse/IR/PulseTypes.h"
 #include "Dialect/QUIR/Utils/Utils.h"
 
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Value.h"
 #include "mlir/IR/Visitors.h"
 #include "mlir/Support/LLVM.h"
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
@@ -37,6 +40,8 @@
 #include <cassert>
 #include <cstdint>
 #include <string>
+#include <unordered_map>
+#include <unordered_set>
 #include <utility>
 
 #define DEBUG_TYPE "SchedulingDebug"
@@ -89,7 +94,7 @@ void QuantumCircuitPulseSchedulingPass::scheduleAlap(
   LLVM_DEBUG(llvm::dbgs() << "\nscheduling " << sequenceName << "\n");
 
   int totalDurationOfQuantumCircuitNegative = 0;
-  portNameToNextAvailabilityMap.clear();
+  mixFrameToNextAvailabilityMap.clear();
 
   // get the MLIR block of the quantum circuit
   auto quantumCircuitSequenceOpBlock =
@@ -113,14 +118,9 @@ void QuantumCircuitPulseSchedulingPass::scheduleAlap(
       LLVM_DEBUG(llvm::dbgs() << "\tprocessing inner sequence "
                               << quantumGateSequenceName << "\n");
 
-      // find ports of the quantum gate SequenceOp
-      auto portsOrError =
-          PulseOpSchedulingInterface::getPorts(quantumGateSequenceOp);
-      if (auto err = portsOrError.takeError()) {
-        quantumGateSequenceOp.emitError() << toString(std::move(err));
-        signalPassFailure();
-      }
-      auto ports = portsOrError.get();
+      // find the block argument numbers of the mix frames
+      std::unordered_set<unsigned int> mixFramesBlockArgNums =
+          getMixFramesBlockArgNums(quantumGateCallSequenceOp);
 
       // find duration of the quantum gate callSequenceOp
       llvm::Expected<uint64_t> durOrError =
@@ -133,25 +133,26 @@ void QuantumCircuitPulseSchedulingPass::scheduleAlap(
       LLVM_DEBUG(llvm::dbgs() << "\t\tduration "
                               << quantumGateCallSequenceOpDuration << "\n");
 
-      // find next available time for all the ports
-      const int64_t nextAvailableTimeOfAllPorts =
-          getNextAvailableTimeOfPorts(ports);
+      // find next available time for all the mix frames
+      const int64_t nextAvailableTimeOfAllMixFrames =
+          getNextAvailableTimeOfMixFrames(mixFramesBlockArgNums);
       LLVM_DEBUG(llvm::dbgs() << "\t\tnext availability is at "
-                              << nextAvailableTimeOfAllPorts << "\n");
+                              << nextAvailableTimeOfAllMixFrames << "\n");
 
       // find the updated available time, i.e., when the current quantum gate
       // will be scheduled
       int64_t updatedAvailableTime =
-          nextAvailableTimeOfAllPorts - quantumGateCallSequenceOpDuration;
+          nextAvailableTimeOfAllMixFrames - quantumGateCallSequenceOpDuration;
       // set the timepoint of quantum gate
       PulseOpSchedulingInterface::setTimepoint(quantumGateCallSequenceOp,
                                                updatedAvailableTime);
       LLVM_DEBUG(llvm::dbgs() << "\t\tcurrent gate scheduled at "
                               << updatedAvailableTime << "\n");
-      // update the port availability map
+      // update the mix frame availability map
       if (sequenceOpIncludeCapture(quantumGateSequenceOp))
         updatedAvailableTime -= PRE_MEASURE_BUFFER_DELAY;
-      updatePortAvailabilityMap(ports, updatedAvailableTime);
+      updateMixFrameAvailabilityMap(mixFramesBlockArgNums,
+                                    updatedAvailableTime);
 
       // keep track of total duration of the quantum circuit
       if (updatedAvailableTime < totalDurationOfQuantumCircuitNegative)
@@ -178,30 +179,37 @@ void QuantumCircuitPulseSchedulingPass::scheduleAlap(
                                            totalDurationOfQuantumCircuit);
 }
 
-int QuantumCircuitPulseSchedulingPass::getNextAvailableTimeOfPorts(
-    mlir::ArrayAttr ports) {
-  int nextAvailableTimeOfAllPorts = 0;
-  for (auto attr : ports) {
-    const std::string portName = attr.dyn_cast<StringAttr>().getValue().str();
-    if (portName.empty())
-      continue;
-    if (portNameToNextAvailabilityMap.find(portName) !=
-        portNameToNextAvailabilityMap.end()) {
-      if (portNameToNextAvailabilityMap[portName] < nextAvailableTimeOfAllPorts)
-        nextAvailableTimeOfAllPorts = portNameToNextAvailabilityMap[portName];
+int64_t QuantumCircuitPulseSchedulingPass::getNextAvailableTimeOfMixFrames(
+    std::unordered_set<unsigned int> &mixFramesBlockArgNums) {
+  int64_t nextAvailableTimeOfAllMixFrames = 0;
+  for (auto id : mixFramesBlockArgNums)
+    if (mixFrameToNextAvailabilityMap.find(id) !=
+            mixFrameToNextAvailabilityMap.end() &&
+        mixFrameToNextAvailabilityMap[id] < nextAvailableTimeOfAllMixFrames)
+      nextAvailableTimeOfAllMixFrames = mixFrameToNextAvailabilityMap[id];
+  return nextAvailableTimeOfAllMixFrames;
+}
+
+std::unordered_set<unsigned int>
+QuantumCircuitPulseSchedulingPass::getMixFramesBlockArgNums(
+    mlir::pulse::CallSequenceOp quantumGateCallSequenceOp) {
+  std::unordered_set<unsigned int> mixFramesBlockArgNums;
+  for (auto const &argumentResult :
+       llvm::enumerate(quantumGateCallSequenceOp.getOperands())) {
+    auto argType = argumentResult.value().getType();
+    if (auto mixFrameType = argType.dyn_cast<mlir::pulse::MixedFrameType>()) {
+      mixFramesBlockArgNums.insert(
+          argumentResult.value().dyn_cast<BlockArgument>().getArgNumber());
     }
   }
-  return nextAvailableTimeOfAllPorts;
+  return mixFramesBlockArgNums;
 }
 
-void QuantumCircuitPulseSchedulingPass::updatePortAvailabilityMap(
-    mlir::ArrayAttr ports, int updatedAvailableTime) {
-  for (auto attr : ports) {
-    const std::string portName = attr.dyn_cast<StringAttr>().getValue().str();
-    if (portName.empty())
-      continue;
-    portNameToNextAvailabilityMap[portName] = updatedAvailableTime;
-  }
+void QuantumCircuitPulseSchedulingPass::updateMixFrameAvailabilityMap(
+    std::unordered_set<unsigned int> &mixFramesBlockArgNums,
+    int64_t updatedAvailableTime) {
+  for (auto id : mixFramesBlockArgNums)
+    mixFrameToNextAvailabilityMap[id] = updatedAvailableTime;
 }
 
 bool QuantumCircuitPulseSchedulingPass::sequenceOpIncludeCapture(

releasenotes/notes/mixframe-scheduling-9f2a5e5ba922e6a0.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Scheduling of a quantum circuit based on the availability of mix frames
+    is added in order to reduce the overall duration of the scheduled pulse
+    sequence compared to that of port-based scheduling.