* Retain variable names assigned to measurement results

* Disallow measure_result as argument to entry-point kernels
* WIP

Signed-off-by: Pradnya Khalate <pkhalate@nvidia.com>

Author: khalatepradnya

include/cudaq/Optimizer/Builder/Intrinsics.h

@@ -72,6 +72,10 @@ static constexpr const char llvmStackRestore[] = "llvm.stackrestore";
 
 static constexpr const char cudaqConvertToInteger[] =
     "__nvqpp_cudaqConvertToInteger";
+
+static constexpr const char cudaqConvertToBoolVector[] =
+    "__nvqpp_cudaq_to_bool_vector";
+
 /// Builder for lowering the clang AST to an IR for CUDA-Q. Lowering includes
 /// the transformation of both quantum and classical computation. Different
 /// features of the CUDA-Q programming model are lowered into different dialects

lib/Frontend/nvqpp/ASTBridge.cpp

@@ -73,6 +73,19 @@ static bool hasAnyQubitTypes(FunctionType funcTy) {
   return false;
 }
 
+// Check the builtin type `FunctionType` to see if it has any `MeasureType`
+// arguments.
+static bool hasMeasureResultArgs(FunctionType funcTy) {
+  for (auto ty : funcTy.getInputs()) {
+    if (isa<quake::MeasureType>(ty))
+      return true;
+    if (auto vecTy = dyn_cast<cudaq::cc::StdvecType>(ty))
+      if (isa<quake::MeasureType>(vecTy.getElementType()))
+        return true;
+  }
+  return false;
+}
+
 // Remove the Itanium mangling "_ZTS" prefix. This is to match the name returned
 // by `typeid(TYPE).name()`.
 static std::string
@@ -640,6 +653,7 @@ void ASTBridgeAction::ASTBridgeConsumer::HandleTranslationUnit(
       // Flag func as a quantum kernel.
       func->setAttr(kernelAttrName, unitAttr);
       if ((!hasAnyQubitTypes(func.getFunctionType())) &&
+          (!hasMeasureResultArgs(func.getFunctionType())) &&
           (!cudaq::ASTBridgeAction::ASTBridgeConsumer::isCustomOpGenerator(
               fdPair.second))) {
         // Flag func as an entry point to a quantum kernel.

lib/Frontend/nvqpp/ConvertDecl.cpp

@@ -169,6 +169,9 @@ bool QuakeBridgeVisitor::interceptRecordDecl(clang::RecordDecl *x) {
       auto fnTy = cast<FunctionType>(popType());
       return pushType(cc::IndirectCallableType::get(fnTy));
     }
+    // Measurement result type.
+    if (name == "measure_result")
+      return pushType(quake::MeasureType::get(ctx));
     if (!isInNamespace(x, "solvers") && !isInNamespace(x, "qec")) {
       auto loc = toLocation(x);
       TODO_loc(loc, "unhandled type, " + name + ", in cudaq namespace");
@@ -744,6 +747,11 @@ bool QuakeBridgeVisitor::VisitVarDecl(clang::VarDecl *x) {
       // and if so, find the mz and tag it with the variable name
       auto elementType = vecType.getElementType();
 
+      if (auto meas = initVec.getDefiningOp<quake::MeasurementInterface>()) {
+        meas.setRegisterName(builder.getStringAttr(x->getName()));
+        return true;
+      }
+
       // Drop out if this is not an i1
       if (!elementType.isIntOrFloat() ||
           elementType.getIntOrFloatBitWidth() != 1)
@@ -817,9 +825,8 @@ bool QuakeBridgeVisitor::VisitVarDecl(clang::VarDecl *x) {
 
   // If this was an auto var = mz(q), then we want to know the
   // var name, as it will serve as the classical bit register name
-  if (auto discr = initValue.getDefiningOp<quake::DiscriminateOp>())
-    if (auto mz = discr.getMeasurement().getDefiningOp<quake::MzOp>())
-      mz.setRegisterName(builder.getStringAttr(x->getName()));
+  if (auto meas = initValue.getDefiningOp<quake::MeasurementInterface>())
+    meas.setRegisterName(builder.getStringAttr(x->getName()));
 
   assert(initValue && "initializer value must be lowered");
   if (isa<IntegerType>(initValue.getType()) && isa<IntegerType>(type)) {

lib/Frontend/nvqpp/ConvertExpr.cpp

@@ -658,14 +658,34 @@ bool QuakeBridgeVisitor::VisitCastExpr(clang::CastExpr *x) {
   }
   case clang::CastKind::CK_UserDefinedConversion: {
     auto sub = popValue();
-    // castToTy is the converion function signature.
+    // castToTy is the conversion function signature.
     castToTy = popType();
     if (isa<IntegerType>(castToTy) && isa<IntegerType>(sub.getType())) {
       auto locSub = toLocation(x->getSubExpr());
       bool result = intToIntCast(locSub, sub);
       assert(result && "integer conversion failed");
       return result;
     }
+    auto i1Type = builder.getI1Type();
+
+    // Handle conversion of `measure_result` to `bool`.
+    if (isa<quake::MeasureType>(sub.getType()))
+      return pushValue(builder.create<quake::DiscriminateOp>(loc, i1Type, sub));
+
+    // Handle conversion of `std::vector<measure_result>` to `std::vector<bool>`
+    if (auto vecTy = dyn_cast<cc::StdvecType>(sub.getType()))
+      if (isa<quake::MeasureType>(vecTy.getElementType()))
+        return pushValue(builder.create<quake::DiscriminateOp>(
+            loc, cc::StdvecType::get(i1Type), sub));
+
+    // Handle pointer to `measure_result` (from vector element access)
+    if (auto ptrTy = dyn_cast<cc::PointerType>(sub.getType()))
+      if (isa<quake::MeasureType>(ptrTy.getElementType())) {
+        auto loaded = builder.create<cc::LoadOp>(loc, sub);
+        return pushValue(
+            builder.create<quake::DiscriminateOp>(loc, i1Type, loaded));
+      }
+
     TODO_loc(loc, "unhandled user-defined implicit conversion");
   }
   case clang::CastKind::CK_ConstructorConversion: {
@@ -1015,7 +1035,7 @@ bool QuakeBridgeVisitor::VisitMaterializeTemporaryExpr(
   // In those cases, there is nothing to materialize, so we can just pass the
   // Value on the top of the stack.
   if (isa<cc::CallableType, quake::VeqType, quake::RefType, cc::SpanLikeType,
-          quake::StateType>(ty))
+          quake::StateType, quake::MeasureType>(ty))
     return true;
 
   // If not one of the above special cases, then materialize the value to a
@@ -1520,7 +1540,13 @@ bool QuakeBridgeVisitor::VisitCallExpr(clang::CallExpr *x) {
   auto funcArity = func->getNumParams();
   SmallVector<Value> args = lastValues(funcArity);
   if (isa<clang::CXXMethodDecl>(func)) {
-    [[maybe_unused]] auto thisPtrValue = popValue();
+    auto thisPtrValue = popValue();
+
+    // Handle `measure_result` conversion operators
+    if (isa<clang::CXXConversionDecl>(func) &&
+        isInClassInNamespace(func, "measure_result", "cudaq")) {
+      return pushValue(thisPtrValue);
+    }
   }
   auto calleeOp = popValue();
 
@@ -1646,7 +1672,6 @@ 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 =
           (args.size() > 1) ||
           (args.size() == 1 && isa<quake::VeqType>(args[0].getType()));
@@ -1660,11 +1685,8 @@ bool QuakeBridgeVisitor::VisitCallExpr(clang::CallExpr *x) {
           return builder.create<quake::MyOp>(loc, measTy, args).getMeasOut();
         return builder.create<quake::MzOp>(loc, measTy, args).getMeasOut();
       }();
-      Type resTy = builder.getI1Type();
-      if (useStdvec)
-        resTy = cc::StdvecType::get(resTy);
-      return pushValue(
-          builder.create<quake::DiscriminateOp>(loc, resTy, measure));
+      // No more discrimination needed, just return the measurement.
+      return pushValue(measure);
     }
 
     // Handle the quantum gate set.
@@ -2136,6 +2158,28 @@ bool QuakeBridgeVisitor::VisitCallExpr(clang::CallExpr *x) {
               .getResult(0));
     }
 
+    if (funcName == "to_bool_vector") {
+      // args[0] is !cc.stdvec<!quake.measure> from mz()
+      auto arg = args[0];
+      // Insert discriminate if needed
+      if (auto vecTy = dyn_cast<cc::StdvecType>(arg.getType())) {
+        if (isa<quake::MeasureType>(vecTy.getElementType())) {
+          auto i1Ty = builder.getI1Type();
+          arg = builder.create<quake::DiscriminateOp>(
+              loc, cc::StdvecType::get(i1Ty), arg);
+        }
+      }
+      IRBuilder irBuilder(builder.getContext());
+      if (failed(irBuilder.loadIntrinsic(module, cudaqConvertToBoolVector))) {
+        reportClangError(x, mangler, "cannot load cudaqConvertToBoolVector");
+        return false;
+      }
+      return pushValue(builder
+                           .create<func::CallOp>(loc, arg.getType(),
+                                                 cudaqConvertToBoolVector, arg)
+                           .getResult(0));
+    }
+
     if (funcName == "slice_vector") {
       auto svecTy = dyn_cast<cc::SpanLikeType>(args[0].getType());
       auto eleTy = svecTy.getElementType();
@@ -3217,6 +3261,18 @@ bool QuakeBridgeVisitor::VisitCXXConstructExpr(clang::CXXConstructExpr *x) {
     return pushValue(builder.create<cc::LoadOp>(loc, copyObj));
   }
 
+  // Handle `measure_result` copy constructor
+  if (ctor->isCopyConstructor() &&
+      isInClassInNamespace(ctor, "measure_result", "cudaq")) {
+    // The source is a pointer to measure_result (!cc.ptr<!quake.measure>)
+    // Just load and return the value
+    assert(x->getNumArgs() == 1);
+    auto srcPtr = popValue();
+    if (auto ptrTy = dyn_cast<cc::PointerType>(srcPtr.getType()))
+      if (isa<quake::MeasureType>(ptrTy.getElementType()))
+        return pushValue(builder.create<cc::LoadOp>(loc, srcPtr));
+  }
+
   // TODO: remove this when we can handle ctors more generally.
   if (!ctor->isDefaultConstructor()) {
     LLVM_DEBUG(llvm::dbgs() << ctorName << " - unhandled ctor:\n"; x->dump());

lib/Frontend/nvqpp/ConvertStmt.cpp

@@ -361,6 +361,12 @@ bool QuakeBridgeVisitor::VisitReturnStmt(clang::ReturnStmt *x) {
       Value tySize;
       if (!cudaq::cc::isDynamicType(eleTy))
         tySize = irb.getByteSizeOfType(loc, eleTy);
+      if (isa<quake::MeasureType>(eleTy)) {
+        /// FIXME: Confirm that this is okay.
+        tySize = irb.getByteSizeOfType(loc, builder.getI32Type());
+      } else {
+        tySize = irb.getByteSizeOfType(loc, eleTy);
+      }
       if (!tySize) {
         // TODO: we need to recursively create copies of all
         // dynamic memory used within the type. See the

lib/Frontend/nvqpp/ConvertType.cpp

@@ -124,13 +124,24 @@ static bool isFunctionCallable(Type t) {
   return false;
 }
 
+static bool isMeasureResultType(Type t) { return isa<quake::MeasureType>(t); }
+
+static bool isMeasureResultSequenceType(Type t) {
+  if (auto vec = dyn_cast<cudaq::cc::SpanLikeType>(t)) {
+    auto eleTy = vec.getElementType();
+    return isMeasureResultType(eleTy) || isMeasureResultSequenceType(eleTy);
+  }
+  return isMeasureResultType(t);
+}
+
 /// Return true if and only if \p t is a (simple) arithmetic type, an arithmetic
 /// sequence type (possibly dynamic in length), or a static product type of
 /// arithmetic types. Note that this means a product type with a dynamic
 /// sequence of arithmetic types is \em disallowed.
 static bool isKernelResultType(Type t) {
   return isArithmeticType(t) || isArithmeticSequenceType(t) ||
-         isStaticArithmeticProductType(t);
+         isStaticArithmeticProductType(t) || isMeasureResultType(t) ||
+         isMeasureResultSequenceType(t);
 }
 
 /// Return true if and only if \p t is a (simple) arithmetic type, an possibly

lib/Optimizer/Builder/Intrinsics.cpp

@@ -337,6 +337,13 @@ static constexpr IntrinsicCode intrinsicTable[] = {
   func.func private @__nvqpp_cudaq_state_numberOfQubits(%p : !cc.ptr<!quake.state>) -> i64
 )#"},
 
+    {cudaq::cudaqConvertToBoolVector, {}, R"#(
+  func.func private @__nvqpp_cudaq_to_bool_vector(%arg : !cc.stdvec<i1>) -> !cc.stdvec<i1> {
+    // TODO: Implement
+    return %arg : !cc.stdvec<i1>
+  }
+)#"},
+
     {cudaq::runtime::bindingDeconstructString,
      {},
      "func.func private @__nvqpp_deconstructString(!cc.ptr<i8>)"},

lib/Optimizer/CodeGen/ConvertToQIRAPI.cpp

@@ -579,8 +579,9 @@ struct DiscriminateOpRewrite
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = disc.getLoc();
     Value m = adaptor.getMeasurement();
-    auto i1PtrTy = cudaq::cc::PointerType::get(rewriter.getI1Type());
-    auto cast = rewriter.create<cudaq::cc::CastOp>(loc, i1PtrTy, m);
+    auto resultTy = typeConverter->convertType(disc.getResult().getType());
+    auto ptrTy = cudaq::cc::PointerType::get(resultTy);
+    auto cast = rewriter.create<cudaq::cc::CastOp>(loc, ptrTy, m);
     rewriter.replaceOpWithNewOp<cudaq::cc::LoadOp>(disc, cast);
     return success();
   }

lib/Optimizer/Dialect/CC/CCOps.cpp

@@ -20,6 +20,8 @@
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+/// FIXME: This seems wrong to add in `cc` dialect!
+#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
 
 using namespace mlir;
 
@@ -77,6 +79,10 @@ Value cudaq::cc::getByteSizeOfType(OpBuilder &builder, Location loc, Type ty,
                 // we're assuming pointers are 64 bits.
                 return {8};
               })
+          .Case([](quake::MeasureType) -> std::optional<std::int32_t> {
+            /// FIXME: What should this be?
+            return {4};
+          })
           .Default({});
 
   if (rawSize)

runtime/cudaq/qis/measure_result.h

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include <optional>
 #include <vector>
 
 namespace cudaq {
@@ -26,24 +27,14 @@ class measure_result {
   int result = 0;
 
   /// Unique integer for measure result identification
-  std::size_t uniqueId = 0;
+  std::optional<std::size_t> uniqueId = std::nullopt;
 
 public:
-  measure_result(int res, std::size_t id) : result(res), uniqueId(id) {}
   measure_result(int res) : result(res) {}
+  measure_result(int res, std::size_t id) : result(res), uniqueId(id) {}
 
   operator int() const { return result; }
   operator bool() const { return __nvqpp__MeasureResultBoolConversion(result); }
-
-  /// TODO: This needs to be revisited to support MLIR mode properly.
-  static std::vector<bool>
-  to_bool_vector(const std::vector<measure_result> &results) {
-    std::vector<bool> boolResults;
-    boolResults.reserve(results.size());
-    for (const auto &res : results)
-      boolResults.push_back(static_cast<bool>(res));
-    return boolResults;
-  }
 };
 
 } // namespace cudaq