Use collection of measurements

include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h

@@ -81,6 +81,8 @@ static constexpr const char QIRArrayConcatArray[] =
     "__quantum__rt__array_concatenate";
 static constexpr const char QIRArrayCreateArray[] =
     "__quantum__rt__array_create_1d";
+static constexpr const char QIRResultArrayCreate[] =
+    "__quantum__rt__result_array_create_1d";
 
 /// Dynamic qubit management helper functions. These are currently only used by
 /// the NVQIR simulator.

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

@@ -1066,7 +1066,7 @@ class Measurement<string mnemonic> : QuakeOp<mnemonic, [MeasurementInterface,
     OptionalAttr<StrAttr>:$registerName
   );
   let results = (outs
-    AnyTypeOf<[MeasureType, StdvecOf<[MeasureType]>]>:$measOut,
+    AnyTypeOf<[MeasureType, MeasurementsType]>:$measOut,
     Variadic<WireType>:$wires
   );
 
@@ -1083,6 +1083,7 @@ class Measurement<string mnemonic> : QuakeOp<mnemonic, [MeasurementInterface,
   }];
 
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 def MxOp : Measurement<"mx"> {
@@ -1143,7 +1144,7 @@ def quake_DiscriminateOp : QuakeOp<"discriminate", [Pure]> {
   }];
 
   let arguments = (ins
-    AnyTypeOf<[MeasureType, StdvecOf<[MeasureType]>]>:$measurement
+    AnyTypeOf<[MeasureType, MeasurementsType]>:$measurement
   );
   let results = (outs
     AnyTypeOf<[AnySignlessInteger, StdvecOf<[AnySignlessInteger]>]>

include/cudaq/Optimizer/Transforms/Passes.td

@@ -219,15 +219,15 @@ def CombineMeasurements :
       %1 = ... : !quake.veq<4>
       %2 = quake.subveq %1, %c2, %c3 : (!quake.veq<4>, i32, i32) ->
             !quake.veq<2>
-      %measOut = quake.mz %2 : (!quake.veq<2>) -> !cc.stdvec<!quake.measure>
+      %measOut = quake.mz %2 : (!quake.veq<2>) -> !quake.measurements<2>
     }
     ```
     with:
     ```
     func.func @kernel() attributes {"cudaq-entrypoint", ["output_names",
     "[[[0,[1,\22q0\22]],[1,[2,\22q1\22]]]]"]} {
       %1 = ... : !quake.veq<4>
-      %measOut = quake.mz %1 : (!quake.veq<4>) -> !cc.stdvec<!quake.measure>
+      %measOut = quake.mz %1 : (!quake.veq<4>) -> !quake.measurements<4>
     }
     ```
   }];

lib/Frontend/nvqpp/ConvertExpr.cpp

@@ -1647,21 +1647,34 @@ bool QuakeBridgeVisitor::VisitCallExpr(clang::CallExpr *x) {
 
     if (funcName == "mx" || funcName == "my" || funcName == "mz") {
       // Measurements always return a bool or a std::vector<bool>.
-      bool useStdvec =
+      bool useMeasurements =
           (args.size() > 1) ||
           (args.size() == 1 && isa<quake::VeqType>(args[0].getType()));
       auto measure = [&]() -> Value {
         Type measTy = quake::MeasureType::get(builder.getContext());
-        if (useStdvec)
-          measTy = cc::StdvecType::get(measTy);
+        if (useMeasurements) {
+          std::size_t totalSize = 0;
+          bool allKnown = true;
+          for (auto a : args) {
+            if (quake::isConstantQuantumRefType(a.getType()))
+              totalSize += quake::getAllocationSize(a.getType());
+            else
+              allKnown = false;
+          }
+          if (allKnown && totalSize > 0)
+            measTy =
+                quake::MeasurementsType::get(builder.getContext(), totalSize);
+          else
+            measTy = quake::MeasurementsType::getUnsized(builder.getContext());
+        }
         if (funcName == "mx")
           return builder.create<quake::MxOp>(loc, measTy, args).getMeasOut();
         if (funcName == "my")
           return builder.create<quake::MyOp>(loc, measTy, args).getMeasOut();
         return builder.create<quake::MzOp>(loc, measTy, args).getMeasOut();
       }();
       Type resTy = builder.getI1Type();
-      if (useStdvec)
+      if (useMeasurements)
         resTy = cc::StdvecType::get(resTy);
       return pushValue(
           builder.create<quake::DiscriminateOp>(loc, resTy, measure));

lib/Optimizer/Builder/Intrinsics.cpp

@@ -536,6 +536,7 @@ static constexpr IntrinsicCode intrinsicTable[] = {
   func.func private @__quantum__rt__qubit_release(!qir_qubit)
 
   func.func private @__quantum__rt__array_create_1d(i32, i64) -> !qir_array
+  func.func private @__quantum__rt__result_array_create_1d(i64) -> !qir_array
   func.func private @__quantum__rt__array_concatenate(!qir_array, !qir_array) -> !qir_array
   func.func private @__quantum__rt__array_get_size_1d(!qir_array) -> i64
   func.func private @__quantum__rt__array_slice(!qir_array, i32, i64, i64, i64) -> !qir_array

lib/Optimizer/CodeGen/ConvertToExecMgr.cpp

@@ -51,6 +51,9 @@ struct QuakeTypeConverter : public TypeConverter {
     addConversion([](quake::MeasureType ty) {
       return IntegerType::get(ty.getContext(), 64);
     });
+    addConversion([](quake::MeasurementsType ty) {
+      return cudaq::cc::PointerType::get(IntegerType::get(ty.getContext(), 8));
+    });
   }
 };
 

lib/Optimizer/CodeGen/ConvertToQIR.cpp

@@ -209,6 +209,9 @@ void cudaq::opt::initializeTypeConversions(LLVMTypeConverter &typeConverter) {
   typeConverter.addConversion([](quake::MeasureType type) {
     return IntegerType::get(type.getContext(), 1);
   });
+  typeConverter.addConversion([](quake::MeasurementsType type) {
+    return getArrayType(type.getContext());
+  });
   cudaq::opt::populateCCTypeConversions(&typeConverter);
 }
 

lib/Optimizer/CodeGen/ConvertToQIRAPI.cpp

@@ -135,6 +135,9 @@ struct QIRAPITypeConverter : public TypeConverter {
         [&](quake::CableType ty) { return getArrayType(ty.getContext()); });
     addConversion(
         [&](quake::MeasureType ty) { return getResultType(ty.getContext()); });
+    addConversion([&](quake::MeasurementsType ty) {
+      return getArrayType(ty.getContext());
+    });
     addConversion([&](quake::StruqType ty) { return convertStruqType(ty); });
   }
 

lib/Optimizer/CodeGen/WireSetsToProfileQIR.cpp

@@ -70,6 +70,9 @@ struct QuakeTypeConverter : public TypeConverter {
     addConversion([](quake::MeasureType ty) {
       return cudaq::opt::getResultType(ty.getContext());
     });
+    addConversion([](quake::MeasurementsType ty) {
+      return cudaq::opt::getArrayType(ty.getContext());
+    });
   }
 };
 } // namespace

lib/Optimizer/Dialect/Quake/CanonicalPatterns.inc

@@ -715,6 +715,46 @@ struct MergeRotationPattern : public OpRewritePattern<OP> {
   }
 };
 
+// %0 = quake.alloca !quake.veq<2>
+// %1 = quake.mz %0 : (!quake.veq<2>) -> !quake.measurements<?>
+// ────────────────────────────────────────────────────────────
+// %0 = quake.alloca !quake.veq<2>
+// %1 = quake.mz %0 : (!quake.veq<2>) -> !quake.measurements<2>
+template <typename MeasOp>
+struct FuseSizeToMeasurementPattern : public OpRewritePattern<MeasOp> {
+  using OpRewritePattern<MeasOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(MeasOp measOp,
+                                PatternRewriter &rewriter) const override {
+    auto measTy =
+        dyn_cast<quake::MeasurementsType>(measOp.getMeasOut().getType());
+    if (!measTy || measTy.hasSpecifiedSize())
+      return failure();
+
+    std::size_t totalSize = 0;
+    for (auto target : measOp.getTargets()) {
+      if (quake::isConstantQuantumRefType(target.getType()))
+        totalSize += quake::getAllocationSize(target.getType());
+      else
+        return failure();
+    }
+    if (totalSize == 0)
+      return failure();
+
+    auto newMeasTy =
+        quake::MeasurementsType::get(rewriter.getContext(), totalSize);
+    SmallVector<Type> resultTypes;
+    resultTypes.push_back(newMeasTy);
+    for (unsigned i = 1; i < measOp->getNumResults(); ++i)
+      resultTypes.push_back(measOp->getResult(i).getType());
+
+    rewriter.replaceOpWithNewOp<MeasOp>(measOp, TypeRange{resultTypes},
+                                        measOp.getTargets(),
+                                        measOp.getRegisterNameAttr());
+    return success();
+  }
+};
+
 // Forward the argument to a relax_size to the users for all users that are
 // quake operations. All quake ops that take a sized veq argument are
 // polymorphic on all veq types. If the op is not a quake op, then maintain

lib/Optimizer/Dialect/Quake/QuakeOps.cpp

@@ -848,12 +848,12 @@ LogicalResult verifyMeasurements(MEAS op, TypeRange targetsType,
                                  const Type bitsType) {
   if (failed(verifyWireResultsAreLinear(op)))
     return failure();
-  bool mustBeStdvec =
+  bool mustBeCollection =
       targetsType.size() > 1 ||
       (targetsType.size() == 1 && isa<quake::VeqType>(targetsType[0]));
-  if (mustBeStdvec) {
-    if (!isa<cudaq::cc::StdvecType>(op.getMeasOut().getType()))
-      return op.emitOpError("must return `!cc.stdvec<!quake.measure>`, when "
+  if (mustBeCollection) {
+    if (!isa<quake::MeasurementsType>(op.getMeasOut().getType()))
+      return op.emitOpError("must return `!quake.measurements`, when "
                             "measuring a qreg, a series of qubits, or both");
   } else {
     if (!isa<quake::MeasureType>(op.getMeasOut().getType()))
@@ -881,19 +881,34 @@ LogicalResult quake::MzOp::verify() {
                             getMeasOut().getType());
 }
 
+void quake::MxOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                              MLIRContext *context) {
+  patterns.add<FuseSizeToMeasurementPattern<quake::MxOp>>(context);
+}
+
+void quake::MyOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                              MLIRContext *context) {
+  patterns.add<FuseSizeToMeasurementPattern<quake::MyOp>>(context);
+}
+
+void quake::MzOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                              MLIRContext *context) {
+  patterns.add<FuseSizeToMeasurementPattern<quake::MzOp>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // Discriminate
 //===----------------------------------------------------------------------===//
 
 LogicalResult quake::DiscriminateOp::verify() {
-  if (isa<cudaq::cc::StdvecType>(getMeasurement().getType())) {
+  if (isa<quake::MeasurementsType>(getMeasurement().getType())) {
     auto stdvecTy = dyn_cast<cudaq::cc::StdvecType>(getResult().getType());
     if (!stdvecTy || !isa<IntegerType>(stdvecTy.getElementType()))
       return emitOpError("must return a !cc.stdvec<integral> type, when "
-                         "discriminating a qreg, a series of qubits, or both");
+                         "discriminating a measurements collection");
   } else {
-    auto measTy = isa<quake::MeasureType>(getMeasurement().getType());
-    if (!measTy || !isa<IntegerType>(getResult().getType()))
+    if (!isa<quake::MeasureType>(getMeasurement().getType()) ||
+        !isa<IntegerType>(getResult().getType()))
       return emitOpError(
           "must return integral type when discriminating exactly one qubit");
   }

lib/Optimizer/Transforms/AddMeasurements.cpp

@@ -86,9 +86,15 @@ addMeasurements(func::FuncOp funcOp, SmallVector<quake::AllocaOp> &allocations,
   builder.setInsertionPointToEnd(newBlock);
   auto measTy = quake::MeasureType::get(builder.getContext());
   for (auto &[index, alloca] : llvm::enumerate(allocations)) {
-    if (isa<quake::VeqType>(alloca.getType())) {
-      auto stdvecTy = cudaq::cc::StdvecType::get(measTy);
-      builder.create<quake::MzOp>(loc, stdvecTy,
+    if (auto veqTy = dyn_cast<quake::VeqType>(alloca.getType())) {
+      Type measurementsTy;
+      if (veqTy.hasSpecifiedSize())
+        measurementsTy =
+            quake::MeasurementsType::get(builder.getContext(), veqTy.getSize());
+      else
+        measurementsTy =
+            quake::MeasurementsType::getUnsized(builder.getContext());
+      builder.create<quake::MzOp>(loc, measurementsTy,
                                   ValueRange{alloca.getResult()});
     } else {
       builder.create<quake::MzOp>(loc, measTy, alloca.getResult());

lib/Optimizer/Transforms/CombineMeasurements.cpp

@@ -109,7 +109,7 @@ class ExtendQubitMeasurePattern : public OpRewritePattern<quake::MzOp> {
   // with:
   // ```
   //   %1 = ... : !quake.veq<4>
-  //   %measOut = quake.mz %1 : (!quake.veq<4>) -> !cc.stdvec<!quake.measure>
+  //   %measOut = quake.mz %1 : (!quake.veq<4>) -> !quake.measurements<4>
   // ```
   // And collect output names information:  `"[[[0,[1,"q0"]],[1,[2,"q1"]]]]"`
   LogicalResult matchAndRewrite(quake::MzOp measure,
@@ -132,7 +132,12 @@ class ExtendQubitMeasurePattern : public OpRewritePattern<quake::MzOp> {
       analysis.resultQubitVals[offset] =
           std::make_pair(idx, std::to_string(idx));
 
-      auto resultType = cudaq::cc::StdvecType::get(measure.getType(0));
+      Type resultType;
+      if (quake::isConstantQuantumRefType(veq.getType()))
+        resultType = quake::MeasurementsType::get(
+            measure->getContext(), quake::getAllocationSize(veq.getType()));
+      else
+        resultType = quake::MeasurementsType::getUnsized(measure->getContext());
       if (measure == analysis.lastMeasurement) {
         rewriter.replaceOpWithNewOp<quake::MzOp>(measure, TypeRange{resultType},
                                                  ValueRange{veq},
@@ -165,12 +170,12 @@ class ExtendVeqMeasurePattern : public OpRewritePattern<quake::MzOp> {
   //   %1 = ... : !quake.veq<4>
   //   %2 = quake.subveq %1, %c1, %c2 : (!quake.veq<4>, i32, i32) ->
   //        !quake.veq<2>
-  //   %measOut = quake.mz %2 : (!quake.veq<2>) -> !cc.stdvec<!quake.measure>
+  //   %measOut = quake.mz %2 : (!quake.veq<2>) -> !quake.measurements<2>
   // ```
   // with:
   // ```
   //   %1 = ... : !quake.veq<4>
-  //   %measOut = quake.mz %1 : (!quake.veq<4>) -> !cc.stdvec<!quake.measure>
+  //   %measOut = quake.mz %1 : (!quake.veq<4>) -> !quake.measurements<4>
   // ```
   // And collect output names information:  `"[[[0,[1,"q0"]],[1,[2,"q1"]]]]"`
   LogicalResult matchAndRewrite(quake::MzOp measure,
@@ -203,12 +208,21 @@ class ExtendVeqMeasurePattern : public OpRewritePattern<quake::MzOp> {
         analysis.resultQubitVals[offset] = std::make_pair(i, std::to_string(i));
       }
 
-      if (measure == analysis.lastMeasurement)
-        rewriter.replaceOpWithNewOp<quake::MzOp>(
-            measure, measure.getResultTypes(), ValueRange{subveq.getVeq()},
-            measure.getRegisterNameAttr());
-      else if (measure.use_empty())
+      if (measure == analysis.lastMeasurement) {
+        auto veq = subveq.getVeq();
+        Type resultType;
+        if (quake::isConstantQuantumRefType(veq.getType()))
+          resultType = quake::MeasurementsType::get(
+              measure->getContext(), quake::getAllocationSize(veq.getType()));
+        else
+          resultType =
+              quake::MeasurementsType::getUnsized(measure->getContext());
+        rewriter.replaceOpWithNewOp<quake::MzOp>(measure, TypeRange{resultType},
+                                                 ValueRange{veq},
+                                                 measure.getRegisterNameAttr());
+      } else if (measure.use_empty()) {
         rewriter.eraseOp(measure);
+      }
 
       return success();
     }