Candidate for the v0.11.4 release tag

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

@@ -721,7 +721,10 @@ ur_result_t AsanInterceptor::prepareLaunch(
                     ContextInfo->Handle, DeviceInfo->Handle, (uptr)Ptr)) {
                 ReportInvalidKernelArgument(Kernel, ArgIndex, (uptr)Ptr,
                                             ValidateResult, PtrPair.second);
-                exitWithErrors();
+                if (ValidateResult.Type !=
+                    ValidateUSMResult::MAYBE_HOST_POINTER) {
+                    exitWithErrors();
+                }
             }
         }
     }
@@ -864,13 +867,15 @@ AsanInterceptor::findAllocInfoByAddress(uptr Address) {
     std::shared_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
     auto It = m_AllocationMap.upper_bound(Address);
     if (It == m_AllocationMap.begin()) {
-        return std::optional<AllocationIterator>{};
+        return std::nullopt;
     }
     --It;
-    // Make sure we got the right AllocInfo
-    assert(Address >= It->second->AllocBegin &&
-           Address < It->second->AllocBegin + It->second->AllocSize &&
-           "Wrong AllocInfo for the address");
+
+    // Maybe it's a host pointer
+    if (Address < It->second->AllocBegin ||
+        Address >= It->second->AllocBegin + It->second->AllocSize) {
+        return std::nullopt;
+    }
     return It;
 }
 

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

@@ -48,22 +48,67 @@ ur_result_t EnqueueMemCopyRectHelper(
     char *DstOrigin = pDst + DstOffset.x + DstRowPitch * DstOffset.y +
                       DstSlicePitch * DstOffset.z;
 
+    const bool IsDstDeviceUSM = getMsanInterceptor()
+                                    ->findAllocInfoByAddress((uptr)DstOrigin)
+                                    .has_value();
+    const bool IsSrcDeviceUSM = getMsanInterceptor()
+                                    ->findAllocInfoByAddress((uptr)SrcOrigin)
+                                    .has_value();
+
+    ur_device_handle_t Device = GetDevice(Queue);
+    std::shared_ptr<DeviceInfo> DeviceInfo =
+        getMsanInterceptor()->getDeviceInfo(Device);
     std::vector<ur_event_handle_t> Events;
-    Events.reserve(Region.depth);
+
     // For now, USM doesn't support 3D memory copy operation, so we can only
     // loop call 2D memory copy function to implement it.
     for (size_t i = 0; i < Region.depth; i++) {
         ur_event_handle_t NewEvent{};
         UR_CALL(getContext()->urDdiTable.Enqueue.pfnUSMMemcpy2D(
-            Queue, Blocking, DstOrigin + (i * DstSlicePitch), DstRowPitch,
+            Queue, false, DstOrigin + (i * DstSlicePitch), DstRowPitch,
             SrcOrigin + (i * SrcSlicePitch), SrcRowPitch, Region.width,
             Region.height, NumEventsInWaitList, EventWaitList, &NewEvent));
-
         Events.push_back(NewEvent);
+
+        // Update shadow memory
+        if (IsDstDeviceUSM && IsSrcDeviceUSM) {
+            NewEvent = nullptr;
+            uptr DstShadowAddr = DeviceInfo->Shadow->MemToShadow(
+                (uptr)DstOrigin + (i * DstSlicePitch));
+            uptr SrcShadowAddr = DeviceInfo->Shadow->MemToShadow(
+                (uptr)SrcOrigin + (i * SrcSlicePitch));
+            UR_CALL(getContext()->urDdiTable.Enqueue.pfnUSMMemcpy2D(
+                Queue, false, (void *)DstShadowAddr, DstRowPitch,
+                (void *)SrcShadowAddr, SrcRowPitch, Region.width, Region.height,
+                NumEventsInWaitList, EventWaitList, &NewEvent));
+            Events.push_back(NewEvent);
+        } else if (IsDstDeviceUSM && !IsSrcDeviceUSM) {
+            uptr DstShadowAddr = DeviceInfo->Shadow->MemToShadow(
+                (uptr)DstOrigin + (i * DstSlicePitch));
+            const char Val = 0;
+            // opencl & l0 adapter doesn't implement urEnqueueUSMFill2D, so
+            // emulate the operation with urEnqueueUSMFill.
+            for (size_t HeightIndex = 0; HeightIndex < Region.height;
+                 HeightIndex++) {
+                NewEvent = nullptr;
+                UR_CALL(getContext()->urDdiTable.Enqueue.pfnUSMFill(
+                    Queue, (void *)(DstShadowAddr + HeightIndex * DstRowPitch),
+                    1, &Val, Region.width, NumEventsInWaitList, EventWaitList,
+                    &NewEvent));
+                Events.push_back(NewEvent);
+            }
+        }
     }
 
-    UR_CALL(getContext()->urDdiTable.Enqueue.pfnEventsWait(
-        Queue, Events.size(), Events.data(), Event));
+    if (Blocking) {
+        UR_CALL(
+            getContext()->urDdiTable.Event.pfnWait(Events.size(), &Events[0]));
+    }
+
+    if (Event) {
+        UR_CALL(getContext()->urDdiTable.Enqueue.pfnEventsWait(
+            Queue, Events.size(), &Events[0], Event));
+    }
 
     return UR_RESULT_SUCCESS;
 }
@@ -93,7 +138,7 @@ ur_result_t MemBuffer::getHandle(ur_device_handle_t Device, char *&Handle) {
         USMDesc.align = getAlignment();
         ur_usm_pool_handle_t Pool{};
         URes = getMsanInterceptor()->allocateMemory(
-            Context, Device, &USMDesc, Pool, Size,
+            Context, Device, &USMDesc, Pool, Size, AllocType::DEVICE_USM,
             ur_cast<void **>(&Allocation));
         if (URes != UR_RESULT_SUCCESS) {
             getContext()->logger.error(
@@ -112,6 +157,12 @@ ur_result_t MemBuffer::getHandle(ur_device_handle_t Device, char *&Handle) {
                     Size, HostPtr, this);
                 return URes;
             }
+
+            // Update shadow memory
+            std::shared_ptr<DeviceInfo> DeviceInfo =
+                getMsanInterceptor()->getDeviceInfo(Device);
+            UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow(
+                Queue, (uptr)Allocation, Size, 0));
         }
     }
 
@@ -130,8 +181,8 @@ ur_result_t MemBuffer::getHandle(ur_device_handle_t Device, char *&Handle) {
             ur_usm_desc_t USMDesc{};
             USMDesc.align = getAlignment();
             ur_usm_pool_handle_t Pool{};
-            URes = getMsanInterceptor()->allocateMemory(
-                Context, nullptr, &USMDesc, Pool, Size,
+            URes = getContext()->urDdiTable.USM.pfnHostAlloc(
+                Context, &USMDesc, Pool, Size,
                 ur_cast<void **>(&HostAllocation));
             if (URes != UR_RESULT_SUCCESS) {
                 getContext()->logger.error("Failed to allocate {} bytes host "