[DeviceSanitizer] Support out-of-bounds on private memory (#1676)

* support private memory

* fix comment

* enable private by default

* using 24bit ASAN_PRIVATE_SIZE

* use urKernelGetSuggestedLocalWorkSize

* skip checking kernel if shadow memory is allocated failed

* add PrivateShadowOffsetEnd

* change default logger level

* add logger in ManagedQueue

* revert

* revert back

* fix build

* use error

* remove pfnGetSuggestedLocalWorkSize

* fix m_Quarantine

* add comment

* fallback urKernelGetSuggestedLocalWorkSize

Author: AllanZyne

source/loader/layers/sanitizer/asan_interceptor.cpp

@@ -714,16 +714,21 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
         EnqueueWriteGlobal(kSPIR_DeviceType, &DeviceInfo->Type,
                            sizeof(DeviceInfo->Type));
 
-        if (DeviceInfo->Type == DeviceType::CPU) {
-            break;
-        }
-
         if (LaunchInfo.LocalWorkSize.empty()) {
-            LaunchInfo.LocalWorkSize.reserve(3);
-            // FIXME: This is W/A until urKernelSuggestGroupSize is added
-            LaunchInfo.LocalWorkSize[0] = 1;
-            LaunchInfo.LocalWorkSize[1] = 1;
-            LaunchInfo.LocalWorkSize[2] = 1;
+            LaunchInfo.LocalWorkSize.resize(LaunchInfo.WorkDim);
+            auto URes = context.urDdiTable.Kernel.pfnGetSuggestedLocalWorkSize(
+                Kernel, Queue, LaunchInfo.WorkDim, LaunchInfo.GlobalWorkOffset,
+                LaunchInfo.GlobalWorkSize, LaunchInfo.LocalWorkSize.data());
+            if (URes != UR_RESULT_SUCCESS) {
+                if (URes != UR_RESULT_ERROR_UNSUPPORTED_FEATURE) {
+                    return URes;
+                }
+                // If urKernelGetSuggestedLocalWorkSize is not supported by driver, we fallback
+                // to inefficient implementation
+                for (size_t Dim = 0; Dim < LaunchInfo.WorkDim; ++Dim) {
+                    LaunchInfo.LocalWorkSize[Dim] = 1;
+                }
+            }
         }
 
         const size_t *LocalWorkSize = LaunchInfo.LocalWorkSize.data();
@@ -733,56 +738,109 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
                      LocalWorkSize[Dim];
         }
 
-        auto EnqueueAllocateDevice = [Context, &DeviceInfo, Queue,
-                                      NumWG](size_t Size, uptr &Ptr) {
+        auto EnqueueAllocateShadowMemory = [Context, &DeviceInfo,
+                                            Queue](size_t Size, uptr &Ptr) {
+            void *Allocated = nullptr;
             auto URes = context.urDdiTable.USM.pfnDeviceAlloc(
                 Context, DeviceInfo->Handle, nullptr, nullptr, Size,
-                (void **)&Ptr);
+                &Allocated);
             if (URes != UR_RESULT_SUCCESS) {
-                context.logger.error(
-                    "Failed to allocate shadow memory for local memory: {}",
-                    URes);
-                context.logger.error(
-                    "Maybe the number of workgroup ({}) too large", NumWG);
                 return URes;
             }
-            // Initialize shadow memory of local memory
-            URes = urEnqueueUSMSet(Queue, (void *)Ptr, 0, Size);
-            if (URes == UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY) {
-                context.logger.error(
-                    "Failed to allocate shadow memory for local memory: {}",
-                    URes);
-                context.logger.error(
-                    "Maybe the number of workgroup ({}) too large", NumWG);
-                return URes;
+            // Initialize shadow memory
+            URes = urEnqueueUSMSet(Queue, Allocated, 0, Size);
+            if (URes != UR_RESULT_SUCCESS) {
+                [[maybe_unused]] auto URes =
+                    context.urDdiTable.USM.pfnFree(Context, Allocated);
+                assert(URes == UR_RESULT_SUCCESS &&
+                       "urUSMFree failed at allocating shadow memory");
+                Allocated = nullptr;
             }
+            Ptr = (uptr)Allocated;
             return URes;
         };
 
+        auto LocalMemoryUsage =
+            GetKernelLocalMemorySize(Kernel, DeviceInfo->Handle);
+        auto PrivateMemoryUsage =
+            GetKernelPrivateMemorySize(Kernel, DeviceInfo->Handle);
+
+        context.logger.info("KernelInfo {} (LocalMemory={}, PrivateMemory={})",
+                            (void *)Kernel, LocalMemoryUsage,
+                            PrivateMemoryUsage);
+
         // Write shadow memory offset for local memory
         if (Options().DetectLocals) {
             // CPU needn't this
             if (DeviceInfo->Type == DeviceType::GPU_PVC) {
-                size_t LocalMemorySize = GetLocalMemorySize(DeviceInfo->Handle);
-                size_t LocalShadowMemorySize =
+                const size_t LocalMemorySize =
+                    GetDeviceLocalMemorySize(DeviceInfo->Handle);
+                const size_t LocalShadowMemorySize =
                     (NumWG * LocalMemorySize) >> ASAN_SHADOW_SCALE;
 
                 context.logger.debug(
-                    "LocalMemoryInfo(WorkGroup={}, LocalMemorySize={}, "
+                    "LocalMemory(WorkGroup={}, LocalMemorySize={}, "
                     "LocalShadowMemorySize={})",
                     NumWG, LocalMemorySize, LocalShadowMemorySize);
 
-                UR_CALL(EnqueueAllocateDevice(
-                    LocalShadowMemorySize, LaunchInfo.Data->LocalShadowOffset));
-
-                LaunchInfo.Data->LocalShadowOffsetEnd =
-                    LaunchInfo.Data->LocalShadowOffset + LocalShadowMemorySize -
-                    1;
+                if (EnqueueAllocateShadowMemory(
+                        LocalShadowMemorySize,
+                        LaunchInfo.Data->LocalShadowOffset) !=
+                    UR_RESULT_SUCCESS) {
+                    context.logger.warning(
+                        "Failed to allocate shadow memory for local "
+                        "memory, maybe the number of workgroup ({}) is too "
+                        "large",
+                        NumWG);
+                    context.logger.warning(
+                        "Skip checking local memory of kernel <{}>",
+                        GetKernelName(Kernel));
+                } else {
+                    LaunchInfo.Data->LocalShadowOffsetEnd =
+                        LaunchInfo.Data->LocalShadowOffset +
+                        LocalShadowMemorySize - 1;
+
+                    context.logger.info(
+                        "ShadowMemory(Local, {} - {})",
+                        (void *)LaunchInfo.Data->LocalShadowOffset,
+                        (void *)LaunchInfo.Data->LocalShadowOffsetEnd);
+                }
+            }
+        }
 
-                context.logger.info(
-                    "ShadowMemory(Local, {} - {})",
-                    (void *)LaunchInfo.Data->LocalShadowOffset,
-                    (void *)LaunchInfo.Data->LocalShadowOffsetEnd);
+        // Write shadow memory offset for private memory
+        if (Options().DetectPrivates) {
+            if (DeviceInfo->Type == DeviceType::CPU) {
+                LaunchInfo.Data->PrivateShadowOffset = DeviceInfo->ShadowOffset;
+            } else if (DeviceInfo->Type == DeviceType::GPU_PVC) {
+                const size_t PrivateShadowMemorySize =
+                    (NumWG * ASAN_PRIVATE_SIZE) >> ASAN_SHADOW_SCALE;
+
+                context.logger.debug("PrivateMemory(WorkGroup={}, "
+                                     "PrivateShadowMemorySize={})",
+                                     NumWG, PrivateShadowMemorySize);
+
+                if (EnqueueAllocateShadowMemory(
+                        PrivateShadowMemorySize,
+                        LaunchInfo.Data->PrivateShadowOffset) !=
+                    UR_RESULT_SUCCESS) {
+                    context.logger.warning(
+                        "Failed to allocate shadow memory for private "
+                        "memory, maybe the number of workgroup ({}) is too "
+                        "large",
+                        NumWG);
+                    context.logger.warning(
+                        "Skip checking private memory of kernel <{}>",
+                        GetKernelName(Kernel));
+                } else {
+                    LaunchInfo.Data->PrivateShadowOffsetEnd =
+                        LaunchInfo.Data->PrivateShadowOffset +
+                        PrivateShadowMemorySize - 1;
+                    context.logger.info(
+                        "ShadowMemory(Private, {} - {})",
+                        (void *)LaunchInfo.Data->PrivateShadowOffset,
+                        (void *)LaunchInfo.Data->PrivateShadowOffsetEnd);
+                }
             }
         }
     } while (false);

source/loader/layers/sanitizer/asan_libdevice.hpp

@@ -70,8 +70,8 @@ struct LocalArgsInfo {
 constexpr std::size_t ASAN_MAX_NUM_REPORTS = 10;
 
 struct LaunchInfo {
-    // Don't move this field, we use it in AddressSanitizerPass
     uintptr_t PrivateShadowOffset = 0;
+    uintptr_t PrivateShadowOffsetEnd = 0;
 
     uintptr_t LocalShadowOffset = 0;
     uintptr_t LocalShadowOffsetEnd = 0;
@@ -85,6 +85,10 @@ struct LaunchInfo {
 constexpr unsigned ASAN_SHADOW_SCALE = 4;
 constexpr unsigned ASAN_SHADOW_GRANULARITY = 1ULL << ASAN_SHADOW_SCALE;
 
+// Based on the observation, only the last 24 bits of the address of the private
+// variable have changed
+constexpr std::size_t ASAN_PRIVATE_SIZE = 0xffffffULL + 1;
+
 // These magic values are written to shadow for better error
 // reporting.
 constexpr int kUsmDeviceRedzoneMagic = (char)0x81;

source/loader/layers/sanitizer/asan_options.hpp

@@ -37,6 +37,7 @@ struct AsanOptions {
     uint64_t MaxRZSize = 2048;
     uint32_t MaxQuarantineSizeMB = 0;
     bool DetectLocals = true;
+    bool DetectPrivates = true;
 
   private:
     AsanOptions() {
@@ -91,6 +92,7 @@ struct AsanOptions {
 
         SetBoolOption("debug", Debug);
         SetBoolOption("detect_locals", DetectLocals);
+        SetBoolOption("detect_privates", DetectPrivates);
 
         auto KV = OptionsEnvMap->find("quarantine_size_mb");
         if (KV != OptionsEnvMap->end()) {

source/loader/layers/sanitizer/ur_sanddi.cpp

@@ -280,9 +280,10 @@ __urdlllocal ur_result_t UR_APICALL urEnqueueKernelLaunch(
     UR_CALL(context.interceptor->preLaunchKernel(hKernel, hQueue, LaunchInfo));
 
     ur_event_handle_t hEvent{};
-    ur_result_t result = pfnKernelLaunch(
-        hQueue, hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize,
-        pLocalWorkSize, numEventsInWaitList, phEventWaitList, &hEvent);
+    ur_result_t result =
+        pfnKernelLaunch(hQueue, hKernel, workDim, pGlobalWorkOffset,
+                        pGlobalWorkSize, LaunchInfo.LocalWorkSize.data(),
+                        numEventsInWaitList, phEventWaitList, &hEvent);
 
     if (result == UR_RESULT_SUCCESS) {
         UR_CALL(

source/loader/layers/sanitizer/ur_sanitizer_layer.cpp

@@ -18,7 +18,8 @@ context_t context;
 
 ///////////////////////////////////////////////////////////////////////////////
 context_t::context_t()
-    : logger(logger::create_logger("sanitizer")),
+    : logger(logger::create_logger("sanitizer", false, false,
+                                   logger::Level::WARN)),
       interceptor(std::make_unique<SanitizerInterceptor>()) {}
 
 bool context_t::isAvailable() const { return true; }

source/loader/layers/sanitizer/ur_sanitizer_utils.cpp

@@ -19,7 +19,7 @@ ManagedQueue::ManagedQueue(ur_context_handle_t Context,
                            ur_device_handle_t Device) {
     [[maybe_unused]] auto Result =
         context.urDdiTable.Queue.pfnCreate(Context, Device, nullptr, &Handle);
-    assert(Result == UR_RESULT_SUCCESS);
+    assert(Result == UR_RESULT_SUCCESS && "Failed to create ManagedQueue");
     context.logger.debug(">>> ManagedQueue {}", (void *)Handle);
 }
 
@@ -31,9 +31,9 @@ ManagedQueue::~ManagedQueue() {
     if (Result != UR_RESULT_SUCCESS) {
         context.logger.error("Failed to finish ManagedQueue: {}", Result);
     }
-    assert(Result == UR_RESULT_SUCCESS);
+    assert(Result == UR_RESULT_SUCCESS && "Failed to finish ManagedQueue");
     Result = context.urDdiTable.Queue.pfnRelease(Handle);
-    assert(Result == UR_RESULT_SUCCESS);
+    assert(Result == UR_RESULT_SUCCESS && "Failed to release ManagedQueue");
 }
 
 ur_context_handle_t GetContext(ur_queue_handle_t Queue) {
@@ -81,7 +81,7 @@ ur_program_handle_t GetProgram(ur_kernel_handle_t Kernel) {
     return Program;
 }
 
-size_t GetLocalMemorySize(ur_device_handle_t Device) {
+size_t GetDeviceLocalMemorySize(ur_device_handle_t Device) {
     size_t LocalMemorySize{};
     [[maybe_unused]] auto Result = context.urDdiTable.Device.pfnGetInfo(
         Device, UR_DEVICE_INFO_LOCAL_MEM_SIZE, sizeof(LocalMemorySize),
@@ -157,6 +157,26 @@ size_t GetKernelNumArgs(ur_kernel_handle_t Kernel) {
     return NumArgs;
 }
 
+size_t GetKernelLocalMemorySize(ur_kernel_handle_t Kernel,
+                                ur_device_handle_t Device) {
+    size_t Size = 0;
+    [[maybe_unused]] auto Res = context.urDdiTable.Kernel.pfnGetGroupInfo(
+        Kernel, Device, UR_KERNEL_GROUP_INFO_LOCAL_MEM_SIZE, sizeof(size_t),
+        &Size, nullptr);
+    assert(Res == UR_RESULT_SUCCESS);
+    return Size;
+}
+
+size_t GetKernelPrivateMemorySize(ur_kernel_handle_t Kernel,
+                                  ur_device_handle_t Device) {
+    size_t Size = 0;
+    [[maybe_unused]] auto Res = context.urDdiTable.Kernel.pfnGetGroupInfo(
+        Kernel, Device, UR_KERNEL_GROUP_INFO_PRIVATE_MEM_SIZE, sizeof(size_t),
+        &Size, nullptr);
+    assert(Res == UR_RESULT_SUCCESS);
+    return Size;
+}
+
 size_t GetVirtualMemGranularity(ur_context_handle_t Context,
                                 ur_device_handle_t Device) {
     size_t Size;

source/loader/layers/sanitizer/ur_sanitizer_utils.hpp

@@ -36,12 +36,16 @@ ur_context_handle_t GetContext(ur_kernel_handle_t Kernel);
 ur_device_handle_t GetDevice(ur_queue_handle_t Queue);
 DeviceType GetDeviceType(ur_device_handle_t Device);
 std::string GetKernelName(ur_kernel_handle_t Kernel);
-size_t GetLocalMemorySize(ur_device_handle_t Device);
+size_t GetDeviceLocalMemorySize(ur_device_handle_t Device);
 ur_program_handle_t GetProgram(ur_kernel_handle_t Kernel);
 std::vector<ur_device_handle_t> GetProgramDevices(ur_program_handle_t Program);
 ur_device_handle_t GetUSMAllocDevice(ur_context_handle_t Context,
                                      const void *MemPtr);
 size_t GetKernelNumArgs(ur_kernel_handle_t Kernel);
+size_t GetKernelLocalMemorySize(ur_kernel_handle_t Kernel,
+                                ur_device_handle_t Device);
+size_t GetKernelPrivateMemorySize(ur_kernel_handle_t Kernel,
+                                  ur_device_handle_t Device);
 size_t GetVirtualMemGranularity(ur_context_handle_t Context,
                                 ur_device_handle_t Device);