wip

Author: AllanZyne

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

@@ -144,6 +144,12 @@ ur_result_t MsanInterceptor::registerProgram(ur_program_handle_t Program) {
         return Result;
     }
 
+    getContext()->logger.info("registerDeviceGlobals");
+    Result = registerDeviceGlobals(Program);
+    if (Result != UR_RESULT_SUCCESS) {
+        return Result;
+    }
+
     return Result;
 }
 
@@ -212,6 +218,53 @@ ur_result_t MsanInterceptor::registerSpirKernels(ur_program_handle_t Program) {
     return UR_RESULT_SUCCESS;
 }
 
+ur_result_t
+MsanInterceptor::registerDeviceGlobals(ur_program_handle_t Program) {
+    std::vector<ur_device_handle_t> Devices = GetDevices(Program);
+    assert(Devices.size() != 0 && "No devices in registerDeviceGlobals");
+    auto Context = GetContext(Program);
+    auto ContextInfo = getContextInfo(Context);
+    auto ProgramInfo = getProgramInfo(Program);
+    assert(ProgramInfo != nullptr && "unregistered program!");
+
+    for (auto Device : Devices) {
+        ManagedQueue Queue(Context, Device);
+
+        size_t MetadataSize;
+        void *MetadataPtr;
+        auto Result =
+            getContext()->urDdiTable.Program.pfnGetGlobalVariablePointer(
+                Device, Program, kSPIR_MsanDeviceGlobalMetadata, &MetadataSize,
+                &MetadataPtr);
+        if (Result != UR_RESULT_SUCCESS) {
+            getContext()->logger.info("No device globals");
+            continue;
+        }
+
+        const uint64_t NumOfDeviceGlobal =
+            MetadataSize / sizeof(DeviceGlobalInfo);
+        assert((MetadataSize % sizeof(DeviceGlobalInfo) == 0) &&
+               "DeviceGlobal metadata size is not correct");
+        std::vector<DeviceGlobalInfo> GVInfos(NumOfDeviceGlobal);
+        Result = getContext()->urDdiTable.Enqueue.pfnUSMMemcpy(
+            Queue, true, &GVInfos[0], MetadataPtr,
+            sizeof(DeviceGlobalInfo) * NumOfDeviceGlobal, 0, nullptr, nullptr);
+        if (Result != UR_RESULT_SUCCESS) {
+            getContext()->logger.error("Device Global[{}] Read Failed: {}",
+                                       kSPIR_MsanDeviceGlobalMetadata, Result);
+            return Result;
+        }
+
+        auto DeviceInfo = getMsanInterceptor()->getDeviceInfo(Device);
+        for (size_t i = 0; i < NumOfDeviceGlobal; i++) {
+            const auto &GVInfo = GVInfos[i];
+            UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow(Queue, GVInfo.Addr, GVInfo.Size, 0));
+        }
+    }
+
+    return UR_RESULT_SUCCESS;
+}
+
 ur_result_t MsanInterceptor::insertContext(ur_context_handle_t Context,
                                            std::shared_ptr<ContextInfo> &CI) {
     std::scoped_lock<ur_shared_mutex> Guard(m_ContextMapMutex);

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

@@ -159,6 +159,11 @@ struct USMLaunchInfo {
     ur_result_t initialize();
 };
 
+struct DeviceGlobalInfo {
+    uptr Size;
+    uptr Addr;
+};
+
 struct SpirKernelInfo {
     uptr KernelName;
     uptr Size;
@@ -261,6 +266,7 @@ class MsanInterceptor {
                       std::shared_ptr<msan::DeviceInfo> &DeviceInfo);
 
     ur_result_t registerSpirKernels(ur_program_handle_t Program);
+    ur_result_t registerDeviceGlobals(ur_program_handle_t Program);
 
   private:
     std::unordered_map<ur_context_handle_t, std::shared_ptr<msan::ContextInfo>>

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

@@ -227,15 +227,10 @@ ur_result_t MsanShadowMemoryGPU::EnqueueMapShadow(
             VirtualMemMaps[MappedPtr].first = PhysicalMem;
         }
 
-        // We don't need to record virtual memory map for null pointer,
-        // since it doesn't have an alloc info.
-        if (Ptr == 0) {
-            continue;
+        auto AllocInfoItOp = getMsanInterceptor()->findAllocInfoByAddress(Ptr);
+        if (AllocInfoItOp) {
+            VirtualMemMaps[MappedPtr].second.insert((*AllocInfoItOp)->second);
         }
-
-        auto AllocInfoIt = getMsanInterceptor()->findAllocInfoByAddress(Ptr);
-        assert(AllocInfoIt);
-        VirtualMemMaps[MappedPtr].second.insert((*AllocInfoIt)->second);
     }
 
     return UR_RESULT_SUCCESS;