[DeviceSanitizer] Support detecting out-of-bounds error on DeviceGlobals

zhaomaosu · zhaomaosu · commit 38e6f788858a · 2024-02-18T09:39:57.000+08:00
diff --git a/source/loader/layers/sanitizer/asan_interceptor.cpp b/source/loader/layers/sanitizer/asan_interceptor.cpp
@@ -24,6 +24,7 @@ constexpr int kUsmDeviceRedzoneMagic = (char)0x81;
 constexpr int kUsmHostRedzoneMagic = (char)0x82;
 constexpr int kUsmSharedRedzoneMagic = (char)0x83;
 constexpr int kMemBufferRedzoneMagic = (char)0x84;
+constexpr int kDeviceGlobalRedZoneMagic = (char)0x85;
 
 constexpr auto kSPIR_AsanShadowMemoryGlobalStart =
     "__AsanShadowMemoryGlobalStart";
@@ -36,6 +37,9 @@ constexpr auto kSPIR_DeviceType = "__DeviceType";
 
 constexpr auto kSPIR_DeviceSanitizerReportMem = "__DeviceSanitizerReportMem";
 
+constexpr auto kSPIR_AsanDeviceGlobalCount = "__AsanDeviceGlobalCount";
+constexpr auto kSPIR_AsanDeviceGlobalMetadata = "__AsanDeviceGlobalMetadata";
+
 DeviceSanitizerReport SPIR_DeviceSanitizerReportMem;
 
 uptr MemToShadow_CPU(uptr USM_SHADOW_BASE, uptr UPtr) {
@@ -78,6 +82,19 @@ ur_program_handle_t getProgram(ur_kernel_handle_t Kernel) {
     return Program;
 }
 
+void getProgramDevices(ur_program_handle_t Program,
+                       std::vector<ur_device_handle_t> &Devices) {
+    size_t PropSize;
+    [[maybe_unused]] ur_result_t Result = context.urDdiTable.Program.pfnGetInfo(
+        Program, UR_PROGRAM_INFO_DEVICES, 0, nullptr, &PropSize);
+    assert(Result == UR_RESULT_SUCCESS);
+
+    Devices.resize(PropSize / sizeof(ur_device_handle_t));
+    Result = context.urDdiTable.Program.pfnGetInfo(
+        Program, UR_PROGRAM_INFO_DEVICES, PropSize, Devices.data(), nullptr);
+    assert(Result == UR_RESULT_SUCCESS);
+}
+
 size_t getLocalMemorySize(ur_device_handle_t Device) {
     size_t LocalMemorySize;
     [[maybe_unused]] auto Result = context.urDdiTable.Device.pfnGetInfo(
@@ -124,7 +141,7 @@ SanitizerInterceptor::~SanitizerInterceptor() {
 ur_result_t SanitizerInterceptor::allocateMemory(
     ur_context_handle_t Context, ur_device_handle_t Device,
     const ur_usm_desc_t *Properties, ur_usm_pool_handle_t Pool, size_t Size,
-    void **ResultPtr, USMMemoryType Type) {
+    void **ResultPtr, AllocType Type) {
     auto Alignment = Properties->align;
     assert(Alignment == 0 || IsPowerOfTwo(Alignment));
 
@@ -147,13 +164,13 @@ ur_result_t SanitizerInterceptor::allocateMemory(
 
     void *Allocated = nullptr;
 
-    if (Type == USMMemoryType::DEVICE) {
+    if (Type == AllocType::DEVICE_USM) {
         UR_CALL(context.urDdiTable.USM.pfnDeviceAlloc(
             Context, Device, Properties, Pool, NeededSize, &Allocated));
-    } else if (Type == USMMemoryType::HOST) {
+    } else if (Type == AllocType::HOST_USM) {
         UR_CALL(context.urDdiTable.USM.pfnHostAlloc(Context, Properties, Pool,
                                                     NeededSize, &Allocated));
-    } else if (Type == USMMemoryType::SHARE) {
+    } else if (Type == AllocType::SHARED_USM) {
         UR_CALL(context.urDdiTable.USM.pfnSharedAlloc(
             Context, Device, Properties, Pool, NeededSize, &Allocated));
     } else {
@@ -173,31 +190,31 @@ ur_result_t SanitizerInterceptor::allocateMemory(
 
     *ResultPtr = reinterpret_cast<void *>(UserBegin);
 
-    auto AllocInfo = std::make_shared<USMAllocInfo>(
-        USMAllocInfo{AllocBegin, UserBegin, UserEnd, NeededSize, Type});
+    auto AI = std::make_shared<AllocInfo>(
+        AllocInfo{AllocBegin, UserBegin, UserEnd, NeededSize, Type});
 
     // For updating shadow memory
     if (DeviceInfo) { // device/shared USM
         std::scoped_lock<ur_shared_mutex> Guard(DeviceInfo->Mutex);
-        DeviceInfo->AllocInfos.emplace_back(AllocInfo);
+        DeviceInfo->AllocInfos.emplace_back(AI);
     } else { // host USM's AllocInfo needs to insert into all devices
         for (auto &pair : ContextInfo->DeviceMap) {
             auto DeviceInfo = pair.second;
             std::scoped_lock<ur_shared_mutex> Guard(DeviceInfo->Mutex);
-            DeviceInfo->AllocInfos.emplace_back(AllocInfo);
+            DeviceInfo->AllocInfos.emplace_back(AI);
         }
     }
 
     // For memory release
     {
         std::scoped_lock<ur_shared_mutex> Guard(ContextInfo->Mutex);
-        ContextInfo->AllocatedUSMMap[AllocBegin] = std::move(AllocInfo);
+        ContextInfo->AllocatedUSMMap[AllocBegin] = std::move(AI);
     }
 
     context.logger.info(
         "AllocInfos(AllocBegin={},  User={}-{}, NeededSize={}, Type={})",
         (void *)AllocBegin, (void *)UserBegin, (void *)UserEnd, NeededSize,
-        Type);
+        ToString(Type));
 
     return UR_RESULT_SUCCESS;
 }
@@ -282,8 +299,8 @@ void SanitizerInterceptor::postLaunchKernel(ur_kernel_handle_t Kernel,
         auto KernelName = getKernelName(Kernel);
 
         context.logger.always("\n====ERROR: DeviceSanitizer: {} on {}",
-                              DeviceSanitizerFormat(AH->ErrorType),
-                              DeviceSanitizerFormat(AH->MemoryType));
+                              ToString(AH->ErrorType),
+                              ToString(AH->MemoryType));
         context.logger.always(
             "{} of size {} at kernel <{}> LID({}, {}, {}) GID({}, "
             "{}, {})",
@@ -475,7 +492,7 @@ ur_result_t SanitizerInterceptor::enqueueMemSetShadow(
 /// ref: https://github.com/google/sanitizers/wiki/AddressSanitizerAlgorithm#mapping
 ur_result_t SanitizerInterceptor::enqueueAllocInfo(
     ur_context_handle_t Context, ur_device_handle_t Device,
-    ur_queue_handle_t Queue, std::shared_ptr<USMAllocInfo> &AllocInfo,
+    ur_queue_handle_t Queue, std::shared_ptr<AllocInfo> &AllocInfo,
     ur_event_handle_t &LastEvent) {
     // Init zero
     UR_CALL(enqueueMemSetShadow(Context, Device, Queue, AllocInfo->AllocBegin,
@@ -496,18 +513,21 @@ ur_result_t SanitizerInterceptor::enqueueAllocInfo(
 
     int ShadowByte;
     switch (AllocInfo->Type) {
-    case USMMemoryType::HOST:
+    case AllocType::HOST_USM:
         ShadowByte = kUsmHostRedzoneMagic;
         break;
-    case USMMemoryType::DEVICE:
+    case AllocType::DEVICE_USM:
         ShadowByte = kUsmDeviceRedzoneMagic;
         break;
-    case USMMemoryType::SHARE:
+    case AllocType::SHARED_USM:
         ShadowByte = kUsmSharedRedzoneMagic;
         break;
-    case USMMemoryType::MEM_BUFFER:
+    case AllocType::MEM_BUFFER:
         ShadowByte = kMemBufferRedzoneMagic;
         break;
+    case AllocType::DEVICE_GLOBAL:
+        ShadowByte = kDeviceGlobalRedZoneMagic;
+        break;
     default:
         ShadowByte = 0xff;
         assert(false && "Unknow AllocInfo Type");
@@ -553,6 +573,62 @@ ur_result_t SanitizerInterceptor::updateShadowMemory(ur_queue_handle_t Queue) {
     return UR_RESULT_SUCCESS;
 }
 
+ur_result_t
+SanitizerInterceptor::registerDeviceGlobals(ur_context_handle_t Context,
+                                         ur_program_handle_t Program) {
+    std::vector<ur_device_handle_t> Devices;
+    getProgramDevices(Program, Devices);
+
+    for (auto Device : Devices) {
+        ur_queue_handle_t Queue;
+        ur_result_t Result = context.urDdiTable.Queue.pfnCreate(
+            Context, Device, nullptr, &Queue);
+        if (Result != UR_RESULT_SUCCESS) {
+            context.logger.error("Failed to create command queue: {}", Result);
+            return Result;
+        }
+
+        uint64_t NumOfDeviceGlobal;
+        Result = context.urDdiTable.Enqueue.pfnDeviceGlobalVariableRead(
+            Queue, Program, kSPIR_AsanDeviceGlobalCount, true,
+            sizeof(NumOfDeviceGlobal), 0, &NumOfDeviceGlobal, 0, nullptr,
+            nullptr);
+        if (Result == UR_RESULT_ERROR_INVALID_ARGUMENT) {
+            context.logger.info("No device globals");
+            continue;
+        } else if (Result != UR_RESULT_SUCCESS) {
+            context.logger.error("Device Global[{}] Read Failed: {}",
+                                 kSPIR_AsanDeviceGlobalCount, Result);
+            return Result;
+        }
+
+        std::vector<DeviceGlobalInfo> GVInfos(NumOfDeviceGlobal);
+        Result = context.urDdiTable.Enqueue.pfnDeviceGlobalVariableRead(
+            Queue, Program, kSPIR_AsanDeviceGlobalMetadata, true,
+            sizeof(DeviceGlobalInfo) * NumOfDeviceGlobal, 0, &GVInfos[0], 0,
+            nullptr, nullptr);
+        if (Result != UR_RESULT_SUCCESS) {
+            context.logger.error("Device Global[{}] Read Failed: {}",
+                                 kSPIR_AsanDeviceGlobalMetadata, Result);
+            return Result;
+        }
+
+        auto ContextInfo = getContextInfo(Context);
+        auto DeviceInfo = ContextInfo->getDeviceInfo(Device);
+        for (size_t i = 0; i < NumOfDeviceGlobal; i++) {
+            auto AI = std::make_shared<AllocInfo>(AllocInfo{
+                GVInfos[i].Addr, GVInfos[i].Addr,
+                GVInfos[i].Addr + GVInfos[i].Size, GVInfos[i].SizeWithRedZone,
+                AllocType::DEVICE_GLOBAL});
+
+            std::scoped_lock<ur_shared_mutex> Guard(DeviceInfo->Mutex);
+            DeviceInfo->AllocInfos.emplace_back(AI);
+        }
+    }
+
+    return UR_RESULT_SUCCESS;
+}
+
 ur_result_t SanitizerInterceptor::insertContext(ur_context_handle_t Context) {
     auto ContextInfo = std::make_shared<ur_sanitizer_layer::ContextInfo>();
 
diff --git a/source/loader/layers/sanitizer/asan_interceptor.hpp b/source/loader/layers/sanitizer/asan_interceptor.hpp
@@ -22,14 +22,20 @@
 
 namespace ur_sanitizer_layer {
 
-enum USMMemoryType { DEVICE, SHARE, HOST, MEM_BUFFER };
+enum class AllocType : uint32_t {
+    DEVICE_USM,
+    SHARED_USM,
+    HOST_USM,
+    MEM_BUFFER,
+    DEVICE_GLOBAL
+};
 
-struct USMAllocInfo {
+struct AllocInfo {
     uptr AllocBegin;
     uptr UserBegin;
     uptr UserEnd;
     size_t AllocSize;
-    USMMemoryType Type;
+    AllocType Type;
 };
 
 enum class DeviceType { UNKNOWN, CPU, GPU_PVC, GPU_DG2 };
@@ -42,7 +48,7 @@ struct DeviceInfo {
 
     // Lock InitPool & AllocInfos
     ur_shared_mutex Mutex;
-    std::vector<std::shared_ptr<USMAllocInfo>> AllocInfos;
+    std::vector<std::shared_ptr<AllocInfo>> AllocInfos;
 };
 
 struct QueueInfo {
@@ -64,7 +70,7 @@ struct ContextInfo {
         return QueueMap[Queue];
     }
 
-    std::shared_ptr<USMAllocInfo> getUSMAllocInfo(uptr Address) {
+    std::shared_ptr<AllocInfo> getUSMAllocInfo(uptr Address) {
         std::shared_lock<ur_shared_mutex> Guard(Mutex);
         assert(AllocatedUSMMap.find(Address) != AllocatedUSMMap.end());
         return AllocatedUSMMap[Address];
@@ -78,7 +84,7 @@ struct ContextInfo {
     /// key: USMAllocInfo.AllocBegin
     /// value: USMAllocInfo
     /// Use AllocBegin as key can help to detect underflow pointer
-    std::map<uptr, std::shared_ptr<USMAllocInfo>> AllocatedUSMMap;
+    std::map<uptr, std::shared_ptr<AllocInfo>> AllocatedUSMMap;
 };
 
 struct LaunchInfo {
@@ -95,6 +101,12 @@ struct LaunchInfo {
     ~LaunchInfo();
 };
 
+struct DeviceGlobalInfo {
+    uptr Size;
+    uptr SizeWithRedZone;
+    uptr Addr;
+};
+
 class SanitizerInterceptor {
   public:
     SanitizerInterceptor();
@@ -105,9 +117,12 @@ class SanitizerInterceptor {
                                ur_device_handle_t Device,
                                const ur_usm_desc_t *Properties,
                                ur_usm_pool_handle_t Pool, size_t Size,
-                               void **ResultPtr, USMMemoryType Type);
+                               void **ResultPtr, AllocType Type);
     ur_result_t releaseMemory(ur_context_handle_t Context, void *Ptr);
 
+    ur_result_t registerDeviceGlobals(ur_context_handle_t Context,
+                                      ur_program_handle_t Program);
+
     ur_result_t preLaunchKernel(ur_kernel_handle_t Kernel,
                                 ur_queue_handle_t Queue,
                                 ur_event_handle_t &Event,
@@ -131,7 +146,7 @@ class SanitizerInterceptor {
     ur_result_t enqueueAllocInfo(ur_context_handle_t Context,
                                  ur_device_handle_t Device,
                                  ur_queue_handle_t Queue,
-                                 std::shared_ptr<USMAllocInfo> &AlloccInfo,
+                                 std::shared_ptr<AllocInfo> &AI,
                                  ur_event_handle_t &LastEvent);
 
     /// Initialize Global Variables & Kernel Name at first Launch
@@ -163,4 +178,21 @@ class SanitizerInterceptor {
     bool m_ShadowMemInited;
 };
 
+inline const char *ToString(AllocType Type) {
+    switch (Type) {
+    case AllocType::DEVICE_USM:
+        return "Device USM";
+    case AllocType::HOST_USM:
+        return "Host USM";
+    case AllocType::SHARED_USM:
+        return "Shared USM";
+    case AllocType::MEM_BUFFER:
+        return "Memory Buffer";
+    case AllocType::DEVICE_GLOBAL:
+        return "Device Global";
+    default:
+        return "Unknown Type";
+    }
+}
+
 } // namespace ur_sanitizer_layer
diff --git a/source/loader/layers/sanitizer/device_sanitizer_report.hpp b/source/loader/layers/sanitizer/device_sanitizer_report.hpp
@@ -32,6 +32,7 @@ enum class DeviceSanitizerMemoryType : int32_t {
     LOCAL,
     PRIVATE,
     MEM_BUFFER,
+    DEVICE_GLOBAL,
 };
 
 struct DeviceSanitizerReport {
@@ -58,26 +59,28 @@ struct DeviceSanitizerReport {
     bool IsRecover = false;
 };
 
-inline const char *DeviceSanitizerFormat(DeviceSanitizerMemoryType MemoryType) {
+inline const char *ToString(DeviceSanitizerMemoryType MemoryType) {
     switch (MemoryType) {
     case DeviceSanitizerMemoryType::USM_DEVICE:
-        return "USM Device Memory";
+        return "Device USM";
     case DeviceSanitizerMemoryType::USM_HOST:
-        return "USM Host Memory";
+        return "Host USM";
     case DeviceSanitizerMemoryType::USM_SHARED:
-        return "USM Shared Memory";
+        return "Shared USM";
     case DeviceSanitizerMemoryType::LOCAL:
         return "Local Memory";
     case DeviceSanitizerMemoryType::PRIVATE:
         return "Private Memory";
     case DeviceSanitizerMemoryType::MEM_BUFFER:
         return "Memory Buffer";
+    case DeviceSanitizerMemoryType::DEVICE_GLOBAL:
+        return "Device Global";
     default:
         return "Unknown Memory";
     }
 }
 
-inline const char *DeviceSanitizerFormat(DeviceSanitizerErrorType ErrorType) {
+inline const char *ToString(DeviceSanitizerErrorType ErrorType) {
     switch (ErrorType) {
     case DeviceSanitizerErrorType::OUT_OF_BOUNDS:
         return "out-of-bounds-access";
diff --git a/source/loader/layers/sanitizer/ur_sanddi.cpp b/source/loader/layers/sanitizer/ur_sanddi.cpp
