[DevSAN] Correctly intercept UR API urEnqueueKernelLaunchWithArgsExp

unified-runtime/source/loader/layers/sanitizer/asan/asan_ddi.cpp

@@ -1739,6 +1739,11 @@ __urdlllocal ur_result_t UR_APICALL urEnqueueKernelLaunchWithArgsExp(
   UR_LOG_L(getContext()->logger, DEBUG,
            "==== urEnqueueKernelLaunchWithArgsExp");
 
+  auto &KernelInfo = getAsanInterceptor()->getOrCreateKernelInfo(hKernel);
+  KernelInfo.ArgProps.resize(numArgs);
+  std::memcpy(KernelInfo.ArgProps.data(), pArgs,
+              numArgs * sizeof(ur_exp_kernel_arg_properties_t));
+
   // We need to set all the args now rather than letting LaunchWithArgs handle
   // them. This is because some implementations of
   // urKernelGetSuggestedLocalWorkSize, which is used in preLaunchKernel, rely
@@ -1749,6 +1754,8 @@ __urdlllocal ur_result_t UR_APICALL urEnqueueKernelLaunchWithArgsExp(
       UR_CALL(ur_sanitizer_layer::asan::urKernelSetArgLocal(
           hKernel, pArgs[ArgPropIndex].index, pArgs[ArgPropIndex].size,
           nullptr));
+      KernelInfo.ArgProps[ArgPropIndex].size =
+          KernelInfo.LocalArgs[ArgPropIndex].SizeWithRedZone;
       break;
     }
     case UR_EXP_KERNEL_ARG_TYPE_POINTER: {
@@ -1770,6 +1777,15 @@ __urdlllocal ur_result_t UR_APICALL urEnqueueKernelLaunchWithArgsExp(
       UR_CALL(ur_sanitizer_layer::asan::urKernelSetArgMemObj(
           hKernel, pArgs[ArgPropIndex].index, &Properties,
           pArgs[ArgPropIndex].value.memObjTuple.hMem));
+      if (std::shared_ptr<MemBuffer> MemBuffer =
+              getAsanInterceptor()->getMemBuffer(
+                  pArgs[ArgPropIndex].value.memObjTuple.hMem)) {
+        char *Handle = nullptr;
+        UR_CALL(MemBuffer->getHandle(GetDevice(hQueue), Handle));
+        KernelInfo.ArgProps[ArgPropIndex].type =
+            ur_exp_kernel_arg_type_t::UR_EXP_KERNEL_ARG_TYPE_POINTER;
+        KernelInfo.ArgProps[ArgPropIndex].value.pointer = Handle;
+      }
       break;
     }
     case UR_EXP_KERNEL_ARG_TYPE_SAMPLER: {
@@ -1791,18 +1807,10 @@ __urdlllocal ur_result_t UR_APICALL urEnqueueKernelLaunchWithArgsExp(
 
   UR_CALL(getAsanInterceptor()->preLaunchKernel(hKernel, hQueue, LaunchInfo));
 
-  /*
-    // TODO: revert to the correct call to pfnKernelLaunchWithArgsExp():
-    UR_CALL(getContext()->urDdiTable.EnqueueExp.pfnKernelLaunchWithArgsExp(
-        hQueue, hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize,
-        LaunchInfo.LocalWorkSize.data(), numArgs, pArgs,
-    launchPropList, numEventsInWaitList,
-    phEventWaitList, phEvent));
-  */
-
-  UR_CALL(getContext()->urDdiTable.Enqueue.pfnKernelLaunch(
+  UR_CALL(getContext()->urDdiTable.EnqueueExp.pfnKernelLaunchWithArgsExp(
       hQueue, hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize,
-      LaunchInfo.LocalWorkSize.data(), launchPropList, numEventsInWaitList,
+      LaunchInfo.LocalWorkSize.data(), KernelInfo.ArgProps.size(),
+      KernelInfo.ArgProps.data(), launchPropList, numEventsInWaitList,
       phEventWaitList, phEvent));
 
   UR_CALL(getAsanInterceptor()->postLaunchKernel(hKernel, hQueue, LaunchInfo));

unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp

@@ -784,6 +784,13 @@ ur_result_t AsanInterceptor::prepareLaunch(
 
     ur_result_t URes = getContext()->urDdiTable.Kernel.pfnSetArgPointer(
         Kernel, ArgNums - 1, nullptr, LaunchInfo.Data.getDevicePtr());
+    KernelInfo.ArgProps.push_back(ur_exp_kernel_arg_properties_t{
+        UR_STRUCTURE_TYPE_EXP_KERNEL_ARG_PROPERTIES,
+        nullptr,
+        UR_EXP_KERNEL_ARG_TYPE_POINTER,
+        ArgNums - 1,
+        sizeof(void *),
+        {LaunchInfo.Data.getDevicePtr()}});
     if (URes != UR_RESULT_SUCCESS) {
       UR_LOG_L(getContext()->logger, ERR, "Failed to set launch info: {}",
                URes);

unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp

@@ -106,6 +106,8 @@ struct KernelInfo {
   // Need preserve the order of local arguments
   std::map<uint32_t, LocalArgsInfo> LocalArgs;
 
+  std::vector<ur_exp_kernel_arg_properties_t> ArgProps;
+
   explicit KernelInfo(ur_kernel_handle_t Kernel) : Handle(Kernel) {
     [[maybe_unused]] auto Result =
         getContext()->urDdiTable.Kernel.pfnRetain(Kernel);

unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp

@@ -500,10 +500,8 @@ ur_result_t urEnqueueKernelLaunch(
 
   UR_LOG_L(getContext()->logger, DEBUG, "==== urEnqueueKernelLaunch");
 
-  USMLaunchInfo LaunchInfo(GetContext(hQueue), GetDevice(hQueue),
-                           pGlobalWorkSize, pLocalWorkSize, pGlobalWorkOffset,
-                           workDim);
-  UR_CALL(LaunchInfo.initialize());
+  LaunchInfo LaunchInfo(GetContext(hQueue), GetDevice(hQueue), pGlobalWorkSize,
+                        pLocalWorkSize, pGlobalWorkOffset, workDim);
 
   UR_CALL(getMsanInterceptor()->preLaunchKernel(hKernel, hQueue, LaunchInfo));
 
@@ -1883,6 +1881,11 @@ ur_result_t urEnqueueKernelLaunchWithArgsExp(
   UR_LOG_L(getContext()->logger, DEBUG,
            "==== urEnqueueKernelLaunchWithArgsExp");
 
+  auto &KernelInfo = getMsanInterceptor()->getOrCreateKernelInfo(hKernel);
+  KernelInfo.ArgProps.resize(numArgs);
+  std::memcpy(KernelInfo.ArgProps.data(), pArgs,
+              numArgs * sizeof(ur_exp_kernel_arg_properties_t));
+
   // We need to set all the args now rather than letting LaunchWithArgs handle
   // them. This is because some implementations of
   // urKernelGetSuggestedLocalWorkSize, which is used in preLaunchKernel, rely
@@ -1916,6 +1919,15 @@ ur_result_t urEnqueueKernelLaunchWithArgsExp(
       UR_CALL(ur_sanitizer_layer::msan::urKernelSetArgMemObj(
           hKernel, pArgs[ArgPropIndex].index, &Properties,
           pArgs[ArgPropIndex].value.memObjTuple.hMem));
+      if (std::shared_ptr<MemBuffer> MemBuffer =
+              getMsanInterceptor()->getMemBuffer(
+                  pArgs[ArgPropIndex].value.memObjTuple.hMem)) {
+        char *Handle = nullptr;
+        UR_CALL(MemBuffer->getHandle(GetDevice(hQueue), Handle));
+        KernelInfo.ArgProps[ArgPropIndex].type =
+            ur_exp_kernel_arg_type_t::UR_EXP_KERNEL_ARG_TYPE_POINTER;
+        KernelInfo.ArgProps[ArgPropIndex].value.pointer = Handle;
+      }
       break;
     }
     case UR_EXP_KERNEL_ARG_TYPE_SAMPLER: {
@@ -1931,26 +1943,15 @@ ur_result_t urEnqueueKernelLaunchWithArgsExp(
     }
   }
 
-  USMLaunchInfo LaunchInfo(GetContext(hQueue), GetDevice(hQueue),
-                           pGlobalWorkSize, pLocalWorkSize, pGlobalWorkOffset,
-                           workDim);
-  UR_CALL(LaunchInfo.initialize());
+  LaunchInfo LaunchInfo(GetContext(hQueue), GetDevice(hQueue), pGlobalWorkSize,
+                        pLocalWorkSize, pGlobalWorkOffset, workDim);
 
   UR_CALL(getMsanInterceptor()->preLaunchKernel(hKernel, hQueue, LaunchInfo));
 
-  /*
-    // TODO: revert to the correct call to pfnKernelLaunchWithArgsExp():
-    UR_CALL(getContext()->urDdiTable.EnqueueExp.pfnKernelLaunchWithArgsExp(
-        hQueue, hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize,
-        LaunchInfo.LocalWorkSize.data(), numArgs, pArgs,
-    launchPropList, numEventsInWaitList,
-    phEventWaitList, phEvent));
-  */
-
-  UR_CALL(getContext()->urDdiTable.Enqueue.pfnKernelLaunch(
+  UR_CALL(getContext()->urDdiTable.EnqueueExp.pfnKernelLaunchWithArgsExp(
       hQueue, hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize,
-      LaunchInfo.LocalWorkSize.data(), launchPropList, numEventsInWaitList,
-      phEventWaitList, phEvent));
+      LaunchInfo.LocalWorkSize.data(), numArgs, KernelInfo.ArgProps.data(),
+      launchPropList, numEventsInWaitList, phEventWaitList, phEvent));
 
   UR_CALL(getMsanInterceptor()->postLaunchKernel(hKernel, hQueue, LaunchInfo));
 

unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.cpp

@@ -137,7 +137,7 @@ ur_result_t MsanInterceptor::releaseMemory(ur_context_handle_t Context,
 
 ur_result_t MsanInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
                                              ur_queue_handle_t Queue,
-                                             USMLaunchInfo &LaunchInfo) {
+                                             LaunchInfo &LaunchInfo) {
   auto Context = GetContext(Queue);
   auto Device = GetDevice(Queue);
   auto ContextInfo = getContextInfo(Context);
@@ -156,7 +156,7 @@ ur_result_t MsanInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
 
 ur_result_t MsanInterceptor::postLaunchKernel(ur_kernel_handle_t Kernel,
                                               ur_queue_handle_t Queue,
-                                              USMLaunchInfo &LaunchInfo) {
+                                              LaunchInfo &LaunchInfo) {
   // FIXME: We must use block operation here, until we support
   // urEventSetCallback
   auto Result = getContext()->urDdiTable.Queue.pfnFinish(Queue);
@@ -458,7 +458,7 @@ MsanInterceptor::getMemBuffer(ur_mem_handle_t MemHandle) {
 
 ur_result_t MsanInterceptor::prepareLaunch(
     std::shared_ptr<DeviceInfo> &DeviceInfo, ur_queue_handle_t Queue,
-    ur_kernel_handle_t Kernel, USMLaunchInfo &LaunchInfo) {
+    ur_kernel_handle_t Kernel, LaunchInfo &LaunchInfo) {
   auto Program = GetProgram(Kernel);
 
   // Set membuffer arguments
@@ -653,12 +653,6 @@ ContextInfo::~ContextInfo() {
   assert(Result == UR_RESULT_SUCCESS);
 }
 
-ur_result_t USMLaunchInfo::initialize() {
-  UR_CALL(getContext()->urDdiTable.Context.pfnRetain(Context));
-  UR_CALL(getContext()->urDdiTable.Device.pfnRetain(Device));
-  return UR_RESULT_SUCCESS;
-}
-
 MsanRuntimeDataWrapper::~MsanRuntimeDataWrapper() {
   if (Host.CleanShadow) {
     [[maybe_unused]] auto Result =
@@ -672,14 +666,6 @@ MsanRuntimeDataWrapper::~MsanRuntimeDataWrapper() {
   }
 }
 
-USMLaunchInfo::~USMLaunchInfo() {
-  [[maybe_unused]] ur_result_t Result;
-  Result = getContext()->urDdiTable.Context.pfnRelease(Context);
-  assert(Result == UR_RESULT_SUCCESS);
-  Result = getContext()->urDdiTable.Device.pfnRelease(Device);
-  assert(Result == UR_RESULT_SUCCESS);
-}
-
 } // namespace msan
 
 using namespace msan;

unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.hpp

@@ -93,6 +93,8 @@ struct KernelInfo {
   // Need preserve the order of local arguments
   std::map<uint32_t, MsanLocalArgsInfo> LocalArgs;
 
+  std::vector<ur_exp_kernel_arg_properties_t> ArgProps;
+
   explicit KernelInfo(ur_kernel_handle_t Kernel) : Handle(Kernel) {
     [[maybe_unused]] auto Result =
         getContext()->urDdiTable.Kernel.pfnRetain(Kernel);
@@ -219,7 +221,7 @@ struct MsanRuntimeDataWrapper {
   }
 };
 
-struct USMLaunchInfo {
+struct LaunchInfo {
   MsanRuntimeDataWrapper Data;
 
   ur_context_handle_t Context = nullptr;
@@ -229,11 +231,16 @@ struct USMLaunchInfo {
   std::vector<size_t> LocalWorkSize;
   uint32_t WorkDim = 0;
 
-  USMLaunchInfo(ur_context_handle_t Context, ur_device_handle_t Device,
-                const size_t *GlobalWorkSize, const size_t *LocalWorkSize,
-                const size_t *GlobalWorkOffset, uint32_t WorkDim)
+  LaunchInfo(ur_context_handle_t Context, ur_device_handle_t Device,
+             const size_t *GlobalWorkSize, const size_t *LocalWorkSize,
+             const size_t *GlobalWorkOffset, uint32_t WorkDim)
       : Data(Context, Device), Context(Context), Device(Device),
         GlobalWorkSize(GlobalWorkSize), WorkDim(WorkDim) {
+    [[maybe_unused]] auto Result =
+        getContext()->urDdiTable.Context.pfnRetain(Context);
+    assert(Result == UR_RESULT_SUCCESS);
+    Result = getContext()->urDdiTable.Device.pfnRetain(Device);
+    assert(Result == UR_RESULT_SUCCESS);
     if (LocalWorkSize) {
       this->LocalWorkSize =
           std::vector<size_t>(LocalWorkSize, LocalWorkSize + WorkDim);
@@ -247,9 +254,17 @@ struct USMLaunchInfo {
       this->GlobalWorkOffset = std::vector<size_t>(WorkDim, 0);
     }
   }
-  ~USMLaunchInfo();
+  ~LaunchInfo() {
+    [[maybe_unused]] ur_result_t Result;
+    Result = getContext()->urDdiTable.Context.pfnRelease(Context);
+    assert(Result == UR_RESULT_SUCCESS);
+    Result = getContext()->urDdiTable.Device.pfnRelease(Device);
+    assert(Result == UR_RESULT_SUCCESS);
+  }
+
+  LaunchInfo(const LaunchInfo &) = delete;
 
-  ur_result_t initialize();
+  LaunchInfo &operator=(const LaunchInfo &) = delete;
 };
 
 struct DeviceGlobalInfo {
@@ -280,11 +295,9 @@ class MsanInterceptor {
   ur_result_t unregisterProgram(ur_program_handle_t Program);
 
   ur_result_t preLaunchKernel(ur_kernel_handle_t Kernel,
-                              ur_queue_handle_t Queue,
-                              USMLaunchInfo &LaunchInfo);
+                              ur_queue_handle_t Queue, LaunchInfo &LaunchInfo);
   ur_result_t postLaunchKernel(ur_kernel_handle_t Kernel,
-                               ur_queue_handle_t Queue,
-                               USMLaunchInfo &LaunchInfo);
+                               ur_queue_handle_t Queue, LaunchInfo &LaunchInfo);
 
   ur_result_t insertContext(ur_context_handle_t Context,
                             std::shared_ptr<ContextInfo> &CI);
@@ -345,7 +358,7 @@ class MsanInterceptor {
   /// Initialize Global Variables & Kernel Name at first Launch
   ur_result_t prepareLaunch(std::shared_ptr<DeviceInfo> &DeviceInfo,
                             ur_queue_handle_t Queue, ur_kernel_handle_t Kernel,
-                            USMLaunchInfo &LaunchInfo);
+                            LaunchInfo &LaunchInfo);
 
   ur_result_t allocShadowMemory(ur_context_handle_t Context,
                                 std::shared_ptr<DeviceInfo> &DeviceInfo);

unified-runtime/source/loader/layers/sanitizer/tsan/tsan_ddi.cpp

@@ -1407,6 +1407,11 @@ ur_result_t urEnqueueKernelLaunchWithArgsExp(
   UR_LOG_L(getContext()->logger, DEBUG,
            "==== urEnqueueKernelLaunchWithArgsExp");
 
+  auto &KernelInfo = getTsanInterceptor()->getKernelInfo(hKernel);
+  KernelInfo.ArgProps.resize(numArgs);
+  std::memcpy(KernelInfo.ArgProps.data(), pArgs,
+              numArgs * sizeof(ur_exp_kernel_arg_properties_t));
+
   // We need to set all the args now rather than letting LaunchWithArgs handle
   // them. This is because some implementations of
   // urKernelGetSuggestedLocalWorkSize, which is used in preLaunchKernel, rely
@@ -1440,6 +1445,15 @@ ur_result_t urEnqueueKernelLaunchWithArgsExp(
       UR_CALL(ur_sanitizer_layer::tsan::urKernelSetArgMemObj(
           hKernel, pArgs[ArgPropIndex].index, &Properties,
           pArgs[ArgPropIndex].value.memObjTuple.hMem));
+      if (std::shared_ptr<MemBuffer> MemBuffer =
+              getTsanInterceptor()->getMemBuffer(
+                  pArgs[ArgPropIndex].value.memObjTuple.hMem)) {
+        char *Handle = nullptr;
+        UR_CALL(MemBuffer->getHandle(GetDevice(hQueue), Handle));
+        KernelInfo.ArgProps[ArgPropIndex].type =
+            ur_exp_kernel_arg_type_t::UR_EXP_KERNEL_ARG_TYPE_POINTER;
+        KernelInfo.ArgProps[ArgPropIndex].value.pointer = Handle;
+      }
       break;
     }
     case UR_EXP_KERNEL_ARG_TYPE_SAMPLER: {
@@ -1460,18 +1474,10 @@ ur_result_t urEnqueueKernelLaunchWithArgsExp(
 
   UR_CALL(getTsanInterceptor()->preLaunchKernel(hKernel, hQueue, LaunchInfo));
 
-  /*
-    // TODO: revert to the correct call to pfnKernelLaunchWithArgsExp():
-    UR_CALL(getContext()->urDdiTable.EnqueueExp.pfnKernelLaunchWithArgsExp(
-        hQueue, hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize,
-        pLocalWorkSize, numArgs, pArgs,
-    launchPropList, numEventsInWaitList, phEventWaitList, phEvent));
-  */
-
-  UR_CALL(getContext()->urDdiTable.Enqueue.pfnKernelLaunch(
+  UR_CALL(getContext()->urDdiTable.EnqueueExp.pfnKernelLaunchWithArgsExp(
       hQueue, hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize,
-      pLocalWorkSize, launchPropList, numEventsInWaitList, phEventWaitList,
-      phEvent));
+      LaunchInfo.LocalWorkSize.data(), numArgs, KernelInfo.ArgProps.data(),
+      launchPropList, numEventsInWaitList, phEventWaitList, phEvent));
 
   UR_CALL(getTsanInterceptor()->postLaunchKernel(hKernel, hQueue, LaunchInfo));
 

unified-runtime/source/loader/layers/sanitizer/tsan/tsan_interceptor.hpp

@@ -127,6 +127,8 @@ struct KernelInfo {
   // Need preserve the order of local arguments
   std::map<uint32_t, TsanLocalArgsInfo> LocalArgs;
 
+  std::vector<ur_exp_kernel_arg_properties_t> ArgProps;
+
   KernelInfo() = default;
 
   explicit KernelInfo(ur_kernel_handle_t Kernel) : Handle(Kernel) {