Draft of dependency graph generation algorithm

include/cudaq/Optimizer/Dialect/Quake/QuakeOps.h

@@ -70,6 +70,13 @@ void genericOpPrinter(mlir::OpAsmPrinter &_odsPrinter, mlir::Operation *op,
 // Utility functions to test the form of an operation.
 //===----------------------------------------------------------------------===//
 
+// Is \p op in the Quake dialect?
+inline bool isQuakeOperation(mlir::Operation *op) {
+  if (auto *dialect = op->getDialect())
+    return dialect->getNamespace().equals("quake");
+  return false;
+}
+
 namespace quake {
 /// Returns true if and only if any quantum operand has type `!quake.ref` or
 /// `!quake.veq`.

include/cudaq/Optimizer/Transforms/Passes.h

@@ -30,6 +30,8 @@ void registerAggressiveEarlyInlining();
 
 void registerUnrollingPipeline();
 
+void registerQubitManagementPipeline();
+
 std::unique_ptr<mlir::Pass> createApplyOpSpecializationPass();
 std::unique_ptr<mlir::Pass>
 createApplyOpSpecializationPass(bool computeActionOpt);

include/cudaq/Optimizer/Transforms/Passes.td

@@ -47,6 +47,17 @@ def ApplySpecialization : Pass<"apply-op-specialization", "mlir::ModuleOp"> {
   ];
 }
 
+def AssignIDs : Pass<"assign-ids", "mlir::func::FuncOp"> {
+  let summary = "Generate and assign unique identifiers for virtual qubits.";
+  let description = [{
+   Attachs a matching unique ID attribute to each `quake.null_wire` and corresponding
+   `quake.sink`. This is used in DependencyAnalysis to associate allocations and
+   deallocations for tracking qubit lifetimes.
+  }];
+
+  let dependentDialects = ["quake::QuakeDialect"];
+}
+
 def BasisConversionPass: Pass<"basis-conversion", "mlir::ModuleOp"> {
   let summary = "Converts kernels to a set of basis operations.";
   let description = [{
@@ -161,6 +172,47 @@ def DecompositionPass: Pass<"decomposition", "mlir::ModuleOp"> {
   ];
 }
 
+def DependencyAnalysis : Pass<"dep-analysis", "mlir::func::FuncOp"> {
+  let summary = "Maps qubits and reorders operations based on dependency graph.";
+  let description = [{
+    A dependency graph is a Directed Acyclic Graph (DAG) where each node represents
+    an operation, and each edge represents a "depends on" relation between that operation
+    and another operation. For example, in the following snippet the `x` operation
+    depends on the `h` operation because it is applied to the same qubit (`q`), so the
+    `h` operation must happen before the `x` operation:
+
+    ```
+    cudaq::qubit q;
+    x(q);
+    h(q);
+    ```
+
+    Once a dependency graph is created, it is then scheduled, assigning each operation
+    a virtual cycle. Operations that don't depend on each other may be scheduled at the
+    same cycle. However, an operation that depends on a second operation must be scheduled
+    after the second operation. The scheduling algorithm tries to pack operations as
+    densely as possible, minimizing the number of cycles between operations.
+
+    From this dependency graph, we can calculate the lifetime of a qubit: from the cycle
+    in which it is first used through the cycle in which it is last used. If two virtual
+    qubits have non-overlapping lifetimes, they can be assigned to the same physical qubit,
+    as every virtual qubit is assumed to be fully reset before release. Failure to fully
+    reset virtual qubits before release is undefinied behavior, and will likely lead to
+    incorrect output when running DependencyAnalysis.
+  }];
+
+  let dependentDialects = ["quake::QuakeDialect"];
+
+  let statistics = [
+    Statistic<"numVirtualQubits", "num-virtual-qubits",
+              "Number of virtual qubits used">,
+    Statistic<"numPhysicalQubits", "num-physical-qubits",
+              "Number of phyiscal qubits used">,
+    Statistic<"numCycles", "num-cycles",
+              "Length of kernel in cycles">,
+  ];
+}
+
 def EraseNopCalls : Pass<"erase-nop-calls", "mlir::func::FuncOp"> {
   let summary = "Erase calls to any builtin intrinsics that are NOPs.";
   let description = [{

lib/Optimizer/Dialect/Quake/QuakeOps.cpp

@@ -26,15 +26,6 @@ namespace {
 #include "cudaq/Optimizer/Dialect/Quake/Canonical.inc"
 } // namespace
 
-// Is \p op in the Quake dialect?
-// TODO: Is this StringRef comparison faster than calling MLIRContext::
-// getLoadedDialect("quake")?
-static bool isQuakeOperation(Operation *op) {
-  if (auto *dialect = op->getDialect())
-    return dialect->getNamespace().equals("quake");
-  return false;
-}
-
 static LogicalResult verifyWireResultsAreLinear(Operation *op) {
   for (Value v : op->getOpResults())
     if (isa<quake::WireType>(v.getType())) {

lib/Optimizer/Transforms/AssignIDs.cpp

@@ -0,0 +1,174 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#include "cudaq/Frontend/nvqpp/AttributeNames.h"
+#include "cudaq/Optimizer/Dialect/CC/CCDialect.h"
+#include "cudaq/Optimizer/Dialect/Quake/QuakeDialect.h"
+#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
+#include "cudaq/Optimizer/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Threading.h"
+#include "mlir/InitAllDialects.h"
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Generated logic
+//===----------------------------------------------------------------------===//
+namespace cudaq::opt {
+#define GEN_PASS_DEF_ASSIGNIDS
+#include "cudaq/Optimizer/Transforms/Passes.h.inc"
+} // namespace cudaq::opt
+
+namespace {
+bool isMeasureOp(Operation *op) {
+  return dyn_cast<quake::MxOp>(*op) || dyn_cast<quake::MyOp>(*op) ||
+         dyn_cast<quake::MzOp>(*op);
+}
+
+int numClassicalInput(Operation *op) {
+  if (dyn_cast<quake::RxOp>(*op) || dyn_cast<quake::RyOp>(*op) ||
+      dyn_cast<quake::RzOp>(*op))
+    return 1;
+
+  if (dyn_cast<quake::PhasedRxOp>(*op))
+    return 2;
+
+  return 0;
+}
+
+class NullWirePat : public OpRewritePattern<quake::NullWireOp> {
+public:
+  unsigned *counter;
+
+  NullWirePat(MLIRContext *context, unsigned *c)
+      : OpRewritePattern<quake::NullWireOp>(context), counter(c) {}
+
+  LogicalResult matchAndRewrite(quake::NullWireOp alloc,
+                                PatternRewriter &rewriter) const override {
+    if (alloc->hasAttr("qid"))
+      return failure();
+
+    auto qid = (*counter)++;
+
+    rewriter.startRootUpdate(alloc);
+    alloc->setAttr("qid", rewriter.getUI32IntegerAttr(qid));
+    rewriter.finalizeRootUpdate(alloc);
+
+    return success();
+  }
+};
+
+std::optional<uint> findQid(Value v) {
+  auto defop = v.getDefiningOp();
+  if (!defop)
+    return std::nullopt;
+
+  if (defop->getRegions().size() != 0) {
+    defop->emitOpError(
+        "AssignIDsPass cannot handle non-function operations with regions."
+        " Do you have if statements in a Base Profile QIR program?");
+    return std::nullopt;
+  }
+
+  if (!isa<quake::WireType>(v.getType()))
+    return std::nullopt;
+
+  assert(quake::isLinearValueForm(defop) &&
+         "AssignIDsPass requires operations to be in value form");
+
+  if (defop->hasAttr("qid")) {
+    uint qid = defop->getAttr("qid").cast<IntegerAttr>().getUInt();
+    return std::optional<uint>(qid);
+  }
+
+  // Figure out matching operand
+  size_t i = 0;
+  for (; i < defop->getNumResults(); i++)
+    if (defop->getResult(i) == v)
+      break;
+
+  // Special cases where result # != operand #:
+  // Wire is second output but sole input
+  if (isMeasureOp(defop))
+    i = 0;
+  // Classical values preceding wires as input are consumed and not part of the results
+  i += numClassicalInput(defop);
+  // Swap op swaps wires
+  if (dyn_cast<quake::SwapOp>(defop))
+    i = (i == 1 ? 0 : 1);
+
+  return findQid(defop->getOperand(i));
+}
+
+class SinkOpPat : public OpRewritePattern<quake::SinkOp> {
+public:
+  SinkOpPat(MLIRContext *context) : OpRewritePattern<quake::SinkOp>(context) {}
+
+  LogicalResult matchAndRewrite(quake::SinkOp release,
+                                PatternRewriter &rewriter) const override {
+    auto qid = findQid(release.getOperand());
+
+    if (!qid.has_value())
+      return failure();
+
+    rewriter.startRootUpdate(release);
+    release->setAttr("qid", rewriter.getUI32IntegerAttr(qid.value()));
+    rewriter.finalizeRootUpdate(release);
+
+    return success();
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// Pass implementation
+//===----------------------------------------------------------------------===//
+
+struct AssignIDsPass : public cudaq::opt::impl::AssignIDsBase<AssignIDsPass> {
+  using AssignIDsBase::AssignIDsBase;
+
+  void runOnOperation() override {
+    auto func = getOperation();
+
+    if (!func->hasAttr("cudaq-kernel") || func.getBlocks().empty())
+      return;
+
+    if (!func.getFunctionBody().hasOneBlock()) {
+      func.emitError("AssignIDsPass cannot handle multiple blocks. Do "
+                     "you have if statements in a Base Profile QIR program?");
+      signalPassFailure();
+      return;
+    }
+
+    assign();
+  }
+
+  void assign() {
+    auto *ctx = &getContext();
+    func::FuncOp func = getOperation();
+    RewritePatternSet patterns(ctx);
+    unsigned x = 0;
+    patterns.insert<NullWirePat>(ctx, &x);
+    patterns.insert<SinkOpPat>(ctx);
+    ConversionTarget target(*ctx);
+    target.addLegalDialect<quake::QuakeDialect>();
+    target.addDynamicallyLegalOp<quake::NullWireOp>(
+        [&](quake::NullWireOp alloc) { return alloc->hasAttr("qid"); });
+    target.addDynamicallyLegalOp<quake::SinkOp>(
+        [&](quake::SinkOp sink) { return sink->hasAttr("qid"); });
+    if (failed(applyPartialConversion(func.getOperation(), target,
+                                      std::move(patterns)))) {
+      func.emitOpError("Assigning qids failed");
+      signalPassFailure();
+    }
+  }
+};
+
+} // namespace

lib/Optimizer/Transforms/CMakeLists.txt

@@ -15,6 +15,7 @@ add_cudaq_library(OptTransforms
   AggressiveEarlyInlining.cpp
   ApplyControlNegations.cpp
   ApplyOpSpecialization.cpp
+  AssignIDs.cpp
   BasisConversion.cpp
   CombineQuantumAlloc.cpp
   ConstPropComplex.cpp
@@ -39,6 +40,7 @@ add_cudaq_library(OptTransforms
   LowerUnwind.cpp
   Mapping.cpp
   MemToReg.cpp
+  DependencyAnalysis.cpp
   MultiControlDecomposition.cpp
   ObserveAnsatz.cpp
   PruneCtrlRelations.cpp