🎨 Retain register name

CHANGELOG.md

@@ -96,9 +96,9 @@ releases may include breaking changes.
   [#1728], [#1730], [#1749], [#1751], [#1762], [#1765], [#1780], [#1781],
   [#1782], [#1806], [#1807], [#1815], [#1808], [#1824], [#1869], [#1872],
   [#1886], [#1914], [#1925], [#1927], [#1935], [#1936], [#1938], [#1975],
-  [#1976], [#2006]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**],
-  [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**],
-  [**@simon1hofmann**], [**@J4MMlE**])
+  [#1976], [#2006], [#2014]) ([**@burgholzer**], [**@denialhaag**],
+  [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**],
+  [**@MatthiasReumann**], [**@simon1hofmann**], [**@J4MMlE**])
 
 ### Changed
 
@@ -721,6 +721,7 @@ for previous changelogs._
 
 <!-- PR links -->
 
+[#2014]: https://github.com/munich-quantum-toolkit/core/pull/2014
 [#2011]: https://github.com/munich-quantum-toolkit/core/pull/2011
 [#2007]: https://github.com/munich-quantum-toolkit/core/pull/2007
 [#2006]: https://github.com/munich-quantum-toolkit/core/pull/2006

mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <llvm/ADT/StringSet.h>
 #include <mlir/IR/Builders.h>
 #include <mlir/IR/OwningOpRef.h>
 #include <mlir/IR/Value.h>
@@ -189,6 +190,7 @@ class QCProgramBuilder final : public ImplicitLocOpBuilder {
   /**
    * @brief Allocate a qubit register and eagerly load every element
    * @param size Number of qubits (must be positive)
+   * @param name Optional source-level register name
    * @return A `QubitRegister` containing the backing memref and one reference
    * for every eagerly loaded element
    *
@@ -203,18 +205,19 @@ class QCProgramBuilder final : public ImplicitLocOpBuilder {
    * %q2 = memref.load %memref[%c2] : memref<3x!qc.qubit>
    * ```
    */
-  QubitRegister allocQubitRegister(int64_t size);
+  QubitRegister allocQubitRegister(int64_t size, StringRef name = {});
 
   /**
    * @brief Allocate storage for a qubit register without loading its elements
    * @param size Number of qubits (must be positive)
+   * @param name Optional source-level register name
    * @return The memref value representing the qubit register
    *
    * @details The register is tracked for automatic deallocation and remains
    * intact until an element is loaded. Use `loadQubit` to obtain references at
    * their points of use.
    */
-  Value allocQubitRegisterStorage(int64_t size);
+  Value allocQubitRegisterStorage(int64_t size, StringRef name = {});
 
   /**
    * @brief Explicitly loads a qubit from a memref
@@ -1378,6 +1381,9 @@ class QCProgramBuilder final : public ImplicitLocOpBuilder {
   /// Track allocated memrefs for automatic deallocation
   DenseSet<Value> allocatedQregs;
 
+  /// Track non-empty source-level quantum register names.
+  llvm::StringSet<> quantumRegisterNames;
+
   /// Check if the builder has been finalized
   void checkFinalized() const;
 

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

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <llvm/ADT/StringSet.h>
 #include <mlir/IR/Builders.h>
 #include <mlir/IR/OwningOpRef.h>
 #include <mlir/IR/Value.h>
@@ -212,6 +213,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder {
   /**
    * @brief Allocate a qubit tensor and eagerly extract every element
    * @param size Number of qubits (must be positive)
+   * @param name Optional source-level register name
    * @return A `QubitRegister` containing the residual tensor and one standalone
    * qubit value for every eagerly extracted element
    *
@@ -226,7 +228,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder {
    * %t3, %q2 = qtensor.extract %t2[%c2]: tensor<3x!qco.qubit>
    * ```
    */
-  QubitRegister allocQubitRegister(int64_t size);
+  QubitRegister allocQubitRegister(int64_t size, StringRef name = {});
 
   /**
    * @brief Allocate a classical bit register
@@ -1823,6 +1825,9 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder {
   MLIRContext* ctx{};
   Operation* module;
 
+  /// Track non-empty source-level quantum register names.
+  llvm::StringSet<> quantumRegisterNames;
+
   /// Check if the builder has been finalized
   void checkFinalized() const;
 

mlir/include/mlir/Dialect/Utils/Utils.h

@@ -45,6 +45,10 @@ namespace mlir::utils {
 inline constexpr llvm::StringLiteral CLASSICAL_REGISTER_NAME_ATTR =
     "mqt.classical_register_name";
 
+/// Attribute used to retain a source-level quantum-register name.
+inline constexpr llvm::StringLiteral QUANTUM_REGISTER_NAME_ATTR =
+    "mqt.quantum_register_name";
+
 /// Check if a floating-point value is an integer.
 [[nodiscard]] inline bool isIntegerExponent(double r) {
   return r == std::floor(r) && std::isfinite(r);

mlir/lib/Conversion/QCOToQC/QCOToQC.cpp

@@ -18,6 +18,7 @@
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
 #include "mlir/Dialect/QTensor/IR/QTensorDialect.h"
 #include "mlir/Dialect/QTensor/IR/QTensorOps.h"
+#include "mlir/Dialect/Utils/Utils.h"
 
 #include <llvm/ADT/STLExtras.h>
 #include <llvm/ADT/TypeSwitch.h>
@@ -279,14 +280,20 @@ struct ConvertQTensorAllocOp final
     auto tensorType = cast<RankedTensorType>(op.getResult().getType());
     auto memrefType = MemRefType::get(tensorType.getShape(), qubitType);
 
+    const auto registerName = op->getAttr(utils::QUANTUM_REGISTER_NAME_ATTR);
+    memref::AllocOp alloc;
     if (tensorType.hasStaticShape()) {
       // Static size: no dynamic size operand needed
-      rewriter.replaceOpWithNewOp<memref::AllocOp>(op, memrefType);
+      alloc = memref::AllocOp::create(rewriter, op.getLoc(), memrefType);
     } else {
       // Dynamic size: forward the runtime size operand
-      rewriter.replaceOpWithNewOp<memref::AllocOp>(op, memrefType,
-                                                   op.getSize());
+      alloc = memref::AllocOp::create(rewriter, op.getLoc(), memrefType,
+                                      op.getSize());
     }
+    if (registerName) {
+      alloc->setAttr(utils::QUANTUM_REGISTER_NAME_ATTR, registerName);
+    }
+    rewriter.replaceOp(op, alloc.getResult());
     return success();
   }
 };

mlir/lib/Conversion/QCToQCO/QCToQCO.cpp

@@ -19,6 +19,7 @@
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
 #include "mlir/Dialect/QTensor/IR/QTensorDialect.h"
 #include "mlir/Dialect/QTensor/IR/QTensorOps.h"
+#include "mlir/Dialect/Utils/Utils.h"
 
 #include <llvm/ADT/DenseSet.h>
 #include <llvm/ADT/STLExtras.h>
@@ -838,21 +839,25 @@ struct ConvertMemRefAllocOp final
       return failure();
     }
 
-    Value qtensor;
+    const auto registerName = op->getAttr(utils::QUANTUM_REGISTER_NAME_ATTR);
+    qtensor::AllocOp alloc;
     if (shape[0] == ShapedType::kDynamic) {
-      qtensor = rewriter.replaceOpWithNewOp<qtensor::AllocOp>(
-          op, adaptor.getDynamicSizes()[0]);
+      alloc = qtensor::AllocOp::create(rewriter, op.getLoc(),
+                                       adaptor.getDynamicSizes()[0]);
     } else {
       auto size =
           arith::ConstantIndexOp::create(rewriter, op.getLoc(), shape[0]);
-      qtensor =
-          rewriter.replaceOpWithNewOp<qtensor::AllocOp>(op, size.getResult());
+      alloc = qtensor::AllocOp::create(rewriter, op.getLoc(), size.getResult());
+    }
+    if (registerName) {
+      alloc->setAttr(utils::QUANTUM_REGISTER_NAME_ATTR, registerName);
     }
 
     auto& state = getState();
     auto memref = op.getResult();
-    assignMappedTensor(state, qtensor.getDefiningOp(),
-                       lookupRegisterId(state, memref), qtensor);
+    assignMappedTensor(state, alloc, lookupRegisterId(state, memref),
+                       alloc.getResult());
+    rewriter.replaceOp(op, alloc.getResult());
 
     return success();
   }

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

@@ -117,8 +117,8 @@ Value QCProgramBuilder::staticQubit(const uint64_t index) {
 }
 
 QCProgramBuilder::QubitRegister
-QCProgramBuilder::allocQubitRegister(const int64_t size) {
-  auto memref = allocQubitRegisterStorage(size);
+QCProgramBuilder::allocQubitRegister(const int64_t size, const StringRef name) {
+  auto memref = allocQubitRegisterStorage(size, name);
 
   SmallVector<Value> qubits;
   qubits.reserve(size);
@@ -130,16 +130,24 @@ QCProgramBuilder::allocQubitRegister(const int64_t size) {
   return {.value = memref, .qubits = std::move(qubits)};
 }
 
-Value QCProgramBuilder::allocQubitRegisterStorage(const int64_t size) {
+Value QCProgramBuilder::allocQubitRegisterStorage(const int64_t size,
+                                                  const StringRef name) {
   checkFinalized();
   ensureAllocationMode(AllocationMode::Dynamic);
 
   if (size <= 0) {
     llvm::reportFatalUsageError("Size must be positive");
   }
+  if (!name.empty() && !quantumRegisterNames.insert(name).second) {
+    llvm::reportFatalUsageError("Quantum register names must be unique");
+  }
 
   auto memrefType = MemRefType::get({size}, QubitType::get(ctx));
-  auto memref = memref::AllocOp::create(*this, memrefType).getResult();
+  auto alloc = memref::AllocOp::create(*this, memrefType);
+  if (!name.empty()) {
+    alloc->setAttr(QUANTUM_REGISTER_NAME_ATTR, getStringAttr(name));
+  }
+  auto memref = alloc.getResult();
   allocatedQregs.insert(memref);
   return memref;
 }

mlir/lib/Dialect/QC/Translation/OpenQASMToQCEmitter.cpp

@@ -2407,7 +2407,7 @@ class OpenQASMToQCEmitter {
         return;
       }
       qubitValues[statement.reg] = builder.allocQubitRegisterStorage(
-          static_cast<int64_t>(declaration.width));
+          static_cast<int64_t>(declaration.width), declaration.name);
       return;
     }
     if (outputBitRegisters[statement.reg]) {

mlir/lib/Dialect/QC/Translation/TranslateQCToOpenQASM3.cpp

@@ -206,6 +206,13 @@ class OpenQASMEmitter {
     return uniqueName("out", nextScalar);
   }
 
+  [[nodiscard]] std::string quantumRegisterName(const StringRef requested) {
+    if (isValidOutputName(requested) && usedNames.insert(requested).second) {
+      return requested.str();
+    }
+    return uniqueName("q", nextQubit);
+  }
+
   [[nodiscard]] LogicalResult preflight() {
     SmallVector<func::FuncOp> functions(moduleOp.getOps<func::FuncOp>());
     if (functions.size() != 1) {
@@ -301,7 +308,12 @@ class OpenQASMEmitter {
         return fail(alloc, "only qubit and i1 memrefs are supported");
       }
       if (resource.kind == ResourceKind::Qubit) {
-        resource.name = uniqueName("q", nextQubit);
+        StringRef requested;
+        if (const auto attr = alloc->getAttrOfType<StringAttr>(
+                utils::QUANTUM_REGISTER_NAME_ATTR)) {
+          requested = attr.getValue();
+        }
+        resource.name = quantumRegisterName(requested);
       } else {
         resource.output = returnedMemrefs.contains(alloc.getResult());
         StringRef requested;

mlir/lib/Dialect/QC/Translation/TranslateQuantumComputationToQC.cpp

@@ -140,8 +140,8 @@ allocateQregs(QCProgramBuilder& builder,
   // Allocate quantum registers using the builder
   SmallVector<QregInfo> qregs;
   for (const auto* qregPtr : qregPtrs) {
-    auto qubitRegister =
-        builder.allocQubitRegister(static_cast<int64_t>(qregPtr->getSize()));
+    auto qubitRegister = builder.allocQubitRegister(
+        static_cast<int64_t>(qregPtr->getSize()), qregPtr->getName());
     qregs.emplace_back(qregPtr, std::move(qubitRegister.qubits));
   }
 

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

@@ -134,14 +134,22 @@ Value QCOProgramBuilder::staticQubit(const uint64_t index) {
 }
 
 QCOProgramBuilder::QubitRegister
-QCOProgramBuilder::allocQubitRegister(const int64_t size) {
+QCOProgramBuilder::allocQubitRegister(const int64_t size,
+                                      const StringRef name) {
   checkFinalized();
 
   if (size <= 0) {
     llvm::reportFatalUsageError("Size must be positive");
   }
+  if (!name.empty() && !quantumRegisterNames.insert(name).second) {
+    llvm::reportFatalUsageError("Quantum register names must be unique");
+  }
 
   auto qtensor = qtensorAlloc(size);
+  if (!name.empty()) {
+    qtensor.getDefiningOp()->setAttr(QUANTUM_REGISTER_NAME_ATTR,
+                                     getStringAttr(name));
+  }
 
   SmallVector<Value> qubits;
   qubits.reserve(size);

mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp

@@ -15,6 +15,7 @@
 #include "mlir/Dialect/QCO/Builder/QCOProgramBuilder.h"
 #include "mlir/Dialect/QCO/IR/QCODialect.h"
 #include "mlir/Dialect/QTensor/IR/QTensorDialect.h"
+#include "mlir/Dialect/Utils/Utils.h"
 #include "mlir/Support/IRVerification.h"
 #include "mlir/Support/Passes.h"
 #include "qc_programs.h"
@@ -25,6 +26,7 @@
 #include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/Dialect/MemRef/IR/MemRef.h>
 #include <mlir/Dialect/SCF/IR/SCF.h>
+#include <mlir/IR/BuiltinAttributes.h>
 #include <mlir/IR/DialectRegistry.h>
 #include <mlir/IR/MLIRContext.h>
 #include <mlir/IR/Value.h>
@@ -37,6 +39,7 @@
 #include <memory>
 #include <ostream>
 #include <string>
+#include <tuple>
 
 using namespace mlir;
 
@@ -83,6 +86,33 @@ static LogicalResult runQCOToQCConversion(ModuleOp module) {
   return pm.run(module);
 }
 
+TEST(QCOToQCRegressionTest, RetainsQuantumRegisterName) {
+  DialectRegistry registry;
+  registry.insert<qc::QCDialect, qco::QCODialect, qtensor::QTensorDialect,
+                  arith::ArithDialect, func::FuncDialect, memref::MemRefDialect,
+                  scf::SCFDialect>();
+  MLIRContext context(registry);
+  context.loadAllAvailableDialects();
+  qco::QCOProgramBuilder builder(&context);
+  builder.initialize();
+  std::ignore = builder.allocQubitRegister(2, "named_qubits");
+  auto moduleOp = builder.finalize();
+  ASSERT_TRUE(moduleOp);
+  ASSERT_TRUE(succeeded(runQCOToQCConversion(*moduleOp)));
+
+  memref::AllocOp allocation;
+  moduleOp->walk([&](memref::AllocOp op) {
+    if (isa<qc::QubitType>(op.getType().getElementType())) {
+      allocation = op;
+    }
+  });
+  ASSERT_TRUE(allocation);
+  const auto name =
+      allocation->getAttrOfType<StringAttr>(utils::QUANTUM_REGISTER_NAME_ATTR);
+  ASSERT_TRUE(name);
+  EXPECT_EQ(name.getValue(), "named_qubits");
+}
+
 static Value
 aliasSafeNestedForLoopCtrlOpWithExtractedQubit(qc::QCProgramBuilder& b) {
   auto reg = b.allocQubitRegister(4);