Fix a potential race condition in quantum_platform thread Id to QPU Id map (#3291)

* Fix potential race condition in quantum_platform threadId to QPU Id map

Signed-off-by: Thien Nguyen <[email protected]>

* Code review: guard the qpu id assignement as well

Signed-off-by: Thien Nguyen <[email protected]>

---------

Signed-off-by: Thien Nguyen <[email protected]>

Author: 1tnguyen

runtime/cudaq/platform/quantum_platform.cpp

@@ -91,13 +91,16 @@ void quantum_platform::set_current_qpu(const std::size_t device_id) {
         " is not valid (greater than number of available QPUs: " +
         std::to_string(platformNumQPUs) + ").");
   }
-  platformCurrentQPU = device_id;
   auto tid = std::hash<std::thread::id>{}(std::this_thread::get_id());
-  auto iter = threadToQpuId.find(tid);
-  if (iter != threadToQpuId.end())
-    iter->second = device_id;
-  else
-    threadToQpuId.emplace(tid, device_id);
+  {
+    std::unique_lock lock(threadToQpuIdMutex);
+    platformCurrentQPU = device_id;
+    auto iter = threadToQpuId.find(tid);
+    if (iter != threadToQpuId.end())
+      iter->second = device_id;
+    else
+      threadToQpuId.emplace(tid, device_id);
+  }
 }
 
 std::size_t quantum_platform::get_current_qpu() { return platformCurrentQPU; }
@@ -150,9 +153,12 @@ void quantum_platform::launchVQE(const std::string kernelName,
   std::size_t qpu_id = 0;
 
   auto tid = std::hash<std::thread::id>{}(std::this_thread::get_id());
-  auto iter = threadToQpuId.find(tid);
-  if (iter != threadToQpuId.end())
-    qpu_id = iter->second;
+  {
+    std::shared_lock lock(threadToQpuIdMutex);
+    auto iter = threadToQpuId.find(tid);
+    if (iter != threadToQpuId.end())
+      qpu_id = iter->second;
+  }
 
   auto &qpu = platformQPUs[qpu_id];
   qpu->launchVQE(kernelName, kernelArgs, gradient, H, optimizer, n_params,
@@ -173,10 +179,12 @@ KernelThunkResultType quantum_platform::launchKernel(
   std::size_t qpu_id = 0;
 
   auto tid = std::hash<std::thread::id>{}(std::this_thread::get_id());
-  auto iter = threadToQpuId.find(tid);
-  if (iter != threadToQpuId.end())
-    qpu_id = iter->second;
-
+  {
+    std::shared_lock lock(threadToQpuIdMutex);
+    auto iter = threadToQpuId.find(tid);
+    if (iter != threadToQpuId.end())
+      qpu_id = iter->second;
+  }
   auto &qpu = platformQPUs[qpu_id];
   return qpu->launchKernel(kernelName, kernelFunc, args, voidStarSize,
                            resultOffset, rawArgs);
@@ -187,10 +195,12 @@ void quantum_platform::launchKernel(std::string kernelName,
   std::size_t qpu_id = 0;
 
   auto tid = std::hash<std::thread::id>{}(std::this_thread::get_id());
-  auto iter = threadToQpuId.find(tid);
-  if (iter != threadToQpuId.end())
-    qpu_id = iter->second;
-
+  {
+    std::shared_lock lock(threadToQpuIdMutex);
+    auto iter = threadToQpuId.find(tid);
+    if (iter != threadToQpuId.end())
+      qpu_id = iter->second;
+  }
   auto &qpu = platformQPUs[qpu_id];
   qpu->launchKernel(kernelName, rawArgs);
 }
@@ -201,10 +211,12 @@ void quantum_platform::launchSerializedCodeExecution(
   std::size_t qpu_id = 0;
 
   auto tid = std::hash<std::thread::id>{}(std::this_thread::get_id());
-  auto iter = threadToQpuId.find(tid);
-  if (iter != threadToQpuId.end())
-    qpu_id = iter->second;
-
+  {
+    std::shared_lock lock(threadToQpuIdMutex);
+    auto iter = threadToQpuId.find(tid);
+    if (iter != threadToQpuId.end())
+      qpu_id = iter->second;
+  }
   auto &qpu = platformQPUs[qpu_id];
   qpu->launchSerializedCodeExecution(name, serializeCodeExecutionObject);
 }

runtime/cudaq/platform/quantum_platform.h

@@ -204,6 +204,9 @@ class quantum_platform {
   /// that it is running in a multi-QPU context.
   std::unordered_map<std::size_t, std::size_t> threadToQpuId;
 
+  /// @brief Mutex to protect access to the thread-QPU map.
+  std::shared_mutex threadToQpuIdMutex;
+
   /// Optional number of shots.
   std::optional<int> platformNumShots;