* Clean-up

* Handle the size for `measure_result` since it behaves differently on
  host and on device
* Log result output

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

Author: khalatepradnya

lib/Frontend/nvqpp/ConvertStmt.cpp

@@ -359,9 +359,8 @@ bool QuakeBridgeVisitor::VisitReturnStmt(clang::ReturnStmt *x) {
       };
       IRBuilder irb(builder);
       Value tySize;
-      if (!cudaq::cc::isDynamicType(eleTy)) {
+      if (!cudaq::cc::isDynamicType(eleTy))
         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/Optimizer/CodeGen/ConvertCCToLLVM.cpp

@@ -81,18 +81,19 @@ void cudaq::opt::populateCCTypeConversions(LLVMTypeConverter *converter) {
     return LLVM::LLVMStructType::getLiteral(type.getContext(), members,
                                             type.getPacked());
   });
-  converter->addConversion([](quake::MeasureType type) -> Type {
-    auto ctx = type.getContext();
-    auto i32Ty = IntegerType::get(ctx, 32);
-    auto i64Ty = IntegerType::get(ctx, 64);
-    return LLVM::LLVMStructType::getLiteral(ctx, {i32Ty, i64Ty});
-  });
 }
 
 std::size_t cudaq::opt::getDataSize(llvm::DataLayout &dataLayout, Type ty) {
   LLVMTypeConverter converter(ty.getContext());
   cudaq::opt::populateCCTypeConversions(&converter);
   auto llvmDialectTy = converter.convertType(ty);
+  // `measure_result` -> struct conversion for size calculation
+  converter.addConversion([](quake::MeasureType type) -> Type {
+    auto ctx = type.getContext();
+    auto i32Ty = IntegerType::get(ctx, 32);
+    auto i64Ty = IntegerType::get(ctx, 64);
+    return LLVM::LLVMStructType::getLiteral(ctx, {i32Ty, i64Ty});
+  });
   llvm::LLVMContext context;
   LLVM::TypeToLLVMIRTranslator translator(context);
   auto llvmTy = translator.translateType(llvmDialectTy);
@@ -103,6 +104,13 @@ std::size_t cudaq::opt::getDataOffset(llvm::DataLayout &dataLayout, Type ty,
                                       std::size_t off) {
   LLVMTypeConverter converter(ty.getContext());
   cudaq::opt::populateCCTypeConversions(&converter);
+  // `measure_result` -> struct conversion for size calculation
+  converter.addConversion([](quake::MeasureType type) -> Type {
+    auto ctx = type.getContext();
+    auto i32Ty = IntegerType::get(ctx, 32);
+    auto i64Ty = IntegerType::get(ctx, 64);
+    return LLVM::LLVMStructType::getLiteral(ctx, {i32Ty, i64Ty});
+  });
   auto llvmDialectTy = converter.convertType(ty);
   llvm::LLVMContext context;
   LLVM::TypeToLLVMIRTranslator translator(context);

lib/Optimizer/CodeGen/ConvertToQIR.cpp

@@ -206,9 +206,8 @@ void cudaq::opt::initializeTypeConversions(LLVMTypeConverter &typeConverter) {
     return LLVM::LLVMStructType::getLiteral(type.getContext(), mems,
                                             /*packed=*/false);
   });
-  typeConverter.addConversion([](quake::MeasureType type) {
-    return IntegerType::get(type.getContext(), 1);
-  });
+  typeConverter.addConversion(
+      [](quake::MeasureType type) { return getResultType(type.getContext()); });
   cudaq::opt::populateCCTypeConversions(&typeConverter);
 }
 

lib/Optimizer/CodeGen/QuakeToLLVM.cpp

@@ -284,8 +284,12 @@ class DiscriminateOpPattern
   LogicalResult
   matchAndRewrite(quake::DiscriminateOp discr, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto m = discr.getMeasurement();
-    rewriter.replaceOp(discr, m);
+    auto loc = discr.getLoc();
+    Value resultPtr = adaptor.getMeasurement();
+    auto i1Ty = rewriter.getI1Type();
+    auto i1PtrTy = LLVM::LLVMPointerType::get(i1Ty);
+    auto cast = rewriter.create<LLVM::BitcastOp>(loc, i1PtrTy, resultPtr);
+    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(discr, i1Ty, cast);
     return success();
   }
 };
@@ -1144,11 +1148,7 @@ class MeasureRewrite : public ConvertOpToLLVMPattern<OP> {
         loc, cudaq::opt::getResultType(context), symbolRef, ValueRange{args});
     if (regName)
       callOp->setAttr("registerName", regName);
-    auto i1Ty = rewriter.getI1Type();
-    auto i1PtrTy = LLVM::LLVMPointerType::get(i1Ty);
-    auto cast =
-        rewriter.create<LLVM::BitcastOp>(loc, i1PtrTy, callOp.getResult());
-    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(measure, i1Ty, cast);
+    rewriter.replaceOp(measure, callOp.getResult());
 
     return success();
   }

lib/Optimizer/CodeGen/ReturnToOutputLog.cpp

@@ -14,6 +14,7 @@
 #include "cudaq/Optimizer/CodeGen/QIRFunctionNames.h"
 #include "cudaq/Optimizer/Dialect/CC/CCOps.h"
 #include "cudaq/Optimizer/Dialect/CC/CCTypes.h"
+#include "cudaq/Optimizer/Dialect/Quake/QuakeTypes.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
@@ -160,6 +161,15 @@ class ReturnRewrite : public OpRewritePattern<cudaq::cc::LogOutputOp> {
               }
             }
         })
+        .Case([&](quake::MeasureType) {
+          std::string labelStr = "result";
+          if (prefix)
+            labelStr = prefix->str();
+          Value label = makeLabel(loc, rewriter, labelStr);
+          rewriter.create<func::CallOp>(loc, TypeRange{},
+                                        cudaq::opt::QIRRecordOutput,
+                                        ArrayRef<Value>{val, label});
+        })
         .Default([&](Type) {
           // If we reach here, we don't know how to handle this type.
           Value one = rewriter.create<arith::ConstantIntOp>(loc, 1, 64);
@@ -195,6 +205,8 @@ class ReturnRewrite : public OpRewritePattern<cudaq::cc::LogOutputOp> {
     if (auto arrTy = dyn_cast<cudaq::cc::StdvecType>(ty))
       return {std::string("array<") + translateType(arrTy.getElementType()) +
               std::string(" x ") + std::to_string(*vecSz) + std::string(">")};
+    if (isa<quake::MeasureType>(ty))
+      return {"Result"};
     return {"error"};
   }
 

runtime/common/RecordLogParser.cpp

@@ -240,7 +240,7 @@ cudaq::RecordLogParser::getDataHandler(const std::string &dataType) {
   static details::DataHandler<double> f64Handler(
       std::make_unique<details::FloatConverter<double>>());
   // Map data type to the corresponding handler
-  if (dataType == "measure_result")
+  if (dataType == "result")
     return measureResultHandler;
   if (dataType == "i1")
     return boolHandler;
@@ -283,7 +283,7 @@ void cudaq::RecordLogParser::processSingleRecord(const std::string &recValue,
   // For result type, we don't use the record label (register name) as the type
   // annotation.
   if (currentOutput == OutputType::RESULT)
-    label = "measure_result";
+    label = "result";
   if (label.empty()) {
     if (currentOutput == OutputType::BOOL)
       label = "i1";

runtime/cudaq/builder/kernel_builder.cpp

@@ -794,8 +794,7 @@ QuakeValue applyMeasure(ImplicitLocOpBuilder &builder, Value value,
     measureResult =
         builder.template create<QuakeMeasureOp>(measTy, value).getMeasOut();
 
-  Value bits = builder.create<quake::DiscriminateOp>(resTy, measureResult);
-  return QuakeValue(builder, bits);
+  return QuakeValue(builder, measureResult);
 }
 
 QuakeValue mx(ImplicitLocOpBuilder &builder, QuakeValue &qubitOrQvec,

runtime/cudaq/qis/measure_result.h

@@ -40,7 +40,11 @@ class measure_result {
   std::size_t getUniqueId() const { return uniqueId; }
 
   // Operator overloads for conversions and comparisons
+#ifdef CUDAQ_LIBRARY_MODE
   operator bool() const { return __nvqpp__MeasureResultBoolConversion(result); }
+#else
+  operator bool() const { return result == 1; }
+#endif
   explicit operator int() const { return result; }
   explicit operator double() const { return static_cast<double>(result); }