[L0 v2] Move locking to top-level functions

to make it consistent and to make sure all the locks
are taken in the same order.

Also, add a few missing TRACK_SCOPE_LATENCY calls

Author: igchor

source/adapters/level_zero/v2/context.cpp

@@ -134,7 +134,8 @@ ur_result_t urContextGetInfo(ur_context_handle_t hContext,
                              void *pContextInfo,
 
                              size_t *pPropSizeRet) {
-  std::shared_lock<ur_shared_mutex> Lock(hContext->Mutex);
+  // No locking needed here, we only read const members
+
   UrReturnHelper ReturnValue(propSize, pContextInfo, pPropSizeRet);
   switch (
       (uint32_t)contextInfoType) { // cast to avoid warnings on EXT enum values

source/adapters/level_zero/v2/event.cpp

@@ -93,8 +93,6 @@ static uint64_t adjustEndEventTimestamp(uint64_t adjustedStartTimestamp,
 }
 
 uint64_t ur_event_handle_t_::getEventEndTimestamp() {
-  std::scoped_lock<ur_shared_mutex> lock(this->Mutex);
-
   // If adjustedEventEndTimestamp on the event is non-zero it means it has
   // collected the result of the queue already. In that case it has been
   // adjusted and is ready for immediate return.
@@ -120,8 +118,6 @@ void ur_event_handle_t_::recordStartTimestamp() {
   UR_CALL_THROWS(ur::level_zero::urDeviceGetGlobalTimestamps(
       getDevice(), &deviceStartTimestamp, nullptr));
 
-  std::scoped_lock<ur_shared_mutex> lock(this->Mutex);
-
   adjustedEventStartTimestamp = deviceStartTimestamp;
 }
 
@@ -183,6 +179,8 @@ ur_result_t urEventGetProfilingInfo(
     size_t *pPropValueSizeRet ///< [out][optional] pointer to the actual size in
                               ///< bytes returned in propValue
 ) {
+  std::scoped_lock<ur_shared_mutex> lock(hEvent->Mutex);
+
   // The event must either have profiling enabled or be recording timestamps.
   bool isTimestampedEvent = hEvent->isTimestamped();
   if (!hEvent->isProfilingEnabled() && !isTimestampedEvent) {

source/adapters/level_zero/v2/kernel.cpp

@@ -182,8 +182,6 @@ ur_result_t ur_kernel_handle_t_::setArgValue(
     return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX;
   }
 
-  std::scoped_lock<ur_shared_mutex> guard(Mutex);
-
   for (auto &singleDeviceKernel : deviceKernels) {
     if (!singleDeviceKernel.has_value()) {
       continue;
@@ -217,8 +215,6 @@ ur_program_handle_t ur_kernel_handle_t_::getProgramHandle() const {
 
 ur_result_t ur_kernel_handle_t_::setExecInfo(ur_kernel_exec_info_t propName,
                                              const void *pPropValue) {
-  std::scoped_lock<ur_shared_mutex> Guard(Mutex);
-
   for (auto &kernel : deviceKernels) {
     if (!kernel.has_value())
       continue;
@@ -315,6 +311,8 @@ ur_result_t urKernelSetArgValue(
         *pArgValue ///< [in] argument value represented as matching arg type.
 ) {
   TRACK_SCOPE_LATENCY("ur_kernel_handle_t_::setArgValue");
+
+  std::scoped_lock<ur_shared_mutex> guard(hKernel->Mutex);
   return hKernel->setArgValue(argIndex, argSize, pProperties, pArgValue);
 }
 
@@ -327,6 +325,8 @@ ur_result_t urKernelSetArgPointer(
         *pArgValue ///< [in] argument value represented as matching arg type.
 ) {
   TRACK_SCOPE_LATENCY("ur_kernel_handle_t_::setArgPointer");
+
+  std::scoped_lock<ur_shared_mutex> guard(hKernel->Mutex);
   return hKernel->setArgPointer(argIndex, pProperties, pArgValue);
 }
 
@@ -336,6 +336,9 @@ urKernelSetArgMemObj(ur_kernel_handle_t hKernel, uint32_t argIndex,
                      ur_mem_handle_t hArgValue) {
   TRACK_SCOPE_LATENCY("ur_kernel_handle_t_::setArgMemObj");
 
+  std::scoped_lock<ur_shared_mutex, ur_shared_mutex> guard(hKernel->Mutex,
+                                                           hArgValue->Mutex);
+
   // TODO: support properties
   std::ignore = pProperties;
 
@@ -367,6 +370,8 @@ urKernelSetArgLocal(ur_kernel_handle_t hKernel, uint32_t argIndex,
                     const ur_kernel_arg_local_properties_t *pProperties) {
   TRACK_SCOPE_LATENCY("ur_kernel_handle_t_::setArgLocal");
 
+  std::scoped_lock<ur_shared_mutex> guard(hKernel->Mutex);
+
   std::ignore = pProperties;
 
   return hKernel->setArgValue(argIndex, argSize, nullptr, nullptr);
@@ -384,6 +389,8 @@ ur_result_t urKernelSetExecInfo(
   std::ignore = propSize;
   std::ignore = pProperties;
 
+  std::scoped_lock<ur_shared_mutex> guard(hKernel->Mutex);
+
   return hKernel->setExecInfo(propName, pPropValue);
 }
 
@@ -401,7 +408,7 @@ ur_result_t urKernelGetGroupInfo(
 ) {
   UrReturnHelper returnValue(paramValueSize, pParamValue, pParamValueSizeRet);
 
-  std::shared_lock<ur_shared_mutex> Guard(hKernel->Mutex);
+  // No locking needed here, we only read const members
   switch (paramName) {
   case UR_KERNEL_GROUP_INFO_GLOBAL_WORK_SIZE: {
     // TODO: To revisit after level_zero/issues/262 is resolved
@@ -483,7 +490,7 @@ ur_result_t urKernelGetSubGroupInfo(
 
   auto props = hKernel->getProperties(hDevice);
 
-  std::shared_lock<ur_shared_mutex> Guard(hKernel->Mutex);
+  // No locking needed here, we only read const members
   if (propName == UR_KERNEL_SUB_GROUP_INFO_MAX_SUB_GROUP_SIZE) {
     returnValue(uint32_t{props.maxSubgroupSize});
   } else if (propName == UR_KERNEL_SUB_GROUP_INFO_MAX_NUM_SUB_GROUPS) {

source/adapters/level_zero/v2/memory.cpp

@@ -125,6 +125,8 @@ static ur_result_t synchronousZeCopy(ur_context_handle_t hContext,
 ur_result_t
 ur_discrete_mem_handle_t::migrateBufferTo(ur_device_handle_t hDevice, void *src,
                                           size_t size) {
+  TRACK_SCOPE_LATENCY("ur_discrete_mem_handle_t::migrateBufferTo");
+
   auto Id = hDevice->Id.value();
 
   if (!deviceAllocations[Id]) {
@@ -163,9 +165,11 @@ ur_discrete_mem_handle_t::~ur_discrete_mem_handle_t() {
   }
 }
 
-void *ur_discrete_mem_handle_t::getDevicePtrUnlocked(
+void *ur_discrete_mem_handle_t::getDevicePtr(
     ur_device_handle_t hDevice, access_mode_t access, size_t offset,
     size_t size, std::function<void(void *src, void *dst, size_t)> migrate) {
+  TRACK_SCOPE_LATENCY("ur_discrete_mem_handle_t::getDevicePtr");
+
   std::ignore = access;
   std::ignore = size;
   std::ignore = migrate;
@@ -198,17 +202,10 @@ void *ur_discrete_mem_handle_t::getDevicePtrUnlocked(
          offset;
 }
 
-void *ur_discrete_mem_handle_t::getDevicePtr(
-    ur_device_handle_t hDevice, access_mode_t access, size_t offset,
-    size_t size, std::function<void(void *src, void *dst, size_t)> migrate) {
-  std::lock_guard lock(this->Mutex);
-  return getDevicePtrUnlocked(hDevice, access, offset, size, migrate);
-}
-
 void *ur_discrete_mem_handle_t::mapHostPtr(
     access_mode_t access, size_t offset, size_t size,
     std::function<void(void *src, void *dst, size_t)> migrate) {
-  std::lock_guard lock(this->Mutex);
+  TRACK_SCOPE_LATENCY("ur_discrete_mem_handle_t::mapHostPtr");
 
   // TODO: use async alloc?
 
@@ -231,7 +228,7 @@ void *ur_discrete_mem_handle_t::mapHostPtr(
 void ur_discrete_mem_handle_t::unmapHostPtr(
     void *pMappedPtr,
     std::function<void(void *src, void *dst, size_t)> migrate) {
-  std::lock_guard lock(this->Mutex);
+  TRACK_SCOPE_LATENCY("ur_discrete_mem_handle_t::unmapHostPtr");
 
   for (auto &hostAllocation : hostAllocations) {
     if (hostAllocation.ptr == pMappedPtr) {
@@ -241,9 +238,9 @@ void ur_discrete_mem_handle_t::unmapHostPtr(
                         deviceAllocations[activeAllocationDevice->Id.value()]) +
                     hostAllocation.offset;
       } else if (hostAllocation.access != access_mode_t::write_invalidate) {
-        devicePtr = ur_cast<char *>(getDevicePtrUnlocked(
-            hContext->getDevices()[0], access_mode_t::read_only,
-            hostAllocation.offset, hostAllocation.size, migrate));
+        devicePtr = ur_cast<char *>(
+            getDevicePtr(hContext->getDevices()[0], access_mode_t::read_only,
+                         hostAllocation.offset, hostAllocation.size, migrate));
       }
 
       if (devicePtr) {
@@ -328,7 +325,8 @@ ur_result_t urMemBufferCreateWithNativeHandle(
 ur_result_t urMemGetInfo(ur_mem_handle_t hMemory, ur_mem_info_t propName,
                          size_t propSize, void *pPropValue,
                          size_t *pPropSizeRet) {
-  std::shared_lock<ur_shared_mutex> Lock(hMemory->Mutex);
+  // No locking needed here, we only read const members
+
   UrReturnHelper returnValue(propSize, pPropValue, pPropSizeRet);
 
   switch (propName) {

source/adapters/level_zero/v2/memory.hpp

@@ -26,6 +26,7 @@ struct ur_mem_handle_t_ : _ur_object {
     write_invalidate
   };
 
+  // Following functions should always be called under the lock.
   virtual void *
   getDevicePtr(ur_device_handle_t, access_mode_t, size_t offset, size_t size,
                std::function<void(void *src, void *dst, size_t)> mecmpy) = 0;
@@ -109,10 +110,6 @@ struct ur_discrete_mem_handle_t : public ur_mem_handle_t_ {
                     std::function<void(void *src, void *dst, size_t)>) override;
 
 private:
-  void *getDevicePtrUnlocked(ur_device_handle_t, access_mode_t, size_t offset,
-                             size_t size,
-                             std::function<void(void *src, void *dst, size_t)>);
-
   // Vector of per-device allocations indexed by device->Id
   std::vector<void *> deviceAllocations;