✨ Add Quantum IPO Pass to mqt-cc

include/mqt-core/qasm3/Types.hpp

@@ -154,7 +154,7 @@ template <typename T> class DesignatedType final : public Type<T> {
   std::string designatorToString();
 };
 
-enum UnsizedTy : uint8_t { Bool, Duration };
+enum UnsizedTy : uint8_t { Bool, Duration, SingleQubit };
 
 template <typename T> class UnsizedType final : public Type<T> {
 public:
@@ -178,6 +178,9 @@ template <typename T> class UnsizedType final : public Type<T> {
   static std::shared_ptr<Type<T>> getDurationTy() {
     return std::make_shared<UnsizedType>(Duration);
   }
+  static std::shared_ptr<Type<T>> getSingleQubitTy() {
+    return std::make_shared<UnsizedType>(SingleQubit);
+  }
 
   T getDesignator() override {
     throw std::runtime_error("Unsized types do not have designators");
@@ -195,6 +198,8 @@ template <typename T> class UnsizedType final : public Type<T> {
       return "bool";
     case Duration:
       return "duration";
+    case SingleQubit:
+      return "qubit";
     }
     throw std::runtime_error("Unhandled type");
   }

mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h

@@ -10,12 +10,14 @@
 
 #pragma once
 
+#include <llvm/ADT/DenseSet.h>
 #include <mlir/IR/Builders.h>
 #include <mlir/IR/OwningOpRef.h>
 #include <mlir/IR/Value.h>
 #include <mlir/Support/LLVM.h>
 
 #include <cstdint>
+#include <optional>
 #include <string>
 #include <utility>
 #include <variant>
@@ -1755,6 +1757,109 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder {
                                   const std::variant<int64_t, Value>& index,
                                   ValueRange yieldedValues);
 
+  //===--------------------------------------------------------------------===//
+  // Additional functions
+  //===--------------------------------------------------------------------===//
+
+  /**
+   * @brief Start building an additional private function in the module
+   *
+   * @details
+   * Creates a new private `func.func` at the end of the module and moves the
+   * insertion point into its entry block. Qubit- and qubit-tensor-typed
+   * arguments are added to the linear-type tracking so that operations can be
+   * applied to them; a tensor argument is treated as a register the callee
+   * owns for the duration of the call. Arguments of other types (e.g., `f64`)
+   * are returned as-is without tracking. Must be called after `initialize()`
+   * and must be paired with a call to `endFunction()`; function definitions
+   * cannot be nested.
+   *
+   * @param name The name of the function
+   * @param argTypes The argument types of the function
+   * @param resultTypes The result types of the function
+   * @return The entry block arguments of the new function
+   *
+   * @par Example:
+   * ```c++
+   * auto args = builder.startFunction(
+   *     "f", {builder.getQubitType()}, {builder.getQubitType()});
+   * auto q = builder.h(args[0]);
+   * builder.endFunction({q});
+   * ```
+   * ```mlir
+   * func.func private @f(%arg0: !qco.qubit) -> !qco.qubit {
+   *   %q = qco.h %arg0 : !qco.qubit -> !qco.qubit
+   *   return %q : !qco.qubit
+   * }
+   * ```
+   */
+  SmallVector<Value> startFunction(StringRef name, TypeRange argTypes,
+                                   TypeRange resultTypes);
+
+  /**
+   * @brief Finish the function started with `startFunction()`
+   *
+   * @details
+   * Creates the `func.return` with the given values and restores the insertion
+   * point to where it was before `startFunction()` was called. All qubits and
+   * tensors that were created within the function (from arguments or
+   * operations) must either be consumed (e.g., by `sink()` or
+   * `qtensorDealloc()`) or returned; otherwise a usage error is reported.
+   *
+   * @param returnValues The values to return from the function
+   */
+  void endFunction(ValueRange returnValues);
+
+  /**
+   * @brief Call a function previously defined in the module
+   *
+   * @details
+   * Creates a `func.call` to the named function. Qubit- and qubit-tensor-typed
+   * operands are validated and consumed; results of those types are added to
+   * the tracking.
+   *
+   * The i-th qubit operand is paired with the i-th qubit result, and likewise
+   * for tensors, so a function that threads its linear values through keeps
+   * the register association intact. Surplus operands are treated as consumed
+   * and surplus results as freshly created. This positional pairing is a
+   * calling convention, not something the IR enforces: a callee that returns
+   * its qubits in a different order than it takes them will be tracked
+   * incorrectly.
+   *
+   * @param callee The name of the function to call
+   * @param operands The operands to pass to the call
+   * @return The results of the call operation
+   *
+   * @par Example:
+   * ```c++
+   * auto results = builder.call("f", {q0});
+   * ```
+   * ```mlir
+   * %q1 = call @f(%q0) : (!qco.qubit) -> !qco.qubit
+   * ```
+   */
+  SmallVector<Value> call(StringRef callee, ValueRange operands);
+
+  /**
+   * @brief Get the qubit type
+   * @return The `!qco.qubit` type
+   */
+  Type getQubitType();
+
+  /**
+   * @brief Get a one-dimensional tensor type holding qubits
+   * @param size The number of qubits in the tensor
+   * @return The `tensor<size x !qco.qubit>` type
+   */
+  Type getQubitTensorType(int64_t size);
+
+  /**
+   * @brief Check whether the given type is a tensor of qubits
+   * @param type The type to check
+   * @return True if @p type is a ranked tensor with `!qco.qubit` elements
+   */
+  static bool isQubitTensor(Type type);
+
   //===--------------------------------------------------------------------===//
   // Finalization
   //===--------------------------------------------------------------------===//
@@ -1943,6 +2048,21 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder {
 
   /// Ensure static and dynamic qubit allocation modes are not mixed.
   void ensureAllocationMode(AllocationMode requestedMode);
+
+  /**
+   * @brief State of an additional function under construction
+   */
+  struct FunctionScope {
+    /// Insertion point to restore when the function is finished
+    OpBuilder::InsertPoint savedInsertPoint;
+    /// Qubit values that were already tracked before the function was started
+    llvm::DenseSet<Value> outerQubits;
+    /// Tensor values that were already tracked before the function was started
+    llvm::DenseSet<Value> outerTensors;
+  };
+
+  /// Active function scope, if a function is currently under construction.
+  std::optional<FunctionScope> functionScope;
 };
 } // namespace qco
 } // namespace mlir

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

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <mlir/IR/BuiltinOps.h>
 #include <mlir/Interfaces/FunctionInterfaces.h>
 #include <mlir/Pass/Pass.h>
 #include <mlir/Pass/PassRegistry.h>
@@ -33,6 +34,10 @@ namespace mlir::qco {
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/QCO/Transforms/Passes.h.inc" // IWYU pragma: export
 
+void runQuantumArgumentPromotion(ModuleOp module);
+void runAuxiliaryQubitHoisting(ModuleOp module);
+void runQuantumFunctionBoundaryCommutation(ModuleOp module,
+                                           SymbolTable& symbolTable);
 /**
  * @brief Create target-independent two-qubit gate fusion.
  */

mlir/include/mlir/Dialect/QCO/Transforms/Passes.td

@@ -305,6 +305,18 @@ def ReuseQubits : Pass<"reuse-qubits", "mlir::ModuleOp"> {
   }];
 }
 
+def QuantumIPO : Pass<"quantum-ipo", "mlir::ModuleOp"> {
+  let summary = "Quantum Interprocedural Optimization Pass";
+  let description = [{
+        Performs interprocedural optimizations on quantum functions.
+    }];
+
+  let dependentDialects = ["::mlir::func::FuncDialect",
+                           "::mlir::arith::ArithDialect",
+                           "::mlir::qtensor::QTensorDialect",
+                           "mlir::qco::QCODialect"];
+}
+
 //===----------------------------------------------------------------------===//
 // Decomposition Passes
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp

@@ -33,10 +33,12 @@
 #include <mlir/IR/Location.h>
 #include <mlir/IR/MLIRContext.h>
 #include <mlir/IR/OwningOpRef.h>
+#include <mlir/IR/SymbolTable.h>
 #include <mlir/IR/Value.h>
 #include <mlir/IR/ValueRange.h>
 #include <mlir/Support/LLVM.h>
 
+#include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
@@ -1374,6 +1376,171 @@ QCOProgramBuilder::scfCondition(Value reg,
   return scfCondition(condition, yieldedValues);
 }
 
+//===----------------------------------------------------------------------===//
+// Additional Functions
+//===----------------------------------------------------------------------===//
+
+Type QCOProgramBuilder::getQubitType() { return QubitType::get(ctx); }
+
+Type QCOProgramBuilder::getQubitTensorType(const int64_t size) {
+  return RankedTensorType::get({size}, getQubitType());
+}
+
+bool QCOProgramBuilder::isQubitTensor(Type type) {
+  const auto tensorType = dyn_cast<RankedTensorType>(type);
+  return tensorType && isa<QubitType>(tensorType.getElementType());
+}
+
+SmallVector<Value> QCOProgramBuilder::startFunction(StringRef name,
+                                                    TypeRange argTypes,
+                                                    TypeRange resultTypes) {
+  checkFinalized();
+
+  if (functionScope.has_value()) {
+    llvm::reportFatalUsageError(
+        "Cannot start a function while another one is being built");
+  }
+
+  FunctionScope scope{.savedInsertPoint = saveInsertionPoint(),
+                      .outerQubits = {},
+                      .outerTensors = {}};
+  for (const auto& [qubit, info] : validQubits) {
+    scope.outerQubits.insert(qubit);
+  }
+  for (const auto& [tensor, info] : validTensors) {
+    scope.outerTensors.insert(tensor);
+  }
+
+  setInsertionPointToEnd(cast<ModuleOp>(module).getBody());
+  auto funcOp =
+      func::FuncOp::create(*this, name, getFunctionType(argTypes, resultTypes));
+  // The interprocedural passes only consider functions that are not externally
+  // visible, so additional functions are private by default.
+  funcOp.setPrivate();
+
+  auto& entryBlock = funcOp.getBody().emplaceBlock();
+  const SmallVector<Location> locs(argTypes.size(), getLoc());
+  entryBlock.addArguments(argTypes, locs);
+  setInsertionPointToStart(&entryBlock);
+
+  SmallVector<Value> args;
+  for (const auto arg : entryBlock.getArguments()) {
+    if (isa<QubitType>(arg.getType())) {
+      validQubits.try_emplace(arg, QubitInfo{});
+    } else if (isQubitTensor(arg.getType())) {
+      // A tensor argument acts like a register the callee owns for the
+      // duration of the call, so give it its own register id.
+      validTensors.try_emplace(arg, TensorInfo{tensorCounter++});
+    }
+    args.emplace_back(arg);
+  }
+
+  functionScope = std::move(scope);
+  return args;
+}
+
+void QCOProgramBuilder::endFunction(ValueRange returnValues) {
+  checkFinalized();
+
+  if (!functionScope.has_value()) {
+    llvm::reportFatalUsageError(
+        "endFunction() called without a matching startFunction()");
+  }
+
+  for (const auto value : returnValues) {
+    if (isa<QubitType>(value.getType())) {
+      validateQubitValue(value);
+      validQubits.erase(value);
+    } else if (isQubitTensor(value.getType())) {
+      validateTensorValue(value);
+      validTensors.erase(value);
+    }
+  }
+
+  for (const auto& [qubit, info] : validQubits) {
+    if (!functionScope->outerQubits.contains(qubit)) {
+      llvm::reportFatalUsageError(
+          "Function body has qubit values that are neither returned nor "
+          "consumed");
+    }
+  }
+  for (const auto& [tensor, info] : validTensors) {
+    if (!functionScope->outerTensors.contains(tensor)) {
+      llvm::reportFatalUsageError(
+          "Function body has tensor values that are neither returned nor "
+          "deallocated");
+    }
+  }
+
+  func::ReturnOp::create(*this, returnValues);
+
+  restoreInsertionPoint(functionScope->savedInsertPoint);
+  functionScope.reset();
+}
+
+SmallVector<Value> QCOProgramBuilder::call(StringRef callee,
+                                           ValueRange operands) {
+  checkFinalized();
+
+  auto funcOp = dyn_cast_or_null<func::FuncOp>(
+      SymbolTable::lookupSymbolIn(module, getStringAttr(callee)));
+  if (!funcOp) {
+    llvm::reportFatalUsageError("Callee not found in module");
+  }
+
+  SmallVector<Value> qubitOperands;
+  SmallVector<Value> tensorOperands;
+  for (const auto operand : operands) {
+    if (isa<QubitType>(operand.getType())) {
+      validateQubitValue(operand);
+      qubitOperands.emplace_back(operand);
+    } else if (isQubitTensor(operand.getType())) {
+      validateTensorValue(operand);
+      tensorOperands.emplace_back(operand);
+    }
+  }
+
+  auto callOp = func::CallOp::create(*this, funcOp, operands);
+
+  SmallVector<Value> qubitResults;
+  SmallVector<Value> tensorResults;
+  for (const auto result : callOp.getResults()) {
+    if (isa<QubitType>(result.getType())) {
+      qubitResults.emplace_back(result);
+    } else if (isQubitTensor(result.getType())) {
+      tensorResults.emplace_back(result);
+    }
+  }
+
+  // Thread the i-th linear operand into the i-th linear result of the same
+  // kind. Any operand without a matching result is consumed by the call, any
+  // result without a matching operand is newly created by it.
+  const auto pairedQubits = std::min(qubitOperands.size(), qubitResults.size());
+  for (size_t i = 0; i < pairedQubits; ++i) {
+    updateQubitTracking(qubitOperands[i], qubitResults[i]);
+  }
+  for (size_t i = pairedQubits; i < qubitOperands.size(); ++i) {
+    validQubits.erase(qubitOperands[i]);
+  }
+  for (size_t i = pairedQubits; i < qubitResults.size(); ++i) {
+    validQubits.try_emplace(qubitResults[i], QubitInfo{});
+  }
+
+  const auto pairedTensors =
+      std::min(tensorOperands.size(), tensorResults.size());
+  for (size_t i = 0; i < pairedTensors; ++i) {
+    updateTensorTracking(tensorOperands[i], tensorResults[i]);
+  }
+  for (size_t i = pairedTensors; i < tensorOperands.size(); ++i) {
+    validTensors.erase(tensorOperands[i]);
+  }
+  for (size_t i = pairedTensors; i < tensorResults.size(); ++i) {
+    validTensors.try_emplace(tensorResults[i], TensorInfo{tensorCounter++});
+  }
+
+  return SmallVector<Value>(callOp.getResults());
+}
+
 //===----------------------------------------------------------------------===//
 // Finalization
 //===----------------------------------------------------------------------===//