wip

Author: AllanZyne

source/loader/layers/sanitizer/asan/asan_shadow.cpp

@@ -223,7 +223,7 @@ ur_result_t ShadowMemoryGPU::EnqueuePoisonShadow(ur_queue_handle_t Queue,
                                            (void *)(MappedPtr + PageSize - 1));
 
                 // Initialize to zero
-                URes = EnqueueUSMSet(Queue, (void *)MappedPtr, 0,
+                URes = EnqueueUSMBlockingSet(Queue, (void *)MappedPtr, 0,
                                              PageSize);
                 if (URes != UR_RESULT_SUCCESS) {
                     getContext()->logger.error("EnqueueUSMBlockingSet(): {}",
@@ -236,7 +236,7 @@ ur_result_t ShadowMemoryGPU::EnqueuePoisonShadow(ur_queue_handle_t Queue,
         }
     }
 
-    auto URes = EnqueueUSMSet(Queue, (void *)ShadowBegin, Value,
+    auto URes = EnqueueUSMBlockingSet(Queue, (void *)ShadowBegin, Value,
                                       ShadowEnd - ShadowBegin + 1);
     getContext()->logger.debug(
         "EnqueuePoisonShadow (addr={}, count={}, value={}): {}",
@@ -272,7 +272,7 @@ ur_result_t ShadowMemoryGPU::AllocLocalShadow(ur_queue_handle_t Queue,
             (void **)&LocalShadowOffset));
 
         // Initialize shadow memory
-        ur_result_t URes = EnqueueUSMSet(
+        ur_result_t URes = EnqueueUSMBlockingSet(
             Queue, (void *)LocalShadowOffset, 0, RequiredShadowSize);
         if (URes != UR_RESULT_SUCCESS) {
             UR_CALL(getContext()->urDdiTable.USM.pfnFree(
@@ -312,7 +312,7 @@ ur_result_t ShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue,
             (void **)&PrivateShadowOffset));
 
         // Initialize shadow memory
-        ur_result_t URes = EnqueueUSMSet(
+        ur_result_t URes = EnqueueUSMBlockingSet(
             Queue, (void *)PrivateShadowOffset, 0, RequiredShadowSize);
         if (URes != UR_RESULT_SUCCESS) {
             UR_CALL(getContext()->urDdiTable.USM.pfnFree(

source/loader/layers/sanitizer/msan/msan_interceptor.cpp

@@ -82,6 +82,19 @@ ur_result_t MsanInterceptor::allocateMemory(ur_context_handle_t Context,
     return UR_RESULT_SUCCESS;
 }
 
+ur_result_t MsanInterceptor::releaseMemory(ur_context_handle_t Context,
+                                           void *Ptr) {
+    auto Addr = reinterpret_cast<uptr>(Ptr);
+    auto AddrInfoItOp = findAllocInfoByAddress(Addr);
+
+    if (AddrInfoItOp) {
+        std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
+        m_AllocationMap.erase(*AddrInfoItOp);
+    }
+
+    return getContext()->urDdiTable.USM.pfnFree(Context, Ptr);
+}
+
 ur_result_t MsanInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
                                              ur_queue_handle_t Queue,
                                              USMLaunchInfo &LaunchInfo) {

source/loader/layers/sanitizer/msan/msan_shadow.cpp

@@ -249,7 +249,7 @@ ur_result_t MsanShadowMemoryGPU::EnqueueMappingShadow(
                                        (void *)(MappedPtr + PageSize - 1));
 
             // Initialize to zero
-            URes = EnqueueUSMSet(Queue, (void *)MappedPtr, 0, PageSize,
+            URes = EnqueueUSMBlockingSet(Queue, (void *)MappedPtr, 0, PageSize,
                                  EventWaitList.size(), EventWaitList.data(),
                                  OutEvent);
             if (URes != UR_RESULT_SUCCESS) {
@@ -296,7 +296,7 @@ ur_result_t MsanShadowMemoryGPU::EnqueuePoisonShadow(
     const uptr ShadowEnd = MemToShadow(Ptr + Size - 1);
     assert(ShadowBegin <= ShadowEnd);
 
-    auto Result = EnqueueUSMSet(Queue, (void *)ShadowBegin, Value,
+    auto Result = EnqueueUSMBlockingSet(Queue, (void *)ShadowBegin, Value,
                                 ShadowEnd - ShadowBegin + 1, Events.size(),
                                 Events.data(), OutEvent);
 

source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.cpp

@@ -243,7 +243,7 @@ size_t GetVirtualMemGranularity(ur_context_handle_t Context,
     return Size;
 }
 
-ur_result_t EnqueueUSMSet(ur_queue_handle_t Queue, void *Ptr, char Value,
+ur_result_t EnqueueUSMBlockingSet(ur_queue_handle_t Queue, void *Ptr, char Value,
                           size_t Size, uint32_t NumEvents,
                           const ur_event_handle_t *EventWaitList,
                           ur_event_handle_t *OutEvent) {

source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.hpp

@@ -59,7 +59,7 @@ size_t GetVirtualMemGranularity(ur_context_handle_t Context,
                                 ur_device_handle_t Device);
 
 ur_result_t
-EnqueueUSMSet(ur_queue_handle_t Queue, void *Ptr, char Value,
+EnqueueUSMBlockingSet(ur_queue_handle_t Queue, void *Ptr, char Value,
                       size_t Size, uint32_t NumEvents = 0,
                       const ur_event_handle_t *EventWaitList = nullptr,
                       ur_event_handle_t *OutEvent = nullptr);