✨ Enhance QubitType to support static and dynamic qubits, updating related operations and conversions in QCO and QC dialects. Add tests for static qubit type propagation and ensure correct handling in various operations.

Author: simon1hofmann

CHANGELOG.md

@@ -15,7 +15,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel
 - ✨ Add conversions between Jeff and QCO ([#1479], [#1548], [#1565]) ([**@denialhaag**])
 - ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568]) ([**@MatthiasReumann**])
 - ✨ Add initial infrastructure for new QC and QCO MLIR dialects
-  ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], [#1513], [#1521], [#1542], [#1548], [#1550], [#1554], [#1570])
+  ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], [#1513], [#1521], [#1542], [#1548], [#1550], [#1554], [#1569], [#1570])
   ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**])
 
 ### Changed
@@ -335,6 +335,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool
 
 [#1571]: https://github.com/munich-quantum-toolkit/core/pull/1571
 [#1570]: https://github.com/munich-quantum-toolkit/core/pull/1570
+[#1569]: https://github.com/munich-quantum-toolkit/core/pull/1569
 [#1568]: https://github.com/munich-quantum-toolkit/core/pull/1568
 [#1565]: https://github.com/munich-quantum-toolkit/core/pull/1565
 [#1564]: https://github.com/munich-quantum-toolkit/core/pull/1564

mlir/include/mlir/Dialect/QC/IR/QCOps.td

@@ -78,18 +78,24 @@ def AllocOp : QCOp<"alloc", [MemoryEffects<[MemAlloc]>]> {
     let hasVerifier = 1;
 }
 
+def DynamicQubit : Type<
+    And<[CPred<"::mlir::isa<::mlir::qc::QubitType>($_self)">,
+         CPred<"!::mlir::cast<::mlir::qc::QubitType>($_self).getIsStatic()">]>,
+    "dynamic qubit type (!qc.qubit)">;
+
 def DeallocOp : QCOp<"dealloc", [MemoryEffects<[MemFree]>]> {
-    let summary = "Deallocate a qubit";
+    let summary = "Deallocate a dynamically allocated qubit";
     let description = [{
-        Deallocates a qubit, releasing its resources.
+        Deallocates a dynamically allocated qubit, releasing its resources.
+        Static qubits (`!qc.qubit<static>`) cannot be deallocated.
 
         Example:
         ```mlir
         qc.dealloc %q : !qc.qubit
         ```
     }];
 
-    let arguments = (ins QubitType:$qubit);
+    let arguments = (ins DynamicQubit:$qubit);
     let assemblyFormat = "$qubit attr-dict `:` type($qubit)";
 
     let hasCanonicalizer = 1;
@@ -104,13 +110,20 @@ def StaticOp : QCOp<"static", [Pure]> {
 
         Example:
         ```mlir
-        %q = qc.static 0 : !qc.qubit
+        %q = qc.static 0 : !qc.qubit<static>
         ```
     }];
 
     let arguments = (ins ConfinedAttr<I64Attr, [IntNonNegative]>:$index);
     let results = (outs QubitType:$qubit);
     let assemblyFormat = "$index attr-dict `:` type($qubit)";
+
+    let builders = [
+        OpBuilder<(ins "int64_t":$index), [{
+            build($_builder, $_state, QubitType::get($_builder.getContext(), /*isStatic=*/true),
+                  $_builder.getI64IntegerAttr(index));
+        }]>
+    ];
 }
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/QCO/IR/QCOOps.td

@@ -56,9 +56,9 @@ class QCOToQCTypeConverter final : public TypeConverter {
     // Identity conversion for all types by default
     addConversion([](Type type) { return type; });
 
-    // Convert QCO qubit values to QC qubit references
-    addConversion([ctx](qco::QubitType /*type*/) -> Type {
-      return qc::QubitType::get(ctx);
+    // Convert QCO qubit values to QC qubit references, preserving isStatic
+    addConversion([ctx](qco::QubitType type) -> Type {
+      return qc::QubitType::get(ctx, type.getIsStatic());
     });
   }
 };
@@ -97,36 +97,25 @@ struct ConvertQCOAllocOp final : OpConversionPattern<qco::AllocOp> {
 };
 
 /**
- * @brief Converts qco.dealloc to qc.dealloc (except for static qubits).
+ * @brief Converts qco.dealloc to qc.dealloc (dynamic) or erases it (static).
  *
  * @details
- * Deallocates a qubit, releasing its resources. The OpAdaptor automatically
- * provides the type-converted qubit operand (!qc.qubit instead of
- * !qco.qubit), so we simply pass it through to the new operation.
- *
- * Example transformation:
- * ```mlir
- * qco.dealloc %q_qco : !qco.qubit
- * // becomes:
- * qc.dealloc %q_qc : !qc.qubit
- * ```
- *
- * For static qubits, we erase `qco.dealloc` because QC does
- * not require explicit sinks for static qubit handles.
+ * For dynamic qubits (`!qco.qubit`), lowers to `qc.dealloc`.
+ * For static qubits (`!qco.qubit<static>`), erases the op since QC does not
+ * require explicit deallocation of static qubits.
  */
 struct ConvertQCODeallocOp final : OpConversionPattern<qco::DeallocOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(qco::DeallocOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter& rewriter) const override {
-    auto qcQubit = adaptor.getQubit();
-    if (llvm::isa_and_nonnull<qc::StaticOp>(qcQubit.getDefiningOp())) {
+    if (op.getQubit().getType().getIsStatic()) {
       rewriter.eraseOp(op);
       return success();
     }
 
-    rewriter.replaceOpWithNewOp<qc::DeallocOp>(op, qcQubit);
+    rewriter.replaceOpWithNewOp<qc::DeallocOp>(op, adaptor.getQubit());
     return success();
   }
 };

mlir/lib/Conversion/QCOToQC/QCOToQC.cpp

@@ -129,42 +129,55 @@ struct ConvertFuncReturnOp final : StatefulOpConversionPattern<func::ReturnOp> {
                   ConversionPatternRewriter& rewriter) const override {
     auto& state = getState();
 
-    // Only insert sinks at function exit (never inside nested modifier
-    // regions).
-    if (state.inNestedRegion == 0) {
-      llvm::SmallVector<Value> liveQubits;
-      liveQubits.reserve(state.qubitMap.size());
+    if (state.inNestedRegion != 0) {
+      rewriter.replaceOpWithNewOp<func::ReturnOp>(op, adaptor.getOperands());
+      return success();
+    }
 
-      llvm::DenseSet<Value> escapedQubits;
-      for (Value returned : adaptor.getOperands()) {
-        escapedQubits.insert(returned);
+    // Build return values from qubitMap (adaptor.getOperands() may carry stale
+    // root values because gate patterns use eraseOp instead of replaceOp).
+    llvm::SmallVector<Value> returnValues;
+    llvm::DenseSet<Value> escapedQubits;
+    returnValues.reserve(op.getNumOperands());
+    for (auto [qcOperand, adaptorOperand] :
+         llvm::zip(op.getOperands(), adaptor.getOperands())) {
+      if (state.qubitMap.contains(qcOperand)) {
+        auto latest = state.qubitMap[qcOperand];
+        returnValues.push_back(latest);
+        escapedQubits.insert(latest);
+      } else {
+        returnValues.push_back(adaptorOperand);
       }
+    }
 
-      llvm::DenseSet<Value> seen;
-      for (const auto& [/*qcQubit*/ _, qcoQubit] : state.qubitMap) {
-        if (!escapedQubits.contains(qcoQubit) && seen.insert(qcoQubit).second) {
-          liveQubits.push_back(qcoQubit);
-        }
+    // Collect non-escaped live qubits for deallocation.
+    llvm::SmallVector<Value> liveQubits;
+    liveQubits.reserve(state.qubitMap.size());
+    llvm::DenseSet<Value> seen;
+    for (const auto& [qcQubit, qcoQubit] : state.qubitMap) {
+      if (escapedQubits.contains(qcoQubit) || !seen.insert(qcoQubit).second) {
+        continue;
       }
+      liveQubits.emplace_back(qcoQubit);
+    }
 
-      // Sort deterministically (mirrors QCOProgramBuilder::finalize()).
-      llvm::sort(liveQubits, [](Value a, Value b) {
-        auto* opA = a.getDefiningOp();
-        auto* opB = b.getDefiningOp();
-        if (!opA || !opB || opA->getBlock() != opB->getBlock()) {
-          return a.getAsOpaquePointer() < b.getAsOpaquePointer();
-        }
-        return opA->isBeforeInBlock(opB);
-      });
-
-      for (auto qubit : liveQubits) {
-        rewriter.create<qco::DeallocOp>(op.getLoc(), qubit);
+    // Sort deterministically (mirrors QCOProgramBuilder::finalize()).
+    llvm::sort(liveQubits, [](Value a, Value b) {
+      auto* opA = a.getDefiningOp();
+      auto* opB = b.getDefiningOp();
+      if (!opA || !opB || opA->getBlock() != opB->getBlock()) {
+        return a.getAsOpaquePointer() < b.getAsOpaquePointer();
       }
+      return opA->isBeforeInBlock(opB);
+    });
 
-      state.qubitMap.clear();
+    for (Value qubit : liveQubits) {
+      rewriter.create<qco::DeallocOp>(op.getLoc(), qubit);
     }
 
-    rewriter.replaceOpWithNewOp<func::ReturnOp>(op, adaptor.getOperands());
+    state.qubitMap.clear();
+
+    rewriter.replaceOpWithNewOp<func::ReturnOp>(op, returnValues);
     return success();
   }
 };
@@ -186,9 +199,9 @@ class QCToQCOTypeConverter final : public TypeConverter {
     // Identity conversion for all types by default
     addConversion([](Type type) { return type; });
 
-    // Convert QC qubit references to QCO qubit values
-    addConversion([ctx](qc::QubitType /*type*/) -> Type {
-      return qco::QubitType::get(ctx);
+    // Convert QC qubit references to QCO qubit values, preserving isStatic
+    addConversion([ctx](qc::QubitType type) -> Type {
+      return qco::QubitType::get(ctx, type.getIsStatic());
     });
   }
 };
@@ -1008,11 +1021,11 @@ struct ConvertQCCtrlOp final : StatefulOpConversionPattern<qc::CtrlOp> {
            "QC ctrl region unexpectedly has entry block arguments");
     SmallVector<Value> qcoTargetAliases;
     qcoTargetAliases.reserve(numTargets);
-    const auto qubitType = qco::QubitType::get(qcoOp.getContext());
     const auto opLoc = op.getLoc();
     rewriter.modifyOpInPlace(qcoOp, [&] {
       for (auto i = 0UL; i < numTargets; i++) {
-        qcoTargetAliases.emplace_back(entryBlock.addArgument(qubitType, opLoc));
+        qcoTargetAliases.emplace_back(
+            entryBlock.addArgument(qcoTargets[i].getType(), opLoc));
       }
     });
     state.targetsIn[state.inNestedRegion] = std::move(qcoTargetAliases);
@@ -1092,11 +1105,11 @@ struct ConvertQCInvOp final : StatefulOpConversionPattern<qc::InvOp> {
            "QC inv region unexpectedly has entry block arguments");
     SmallVector<Value> qcoTargetAliases;
     qcoTargetAliases.reserve(numTargets);
-    const auto qubitType = qco::QubitType::get(qcoOp.getContext());
     const auto opLoc = op.getLoc();
     rewriter.modifyOpInPlace(qcoOp, [&] {
       for (auto i = 0UL; i < numTargets; i++) {
-        qcoTargetAliases.emplace_back(entryBlock.addArgument(qubitType, opLoc));
+        qcoTargetAliases.emplace_back(
+            entryBlock.addArgument(qcoTargets[i].getType(), opLoc));
       }
     });
     targetsIn[inNestedRegion] = std::move(qcoTargetAliases);
@@ -1218,9 +1231,9 @@ struct QCToQCO final : impl::QCToQCOBase<QCToQCO> {
     // Conversion of qc types in func.return
     //
     // Note: `func.return` may already be type-legal even though we still need
-    // to insert `qco.dealloc` sinks for remaining live qubits. Therefore, we
-    // make it dynamically illegal unless the lowering state has no remaining
-    // qubits.
+    // to insert sink operations (`qco.dealloc`) for remaining live
+    // qubits. Therefore, we make it dynamically illegal unless the lowering
+    // state has no remaining qubits.
     patterns.add<ConvertFuncReturnOp>(typeConverter, context, &state);
     populateReturnOpTypeConversionPattern(patterns, typeConverter);
     target.addDynamicallyLegalOp<func::ReturnOp>([&](const func::ReturnOp op) {

mlir/lib/Conversion/QCToQCO/QCToQCO.cpp

@@ -86,9 +86,7 @@ Value QCProgramBuilder::staticQubit(const int64_t index) {
     llvm::reportFatalUsageError("Index must be non-negative");
   }
 
-  // Create the StaticOp with the given index
-  auto indexAttr = getI64IntegerAttr(index);
-  auto staticOp = StaticOp::create(*this, indexAttr);
+  auto staticOp = StaticOp::create(*this, index);
   return staticOp.getQubit();
 }
 

mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp

@@ -46,6 +46,24 @@ void QCDialect::initialize() {
 // Types
 //===----------------------------------------------------------------------===//
 
+/// Print `!qc.qubit` (dynamic, default) or `!qc.qubit<static>`.
+void QubitType::print(AsmPrinter& printer) const {
+  if (getIsStatic()) {
+    printer << "<static>";
+  }
+}
+
+/// Parse `!qc.qubit` or `!qc.qubit<static>`.
+Type QubitType::parse(AsmParser& parser) {
+  if (succeeded(parser.parseOptionalLess())) {
+    if (parser.parseKeyword("static") || parser.parseGreater()) {
+      return {};
+    }
+    return get(parser.getContext(), /*isStatic=*/true);
+  }
+  return get(parser.getContext(), /*isStatic=*/false);
+}
+
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/QC/IR/QCOpsTypes.cpp.inc"
 

mlir/lib/Dialect/QC/IR/QCOps.cpp

@@ -90,8 +90,7 @@ Value QCOProgramBuilder::staticQubit(const int64_t index) {
     llvm::reportFatalUsageError("Index must be non-negative");
   }
 
-  auto indexAttr = getI64IntegerAttr(index);
-  auto staticOp = StaticOp::create(*this, indexAttr);
+  auto staticOp = StaticOp::create(*this, index);
   const auto qubit = staticOp.getQubit();
 
   // Track the static qubit as valid
@@ -747,9 +746,8 @@ std::pair<ValueRange, ValueRange> QCOProgramBuilder::ctrl(
 
   auto ctrlOp = CtrlOp::create(*this, controls, targets);
   auto& block = ctrlOp.getBodyRegion().emplaceBlock();
-  const auto qubitType = QubitType::get(getContext());
   for (const auto target : targets) {
-    const auto arg = block.addArgument(qubitType, getLoc());
+    const auto arg = block.addArgument(target.getType(), getLoc());
     updateQubitTracking(target, arg);
   }
   const InsertionGuard guard(*this);
@@ -785,9 +783,8 @@ ValueRange QCOProgramBuilder::inv(
 
   // Add block arguments for all qubits
   auto& block = invOp.getBodyRegion().emplaceBlock();
-  const auto qubitType = QubitType::get(getContext());
   for (const auto qubit : qubits) {
-    const auto arg = block.addArgument(qubitType, getLoc());
+    const auto arg = block.addArgument(qubit.getType(), getLoc());
     updateQubitTracking(qubit, arg);
   }
 

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

@@ -242,9 +242,8 @@ void CtrlOp::build(
   build(odsBuilder, odsState, controls, targets);
   auto& block = odsState.regions.front()->emplaceBlock();
 
-  const auto qubitType = QubitType::get(odsBuilder.getContext());
   for (size_t i = 0; i < targets.size(); ++i) {
-    block.addArgument(qubitType, odsState.location);
+    block.addArgument(targets[i].getType(), odsState.location);
   }
 
   const OpBuilder::InsertionGuard guard(odsBuilder);
@@ -263,9 +262,8 @@ LogicalResult CtrlOp::verify() {
     return emitOpError(
         "number of block arguments must match the number of targets");
   }
-  const auto qubitType = QubitType::get(getContext());
   for (size_t i = 0; i < numTargets; ++i) {
-    if (block.getArgument(i).getType() != qubitType) {
+    if (block.getArgument(i).getType() != getTargetsIn()[i].getType()) {
       return emitOpError("block argument type at index ")
              << i << " does not match target type";
     }

mlir/lib/Dialect/QCO/IR/Modifiers/CtrlOp.cpp

@@ -326,9 +326,8 @@ void InvOp::build(
   build(odsBuilder, odsState, qubits);
   auto& block = odsState.regions.front()->emplaceBlock();
 
-  const auto qubitType = QubitType::get(odsBuilder.getContext());
   for (size_t i = 0; i < qubits.size(); ++i) {
-    block.addArgument(qubitType, odsState.location);
+    block.addArgument(qubits[i].getType(), odsState.location);
   }
 
   const OpBuilder::InsertionGuard guard(odsBuilder);
@@ -347,9 +346,8 @@ LogicalResult InvOp::verify() {
     return emitOpError(
         "number of block arguments must match the number of targets");
   }
-  const auto qubitType = QubitType::get(getContext());
   for (size_t i = 0; i < numTargets; ++i) {
-    if (block.getArgument(i).getType() != qubitType) {
+    if (block.getArgument(i).getType() != getQubitsIn()[i].getType()) {
       return emitOpError("block argument type at index ")
              << i << " does not match target type";
     }