Use collection of measurements

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

@@ -1064,7 +1064,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
   );
 
@@ -1081,6 +1081,7 @@ class Measurement<string mnemonic> : QuakeOp<mnemonic, [MeasurementInterface,
   }];
 
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 def MxOp : Measurement<"mx"> {
@@ -1141,7 +1142,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,38 @@ 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 useCollection =
           (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 (useCollection) {
+          std::size_t totalSize = 0;
+          bool allKnown = true;
+          for (auto a : args) {
+            if (isa<quake::RefType>(a.getType())) {
+              totalSize++;
+            } else if (auto veqTy = dyn_cast<quake::VeqType>(a.getType())) {
+              if (veqTy.hasSpecifiedSize())
+                totalSize += veqTy.getSize();
+              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 (useCollection)
         resTy = cc::StdvecType::get(resTy);
       return pushValue(
           builder.create<quake::DiscriminateOp>(loc, resTy, measure));

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 LLVM::LLVMPointerType::get(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 LLVM::LLVMPointerType::get(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 LLVM::LLVMPointerType::get(ty.getContext());
+    });
   }
 };
 } // namespace

lib/Optimizer/Dialect/Quake/CanonicalPatterns.inc

@@ -695,6 +695,51 @@ 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 (isa<quake::RefType>(target.getType())) {
+        totalSize++;
+      } else if (auto veqTy = dyn_cast<quake::VeqType>(target.getType())) {
+        if (!veqTy.hasSpecifiedSize())
+          return failure();
+        totalSize += veqTy.getSize();
+      } 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

@@ -798,12 +798,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()))
@@ -831,19 +831,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,13 @@ 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 (auto veqTy = dyn_cast<quake::VeqType>(veq.getType());
+          veqTy && veqTy.hasSpecifiedSize())
+        resultType = quake::MeasurementsType::get(measure->getContext(),
+                                                  veqTy.getSize());
+      else
+        resultType = quake::MeasurementsType::getUnsized(measure->getContext());
       if (measure == analysis.lastMeasurement) {
         rewriter.replaceOpWithNewOp<quake::MzOp>(measure, TypeRange{resultType},
                                                  ValueRange{veq},
@@ -165,12 +171,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,11 +209,20 @@ 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 (auto veqTy = dyn_cast<quake::VeqType>(veq.getType());
+            veqTy && veqTy.hasSpecifiedSize())
+          resultType = quake::MeasurementsType::get(measure->getContext(),
+                                                    veqTy.getSize());
+        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();