From 94e10120f46e1669036081d7f45f3a0a77318170 Mon Sep 17 00:00:00 2001
From: Ben Howe <bhowe@nvidia.com>
Date: Fri, 16 Aug 2024 13:13:39 +0000
Subject: [PATCH] Implement support for __quantum__qis__read_result__body

Failed Tests (8):
  CUDAQ-Target :: execution/conditional_sample-cpp17.cpp
  CUDAQ-Target :: execution/conditional_sample.cpp
  CUDAQ-Target :: execution/graph_coloring.cpp
  CUDAQ-Target :: execution/qir_if_base.cpp
  CUDAQ-Target :: execution/qir_op1_after_measure.cpp
  CUDAQ-Target :: execution/qir_op2_after_measure.cpp
  CUDAQ-Target :: execution/qir_string_labels.cpp
  CUDAQ-Target :: execution/test-6.cpp
---
 .../Optimizer/CodeGen/QIRFunctionNames.h      |  4 ++++
 runtime/common/RuntimeMLIRCommonImpl.h        |  7 ++++++
 runtime/nvqir/NVQIR.cpp                       | 23 +++++++++++++++----
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h b/include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h
index 06383275e8..28e3400ab8 100644
--- a/include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h
+++ b/include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h
@@ -84,6 +84,10 @@ static constexpr const char QIRsetDynamicQubitManagement[] =
 static constexpr const char QIRRecordOutput[] =
     "__quantum__rt__result_record_output";
 
+/// Custom NVQIR method to cleanup result maps in between consecutive programs.
+static constexpr const char QIRClearResultMaps[] =
+    "__quantum__rt__clear_result_maps";
+
 inline mlir::Type getQuantumTypeByName(mlir::StringRef type,
                                        mlir::MLIRContext *context) {
   return mlir::LLVM::LLVMStructType::getOpaque(type, context);
diff --git a/runtime/common/RuntimeMLIRCommonImpl.h b/runtime/common/RuntimeMLIRCommonImpl.h
index f44997175b..a95c357835 100644
--- a/runtime/common/RuntimeMLIRCommonImpl.h
+++ b/runtime/common/RuntimeMLIRCommonImpl.h
@@ -578,6 +578,10 @@ void insertSetupAndCleanupOperations(mlir::Operation *module) {
       cudaq::opt::factory::createLLVMFunctionSymbol(
           cudaq::opt::QIRsetDynamicQubitManagement, {voidTy}, {boolTy},
           dyn_cast<mlir::ModuleOp>(module));
+  mlir::FlatSymbolRefAttr clearResultMapsSymbol =
+      cudaq::opt::factory::createLLVMFunctionSymbol(
+          cudaq::opt::QIRClearResultMaps, {voidTy}, {},
+          dyn_cast<mlir::ModuleOp>(module));
 
   // Iterate through all operations in the ModuleOp
   mlir::SmallVector<mlir::LLVM::LLVMFuncOp> funcs;
@@ -625,6 +629,9 @@ void insertSetupAndCleanupOperations(mlir::Operation *module) {
     builder.create<mlir::LLVM::CallOp>(loc, mlir::TypeRange{voidTy},
                                        setDynamicSymbol,
                                        mlir::ValueRange{origMode.getResult()});
+    builder.create<mlir::LLVM::CallOp>(loc, mlir::TypeRange{voidTy},
+                                       clearResultMapsSymbol,
+                                       mlir::ValueRange{});
   }
 }
 
diff --git a/runtime/nvqir/NVQIR.cpp b/runtime/nvqir/NVQIR.cpp
index b1d77ad86d..a5be67af3a 100644
--- a/runtime/nvqir/NVQIR.cpp
+++ b/runtime/nvqir/NVQIR.cpp
@@ -37,8 +37,13 @@ thread_local nvqir::CircuitSimulator *simulator;
 inline static constexpr std::string_view GetCircuitSimulatorSymbol =
     "getCircuitSimulator";
 
+// The following maps are used to map Qubits to Results, and Results to boolean
+// values. The pointer values may be integers if they are referring to Base
+// Profile or Adaptive Profile QIR programs, so it is generally not safe to
+// dereference them.
 static thread_local std::map<Qubit *, Result *> measQB2Res;
 static thread_local std::map<Result *, Qubit *> measRes2QB;
+static thread_local std::map<Result *, Result> measRes2Val;
 
 /// @brief Provide a holder for externally created
 /// CircuitSimulator pointers (like from Python) that
@@ -560,15 +565,15 @@ Result *__quantum__qis__mz__body(Qubit *q, Result *r) {
   auto qI = qubitToSizeT(q);
   ScopedTraceWithContext("NVQIR::mz", qI);
   auto b = nvqir::getCircuitSimulatorInternal()->mz(qI, "");
+  measRes2Val[r] = b;
   return b ? ResultOne : ResultZero;
 }
 
 bool __quantum__qis__read_result__body(Result *result) {
-  // TODO: implement post-measurement result retrieval. This is not needed for
-  // typical simulator operation (other than to have it defined), but it may be
-  // useful in the future.
-  // https://github.com/NVIDIA/cuda-quantum/issues/758
-  ScopedTraceWithContext("NVQIR::read_result (stubbed out)");
+  ScopedTraceWithContext("NVQIR::read_result");
+  auto iter = measRes2Val.find(result);
+  if (iter != measRes2Val.end())
+    return iter->second;
   return ResultZeroVal;
 }
 
@@ -862,6 +867,14 @@ void __quantum__qis__exp__body(Array *paulis, double angle, Array *qubits) {
   }
 }
 
+/// @brief Cleanup an result maps at the end of a QIR program to avoid leaking
+/// results into the next program.
+void __quantum__rt__clear_result_maps() {
+  measQB2Res.clear();
+  measRes2QB.clear();
+  measRes2Val.clear();
+}
+
 /// @brief Utility function used by Quake->QIR to pack a single Qubit pointer
 /// into an Array pointer.
 Array *packSingleQubitInArray(Qubit *q) {