Add LabelPlayOpDurationsPass (#231)

This PR adds a LabelPlayOpDurationPass which labels all `pulse.play`
operations with the duration of the waveform that the operation plays.
This pass is to be used to provide duration information which is assumed
to be present in the `SchedulePortPass`.

---------

Co-authored-by: reza-j <[email protected]>

Author: bcdonovan

include/Dialect/Pulse/Transforms/LabelPlayOpDurations.h

@@ -0,0 +1,40 @@
+//===- LabelPlayOpDurations.cpp - Label PlayOps with Durations --*- C++ -*-===//
+//
+// (C) Copyright IBM 2024.
+//
+// This code is part of Qiskit.
+//
+// This code is licensed under the Apache License, Version 2.0 with LLVM
+// Exceptions. You may obtain a copy of this license in the LICENSE.txt
+// file in the root directory of this source tree.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+//
+//===----------------------------------------------------------------------===//
+///
+///  This file defines the pass for labeling pulse.play operations with the
+///  duration of the waveform being played.
+//===----------------------------------------------------------------------===//
+
+#ifndef PULSE_LABEL_PLAY_DURATION_H
+#define PULSE_LABEL_PLAY_DURATION_H
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir::pulse {
+
+class LabelPlayOpDurationsPass
+    : public PassWrapper<LabelPlayOpDurationsPass, OperationPass<ModuleOp>> {
+public:
+  void runOnOperation() override;
+
+  llvm::StringRef getArgument() const override;
+  llvm::StringRef getDescription() const override;
+  llvm::StringRef getName() const override;
+};
+} // namespace mlir::pulse
+
+#endif // PULSE_LABEL_PLAY_DURATION_H

include/Dialect/Pulse/Transforms/Passes.h

@@ -24,6 +24,7 @@
 #include "Conversion/QUIRToPulse/LoadPulseCals.h"
 #include "Conversion/QUIRToPulse/QUIRToPulse.h"
 #include "Dialect/Pulse/Transforms/ClassicalOnlyDetection.h"
+#include "Dialect/Pulse/Transforms/LabelPlayOpDurations.h"
 #include "Dialect/Pulse/Transforms/MergeDelays.h"
 #include "Dialect/Pulse/Transforms/RemoveUnusedArguments.h"
 #include "Dialect/Pulse/Transforms/SchedulePort.h"

lib/Dialect/Pulse/Transforms/CMakeLists.txt

@@ -13,6 +13,7 @@
 add_mlir_dialect_library(MLIRPulseTransforms
         ClassicalOnlyDetection.cpp
         InlineRegion.cpp
+        LabelPlayOpDurations.cpp
         MergeDelays.cpp
         Passes.cpp
         RemoveUnusedArguments.cpp

lib/Dialect/Pulse/Transforms/LabelPlayOpDurations.cpp

@@ -0,0 +1,87 @@
+//===- LabelPlayOpDurations.cpp - Label PlayOps with Durations --*- C++ -*-===//
+//
+// (C) Copyright IBM 2024.
+//
+// This code is part of Qiskit.
+//
+// This code is licensed under the Apache License, Version 2.0 with LLVM
+// Exceptions. You may obtain a copy of this license in the LICENSE.txt
+// file in the root directory of this source tree.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+//
+//===----------------------------------------------------------------------===//
+///
+///  This file implements the pass for labeling pulse.play operations with the
+///  duration of the waveform being played.
+///
+//===----------------------------------------------------------------------===//
+
+#include "Dialect/Pulse/Transforms/LabelPlayOpDurations.h"
+#include "Dialect/Pulse/IR/PulseInterfaces.h"
+#include "Dialect/Pulse/IR/PulseOps.h"
+
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+
+#include "llvm/ADT/StringRef.h"
+
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+using namespace mlir;
+using namespace mlir::pulse;
+
+void LabelPlayOpDurationsPass::runOnOperation() {
+
+  // all PlayOps are assumed to be inside of a pulse.sequence
+  // pass builds a mapping of sequence name , argument number to duration
+  // for all play operations using call_sequences
+  //
+  // pass then searches for all play operations and assigns the durations using
+  // the mapping
+
+  Operation *module = getOperation();
+
+  std::unordered_map<std::string, std::vector<uint64_t>> argumentToDuration;
+
+  module->walk([&](CallSequenceOp callSequenceOp) {
+    auto callee = callSequenceOp.getCallee().str();
+
+    for (const auto &operand : callSequenceOp->getOperands()) {
+
+      uint64_t duration = 0;
+      auto *defOp = operand.getDefiningOp();
+      if (defOp)
+        if (auto waveformOp = dyn_cast<Waveform_CreateOp>(defOp))
+          duration = waveformOp.getDuration(nullptr /*callSequenceOp*/).get();
+
+      argumentToDuration[callee].push_back(duration);
+    }
+  });
+
+  module->walk([&](PlayOp playOp) {
+    auto sequenceOp = playOp->getParentOfType<mlir::pulse::SequenceOp>();
+    auto sequenceStr = sequenceOp.getSymName().str();
+    auto wfArgNumber = playOp.getWfr().dyn_cast<BlockArgument>().getArgNumber();
+    auto duration = argumentToDuration[sequenceStr][wfArgNumber];
+    mlir::pulse::PulseOpSchedulingInterface::setDuration(playOp, duration);
+  });
+
+} // runOnOperation
+
+llvm::StringRef LabelPlayOpDurationsPass::getArgument() const {
+  return "pulse-label-play-op-duration";
+}
+
+llvm::StringRef LabelPlayOpDurationsPass::getDescription() const {
+  return "Label PlayOps with duration attributes";
+}
+
+llvm::StringRef LabelPlayOpDurationsPass::getName() const {
+  return "Label PlayOp Durations Pass";
+}

lib/Dialect/Pulse/Transforms/Passes.cpp

@@ -24,6 +24,7 @@
 #include "Conversion/QUIRToPulse/QUIRToPulse.h"
 
 #include "Dialect/Pulse/Transforms/ClassicalOnlyDetection.h"
+#include "Dialect/Pulse/Transforms/LabelPlayOpDurations.h"
 #include "Dialect/Pulse/Transforms/MergeDelays.h"
 #include "Dialect/Pulse/Transforms/RemoveUnusedArguments.h"
 #include "Dialect/Pulse/Transforms/SchedulePort.h"
@@ -37,6 +38,7 @@ namespace mlir::pulse {
 void pulsePassPipelineBuilder(OpPassManager &pm) {}
 
 void registerPulsePasses() {
+  PassRegistration<LabelPlayOpDurationsPass>();
   PassRegistration<LoadPulseCalsPass>();
   PassRegistration<QUIRToPulsePass>();
   PassRegistration<MergeDelayPass>();

releasenotes/notes/add-label-playop-duration-pass-2c79b1506b8b5854.yaml

@@ -0,0 +1,5 @@
+---
+other:
+  - |
+    Add LabelPlayOpsDurationPass to add a pulse.duration label to pulse.play
+    operations based on the duration of the waveform being played.