[ASAN] Intercept urProgramLink to enable asan layer for openmp offloading

Signed-off-by: jinge90 <[email protected]>

Author: jinge90

source/loader/layers/sanitizer/asan_interceptor.cpp

@@ -537,14 +537,28 @@ ur_result_t SanitizerInterceptor::enqueueAllocInfo(
     }
 
     // Left red zone
-    UR_CALL(enqueueMemSetShadow(Context, Device, Queue, AllocInfo->AllocBegin,
-                                AllocInfo->UserBegin - AllocInfo->AllocBegin,
-                                ShadowByte, LastEvent, &LastEvent));
+    uptr LeftRZSize = AllocInfo->UserBegin - AllocInfo->AllocBegin;
+    if (LeftRZSize != 0) {
+        UR_CALL(enqueueMemSetShadow(Context, Device, Queue,
+                                    AllocInfo->AllocBegin, LeftRZSize,
+                                    ShadowByte, LastEvent, &LastEvent));
+    } else {
+        context.logger.debug(
+            "Try to create empty left red zone, Alloc Begin {}",
+            (void *)AllocInfo->AllocBegin);
+    }
 
     // Right red zone
-    UR_CALL(enqueueMemSetShadow(Context, Device, Queue, TailBegin,
-                                TailEnd - TailBegin, ShadowByte, LastEvent,
-                                &LastEvent));
+    uptr RightRZSize = TailEnd - TailBegin;
+    if (RightRZSize != 0) {
+        UR_CALL(enqueueMemSetShadow(Context, Device, Queue, TailBegin,
+                                    RightRZSize, ShadowByte, LastEvent,
+                                    &LastEvent));
+    } else {
+        context.logger.debug(
+            "Try to create empty right red zone, Alloc Begin {}",
+            (void *)TailBegin);
+    }
 
     return UR_RESULT_SUCCESS;
 }

source/loader/layers/sanitizer/ur_sanddi.cpp

@@ -178,6 +178,35 @@ __urdlllocal ur_result_t UR_APICALL urProgramBuild(
     return UR_RESULT_SUCCESS;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urProgramLink
+/// TODO: It is technically possible that user directly calls urProgramLinkExp
+/// instead of urProgramLink, intercept urProgramLinkExp if that happens.
+__urdlllocal ur_result_t UR_APICALL urProgramLink(
+    ur_context_handle_t hContext, ///< [in] handle of the context instance.
+    uint32_t count, ///< [in] number of program handles in `phPrograms`.
+    const ur_program_handle_t *
+        phPrograms, ///< [in][range(0, count)] pointer to array of program handles.
+    const char *
+        pOptions, ///< [in][optional] pointer to linker options null-terminated string.
+    ur_program_handle_t
+        *phProgram ///< [out] pointer to handle of program object created.
+) {
+    auto pfnProgramLink = context.urDdiTable.Program.pfnLink;
+
+    if (nullptr == pfnProgramLink) {
+        return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    }
+
+    context.logger.debug("==== urProgramLink");
+
+    UR_CALL(pfnProgramLink(hContext, count, phPrograms, pOptions, phProgram));
+
+    UR_CALL(context.interceptor->registerDeviceGlobals(hContext, *phProgram));
+
+    return UR_RESULT_SUCCESS;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Intercept function for urEnqueueKernelLaunch
 __urdlllocal ur_result_t UR_APICALL urEnqueueKernelLaunch(
@@ -422,6 +451,7 @@ __urdlllocal ur_result_t UR_APICALL urGetProgramProcAddrTable(
     }
 
     pDdiTable->pfnBuild = ur_sanitizer_layer::urProgramBuild;
+    pDdiTable->pfnLink = ur_sanitizer_layer::urProgramLink;
 
     return UR_RESULT_SUCCESS;
 }