🐛 Enforce terminal measurements in Base QIR

mlir/include/mlir/Conversion/QCToQIR/QIRBase/QCToQIRBase.td

@@ -23,10 +23,18 @@ def QCToQIRBase : Pass<"qc-to-qir-base", "mlir::ModuleOp"> {
     - The input entry function must consist of a single block.
       Multi-block input functions are currently not supported.
     - The program must have straight-line control flow (i.e., Base Profile QIR).
+    - A measured qubit must not be used by another quantum instruction,
+      including another measurement. Gates on independent qubits may follow
+      measurements in the input.
+    - Explicit static qubit IDs cannot be mixed with qubit allocations.
+    - Qubit-register loads require constant, in-bounds indices into statically
+      sized allocations. Register aliases must be resolved before conversion.
 
     Behavior:
 
-    - Each QC quantum operation is replaced by a call to the corresponding QIR function in the LLVM dialect.
+    - Each QC quantum operation is replaced in place by its QIR call. After
+      validating qubit usage in instruction order, the pass moves terminal
+      measurements to the irreversible operations block.
     - Required QIR module flags are attached as attributes to the entry function.
     - The pass transforms the single-block entry function into four blocks to satisfy QIR Base Profile constraints:
       0. Initialization block: Sets up the execution environment and performs required runtime initialization.

mlir/include/mlir/Conversion/QCToQIR/QIRCommon/QIRCommon.h

@@ -12,6 +12,7 @@
 
 #include "mlir/Dialect/QIR/Utils/QIRUtils.h"
 
+#include <llvm/ADT/DenseMap.h>
 #include <llvm/ADT/StringMap.h>
 #include <llvm/Support/Allocator.h>
 #include <llvm/Support/StringSaver.h>
@@ -46,6 +47,9 @@ struct LoweringState {
   /// Cache static qubit pointers for reuse
   DenseMap<int64_t, Value> staticQubits;
 
+  /// Canonical Base-profile pointers for constant qubit-register elements.
+  DenseMap<std::pair<Value, int64_t>, Value> staticRegisterQubits;
+
   /// Cache qubit register sizes for reuse
   DenseMap<Value, Value> qregSizes;
 

mlir/lib/Conversion/QCToQIR/QIRBase/QCToQIRBase.cpp

@@ -20,6 +20,7 @@
 #include "mlir/Dialect/QIR/QIRDefinitions.h"
 #include "mlir/Dialect/QIR/Utils/QIRUtils.h"
 
+#include <llvm/ADT/DenseSet.h>
 #include <mlir/Conversion/ArithToLLVM/ArithToLLVM.h>
 #include <mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h>
 #include <mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h>
@@ -40,6 +41,7 @@
 #include <mlir/IR/OpDefinition.h>
 #include <mlir/IR/PatternMatch.h>
 #include <mlir/IR/Region.h>
+#include <mlir/IR/Value.h>
 #include <mlir/IR/ValueRange.h>
 #include <mlir/Pass/PassManager.h>
 #include <mlir/Support/LLVM.h>
@@ -49,6 +51,7 @@
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
+#include <iterator>
 #include <utility>
 #include <variant>
 
@@ -85,6 +88,39 @@ static FailureOr<Value> resolveRegisterMeasurement(LoweringState& state,
   return results[static_cast<size_t>(*indexValue)];
 }
 
+/// Validates canonical qubit pointers before moving measurements out of order.
+static LogicalResult moveTerminalMeasurements(Block& body,
+                                              Block& measurements) {
+  DenseSet<Value> measuredQubits;
+  SmallVector<LLVM::CallOp> measurementCalls;
+  for (auto call : body.getOps<LLVM::CallOp>()) {
+    if (!call.getCallee() ||
+        !call.getCallee()->starts_with("__quantum__qis__")) {
+      continue;
+    }
+    const bool isMeasurement = call.getCallee() == QIR_MEASURE;
+    /// Measurement's second pointer identifies a result, not a qubit.
+    auto operands = call.getOperands();
+    if (isMeasurement) {
+      operands = operands.take_front(1);
+    }
+    for (auto operand : operands) {
+      if (measuredQubits.contains(operand)) {
+        return call.emitError(
+            "QIR Base Profile forbids using a qubit after measurement");
+      }
+    }
+    if (isMeasurement) {
+      measuredQubits.insert(call.getOperand(0));
+      measurementCalls.push_back(call);
+    }
+  }
+  for (auto call : measurementCalls) {
+    call->moveBefore(measurements.getTerminator());
+  }
+  return success();
+}
+
 namespace {
 
 /**
@@ -171,6 +207,9 @@ struct ConvertMemRefAllocOp final
   LogicalResult
   matchAndRewrite(memref::AllocOp op, OpAdaptor /*adaptor*/,
                   ConversionPatternRewriter& rewriter) const override {
+    if (failed(getState().ensureAllocationMode(AllocationMode::Dynamic, op))) {
+      return failure();
+    }
     rewriter.eraseOp(op);
     return success();
   }
@@ -206,16 +245,23 @@ struct ConvertMemRefLoadOp final : StatefulOpConversionPattern<memref::LoadOp> {
       return rewriter.notifyMatchFailure(
           op, "Only one-dimensional registers are supported");
     }
-    // Save current insertion point
-    const OpBuilder::InsertionGuard guard(rewriter);
-
-    // Switch to entry block
-    rewriter.setInsertionPoint(state.entryBlock->getTerminator());
-
-    auto nqubits = state.staticQubits.size();
-    auto qubit = createPointerFromIndex(rewriter, op.getLoc(),
-                                        static_cast<int64_t>(nqubits));
-    state.staticQubits.try_emplace(static_cast<int64_t>(nqubits), qubit);
+    const auto index = getConstantIntValue(op.getIndices().front());
+    if (!index || ShapedType::isDynamic(shape.front()) ||
+        !op.getMemref().getDefiningOp<memref::AllocOp>()) {
+      return op.emitError("QIR Base Profile requires constant indices into "
+                          "statically allocated qubit registers");
+    }
+    if (*index < 0 || *index >= shape.front()) {
+      return op.emitError("qubit-register index is out of bounds");
+    }
+    auto& qubit = state.staticRegisterQubits[{op.getMemref(), *index}];
+    if (!qubit) {
+      const OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPoint(state.entryBlock->getTerminator());
+      const auto id = static_cast<int64_t>(state.staticQubits.size());
+      qubit = createPointerFromIndex(rewriter, op.getLoc(), id);
+      state.staticQubits.try_emplace(id, qubit);
+    }
     rewriter.replaceOp(op, qubit);
 
     return success();
@@ -262,6 +308,9 @@ struct ConvertQCAllocOp final : StatefulOpConversionPattern<AllocOp> {
   matchAndRewrite(AllocOp op, OpAdaptor /*adaptor*/,
                   ConversionPatternRewriter& rewriter) const override {
     auto& state = getState();
+    if (failed(state.ensureAllocationMode(AllocationMode::Dynamic, op))) {
+      return failure();
+    }
 
     const OpBuilder::InsertionGuard guard(rewriter);
 
@@ -331,8 +380,8 @@ struct ConvertQCMeasureOp final : StatefulOpConversionPattern<MeasureOp> {
       result = getResultPtr(state, op.getOperation(), rewriter);
     }
 
-    // Emit the measurement in the measurements block
-    rewriter.setInsertionPoint(state.measurementsBlock->getTerminator());
+    /// Preserve instruction order until terminal measurements are verified.
+    rewriter.setInsertionPoint(op);
     auto fnSig = LLVM::LLVMFunctionType::get(voidType, {ptrType, ptrType});
     auto fnDec =
         getOrCreateFunctionDeclaration(rewriter, op, QIR_MEASURE, fnSig);
@@ -468,8 +517,8 @@ struct QCToQIRBase final : impl::QCToQIRBaseBase<QCToQIRBase> {
    * Insert the `__quantum__rt__initialize` call.
    *
    * **Stage 4: QC to LLVM**
-   * Convert QC dialect operations to QIR calls and add output recording to the
-   * output block.
+   * Convert QC dialect operations in place, validate and move terminal
+   * measurements, and add output recording to the output block.
    *
    * **Stage 5: Standard dialects to LLVM**
    * Convert arith and control flow dialects to LLVM (for index arithmetic and
@@ -554,6 +603,11 @@ struct QCToQIRBase final : impl::QCToQIRBaseBase<QCToQIRBase> {
         return;
       }
 
+      auto& body = *std::next(main.getBody().begin());
+      if (failed(moveTerminalMeasurements(body, *state.measurementsBlock))) {
+        signalPassFailure();
+        return;
+      }
       addOutputRecording(main, ctx, state);
     }
 

mlir/unittests/Compiler/test_compiler_pipeline.cpp

@@ -834,6 +834,30 @@ if (flag) {
   EXPECT_TRUE(qir->llvmIR().has_value());
 }
 
+TEST_F(CompilerPipelineTest, BaseMeasurementMayBeInsertedIntoFreedQTensor) {
+  auto qco = QCOProgram::fromMLIRString(R"mlir(module {
+    func.func @main() -> i1 attributes {mqt.entry_point} {
+      %c0 = arith.constant 0 : index
+      %c1 = arith.constant 1 : index
+      %reg = qtensor.alloc(%c1) : tensor<1x!qco.qubit>
+      %rest, %qubit = qtensor.extract %reg[%c0] : tensor<1x!qco.qubit>
+      %out, %result = qco.measure %qubit : !qco.qubit
+      %final = qtensor.insert %out into %rest[%c0] : tensor<1x!qco.qubit>
+      qtensor.dealloc %final : tensor<1x!qco.qubit>
+      return %result : i1
+    }
+  })mlir");
+  ASSERT_TRUE(qco);
+  auto qc = std::move(*qco).intoQC();
+  ASSERT_TRUE(qc);
+  auto qir = std::move(*qc).intoQIR(QIRProfile::Base);
+  ASSERT_TRUE(qir);
+  const auto llvmIR = qir->llvmIR();
+  ASSERT_TRUE(llvmIR);
+  EXPECT_NE(llvmIR->find("call void @__quantum__qis__mz__body"),
+            std::string::npos);
+}
+
 TEST_F(CompilerPipelineTest, EmitsQIR21ProfileModuleFlags) {
   constexpr llvm::StringLiteral source = R"qasm(
 OPENQASM 3.0;