Check use-of-uninitialized value on private memory (#17309)

Support check use-of-uninitialized value on private memory

Author: AllanZyne

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

@@ -517,10 +517,30 @@ ur_result_t MsanInterceptor::prepareLaunch(
       getContext()->logger.warning("Skip checking local memory of kernel <{}> ",
                                    GetKernelName(Kernel));
     } else {
-      getContext()->logger.info("ShadowMemory(Local, WorkGroup={}, {} - {})",
-                                NumWG,
-                                (void *)LaunchInfo.Data->LocalShadowOffset,
-                                (void *)LaunchInfo.Data->LocalShadowOffsetEnd);
+      getContext()->logger.debug("ShadowMemory(Local, WorkGroup={}, {} - {})",
+                                 NumWG,
+                                 (void *)LaunchInfo.Data->LocalShadowOffset,
+                                 (void *)LaunchInfo.Data->LocalShadowOffsetEnd);
+    }
+  }
+
+  // Write shadow memory offset for private memory
+  if (KernelInfo.IsCheckPrivates) {
+    if (DeviceInfo->Shadow->AllocPrivateShadow(
+            Queue, NumWG, LaunchInfo.Data->PrivateShadowOffset,
+            LaunchInfo.Data->PrivateShadowOffsetEnd) != UR_RESULT_SUCCESS) {
+      getContext()->logger.warning(
+          "Failed to allocate shadow memory for private "
+          "memory, maybe the number of workgroup ({}) is too "
+          "large",
+          NumWG);
+      getContext()->logger.warning(
+          "Skip checking private memory of kernel <{}>", GetKernelName(Kernel));
+    } else {
+      getContext()->logger.debug(
+          "ShadowMemory(Private, WorkGroup={}, {} - {})", NumWG,
+          (void *)LaunchInfo.Data->PrivateShadowOffset,
+          (void *)LaunchInfo.Data->PrivateShadowOffsetEnd);
     }
     // Write local arguments info
     if (!KernelInfo.LocalArgs.empty()) {
@@ -535,11 +555,11 @@ ur_result_t MsanInterceptor::prepareLaunch(
   }
 
   getContext()->logger.info(
-      "LaunchInfo {} (GlobalShadow={}, LocalShadow={}, CleanShadow={}, "
-      "LocalArgs={}, NumLocalArgs={}, "
-      "Device={}, Debug={})",
+      "LaunchInfo {} (GlobalShadow={}, LocalShadow={}, PrivateShadow={}, "
+      "CleanShadow={}, LocalArgs={}, NumLocalArgs={}, Device={}, Debug={})",
       (void *)LaunchInfo.Data, (void *)LaunchInfo.Data->GlobalShadowOffset,
       (void *)LaunchInfo.Data->LocalShadowOffset,
+      (void *)LaunchInfo.Data->PrivateShadowOffset,
       (void *)LaunchInfo.Data->CleanShadow, (void *)LaunchInfo.Data->LocalArgs,
       LaunchInfo.Data->NumLocalArgs, ToString(LaunchInfo.Data->DeviceTy),
       LaunchInfo.Data->Debug);

source/loader/layers/sanitizer/msan/msan_libdevice.hpp

@@ -37,6 +37,7 @@ struct MsanErrorReport {
 
   uint32_t AccessSize = 0;
   ErrorType ErrorTy = ErrorType::UNKNOWN;
+  uintptr_t Origin;
 };
 
 struct MsanLocalArgsInfo {
@@ -50,6 +51,9 @@ struct MsanLaunchInfo {
   uintptr_t LocalShadowOffset = 0;
   uintptr_t LocalShadowOffsetEnd = 0;
 
+  uintptr_t PrivateShadowOffset = 0;
+  uintptr_t PrivateShadowOffsetEnd = 0;
+
   uintptr_t CleanShadow = 0;
 
   DeviceType DeviceTy = DeviceType::UNKNOWN;

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

@@ -30,8 +30,15 @@ void ReportUsesUninitializedValue(const MsanErrorReport &Report,
   // Try to demangle the kernel name
   KernelName = DemangleName(KernelName);
 
-  getContext()->logger.always(
-      "====WARNING: DeviceSanitizer: use-of-uninitialized-value");
+  if (Report.Origin) {
+    getContext()->logger.always(
+        "====WARNING: DeviceSanitizer: use-of-uninitialized-value (shadow: {})",
+        (void *)Report.Origin);
+  } else {
+    getContext()->logger.always(
+        "====WARNING: DeviceSanitizer: use-of-uninitialized-value)");
+  }
+
   getContext()->logger.always(
       "use of size {} at kernel <{}> LID({}, {}, {}) GID({}, "
       "{}, {})",

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

@@ -333,6 +333,42 @@ ur_result_t MsanShadowMemoryGPU::AllocLocalShadow(ur_queue_handle_t Queue,
   return UR_RESULT_SUCCESS;
 }
 
+ur_result_t MsanShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue,
+                                                    uint32_t NumWG, uptr &Begin,
+                                                    uptr &End) {
+  const size_t RequiredShadowSize = NumWG * MSAN_PRIVATE_SIZE;
+  static size_t LastAllocedSize = 0;
+  if (RequiredShadowSize > LastAllocedSize) {
+    auto ContextInfo = getMsanInterceptor()->getContextInfo(Context);
+    if (PrivateShadowOffset) {
+      UR_CALL(getContext()->urDdiTable.USM.pfnFree(
+          Context, (void *)PrivateShadowOffset));
+      PrivateShadowOffset = 0;
+      LastAllocedSize = 0;
+    }
+
+    UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc(
+        Context, Device, nullptr, nullptr, RequiredShadowSize,
+        (void **)&PrivateShadowOffset));
+
+    // Initialize shadow memory
+    ur_result_t URes = EnqueueUSMBlockingSet(Queue, (void *)PrivateShadowOffset,
+                                             0, RequiredShadowSize);
+    if (URes != UR_RESULT_SUCCESS) {
+      UR_CALL(getContext()->urDdiTable.USM.pfnFree(
+          Context, (void *)PrivateShadowOffset));
+      PrivateShadowOffset = 0;
+      LastAllocedSize = 0;
+    }
+
+    LastAllocedSize = RequiredShadowSize;
+  }
+
+  Begin = PrivateShadowOffset;
+  End = PrivateShadowOffset + RequiredShadowSize - 1;
+  return UR_RESULT_SUCCESS;
+}
+
 uptr MsanShadowMemoryPVC::MemToShadow(uptr Ptr) {
   assert(MsanShadowMemoryPVC::IsDeviceUSM(Ptr) && "Ptr must be device USM");
   if (Ptr < ShadowBegin) {

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

@@ -46,6 +46,10 @@ struct MsanShadowMemory {
   virtual ur_result_t AllocLocalShadow(ur_queue_handle_t Queue, uint32_t NumWG,
                                        uptr &Begin, uptr &End) = 0;
 
+  virtual ur_result_t AllocPrivateShadow(ur_queue_handle_t Queue,
+                                         uint32_t NumWG, uptr &Begin,
+                                         uptr &End) = 0;
+
   ur_context_handle_t Context{};
 
   ur_device_handle_t Device{};
@@ -95,6 +99,13 @@ struct MsanShadowMemoryCPU final : public MsanShadowMemory {
     End = ShadowEnd;
     return UR_RESULT_SUCCESS;
   }
+
+  ur_result_t AllocPrivateShadow(ur_queue_handle_t, uint32_t, uptr &Begin,
+                                 uptr &End) override {
+    Begin = ShadowBegin;
+    End = ShadowEnd;
+    return UR_RESULT_SUCCESS;
+  }
 };
 
 struct MsanShadowMemoryGPU : public MsanShadowMemory {
@@ -116,6 +127,9 @@ struct MsanShadowMemoryGPU : public MsanShadowMemory {
   ur_result_t AllocLocalShadow(ur_queue_handle_t Queue, uint32_t NumWG,
                                uptr &Begin, uptr &End) override final;
 
+  ur_result_t AllocPrivateShadow(ur_queue_handle_t Queue, uint32_t NumWG,
+                                 uptr &Begin, uptr &End) override final;
+
   virtual size_t GetShadowSize() = 0;
 
   virtual uptr GetStartAddress() { return 0; }
@@ -132,6 +146,8 @@ struct MsanShadowMemoryGPU : public MsanShadowMemory {
   ur_mutex VirtualMemMapsMutex;
 
   uptr LocalShadowOffset = 0;
+
+  uptr PrivateShadowOffset = 0;
 };
 
 // clang-format off