Merge pull request #2585 from yingcong-wu/yc-test-main/0120-on-shadow-per-type

[DeivceASAN] Make ShadowMemory one instance per type

Author: kbenzie

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

@@ -28,26 +28,36 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices,
                          const ur_device_handle_t *phDevices) {
     std::shared_ptr<ContextInfo> CI;
     UR_CALL(getAsanInterceptor()->insertContext(Context, CI));
-    for (uint32_t i = 0; i < numDevices; ++i) {
-        auto hDevice = phDevices[i];
-        std::shared_ptr<DeviceInfo> DI;
-        UR_CALL(getAsanInterceptor()->insertDevice(hDevice, DI));
-        DI->Type = GetDeviceType(Context, hDevice);
-        if (DI->Type == DeviceType::UNKNOWN) {
-            getContext()->logger.error("Unsupport device");
-            return UR_RESULT_ERROR_INVALID_DEVICE;
-        }
-        getContext()->logger.info(
-            "DeviceInfo {} (Type={}, IsSupportSharedSystemUSM={})",
-            (void *)DI->Handle, ToString(DI->Type),
-            DI->IsSupportSharedSystemUSM);
-        getContext()->logger.info("Add {} into context {}", (void *)DI->Handle,
-                                  (void *)Context);
-        if (!DI->Shadow) {
-            UR_CALL(DI->allocShadowMemory(Context));
+
+    if (numDevices > 0) {
+        auto DeviceType = GetDeviceType(Context, phDevices[0]);
+        auto ShadowMemory = getAsanInterceptor()->getOrCreateShadowMemory(
+            phDevices[0], DeviceType);
+
+        for (uint32_t i = 0; i < numDevices; ++i) {
+            auto hDevice = phDevices[i];
+            std::shared_ptr<DeviceInfo> DI;
+            UR_CALL(getAsanInterceptor()->insertDevice(hDevice, DI));
+            DI->Type = GetDeviceType(Context, hDevice);
+            if (DI->Type == DeviceType::UNKNOWN) {
+                getContext()->logger.error("Unsupport device");
+                return UR_RESULT_ERROR_INVALID_DEVICE;
+            }
+            if (DI->Type != DeviceType) {
+                getContext()->logger.error(
+                    "Different device type in the same context");
+                return UR_RESULT_ERROR_INVALID_DEVICE;
+            }
+            getContext()->logger.info(
+                "DeviceInfo {} (Type={}, IsSupportSharedSystemUSM={})",
+                (void *)DI->Handle, ToString(DI->Type),
+                DI->IsSupportSharedSystemUSM);
+            getContext()->logger.info("Add {} into context {}",
+                                      (void *)DI->Handle, (void *)Context);
+            DI->Shadow = ShadowMemory;
+            CI->DeviceList.emplace_back(hDevice);
+            CI->AllocInfosMap[hDevice];
         }
-        CI->DeviceList.emplace_back(hDevice);
-        CI->AllocInfosMap[hDevice];
     }
     return UR_RESULT_SUCCESS;
 }

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

@@ -36,8 +36,7 @@ AsanInterceptor::~AsanInterceptor() {
     // We must release these objects before releasing adapters, since
     // they may use the adapter in their destructor
     for (const auto &[_, DeviceInfo] : m_DeviceMap) {
-        [[maybe_unused]] auto URes = DeviceInfo->Shadow->Destory();
-        assert(URes == UR_RESULT_SUCCESS);
+        DeviceInfo->Shadow = nullptr;
     }
 
     m_Quarantine = nullptr;
@@ -48,6 +47,11 @@ AsanInterceptor::~AsanInterceptor() {
     // detection depends on it.
     m_AllocationMap.clear();
 
+    for (auto &[_, ShadowMemory] : m_ShadowMap) {
+        ShadowMemory->Destory();
+        getContext()->urDdiTable.Context.pfnRelease(ShadowMemory->Context);
+    }
+
     for (auto Adapter : m_Adapters) {
         getContext()->urDdiTable.Global.pfnAdapterRelease(Adapter);
     }
@@ -301,14 +305,24 @@ ur_result_t AsanInterceptor::postLaunchKernel(ur_kernel_handle_t Kernel,
     return Result;
 }
 
-ur_result_t DeviceInfo::allocShadowMemory(ur_context_handle_t Context) {
-    Shadow = GetShadowMemory(Context, Handle, Type);
-    assert(Shadow && "Failed to get shadow memory");
-    UR_CALL(Shadow->Setup());
-    getContext()->logger.info("ShadowMemory(Global): {} - {}",
-                              (void *)Shadow->ShadowBegin,
-                              (void *)Shadow->ShadowEnd);
-    return UR_RESULT_SUCCESS;
+std::shared_ptr<ShadowMemory>
+AsanInterceptor::getOrCreateShadowMemory(ur_device_handle_t Device,
+                                         DeviceType Type) {
+    std::scoped_lock<ur_shared_mutex> Guard(m_ShadowMapMutex);
+    if (m_ShadowMap.find(Type) == m_ShadowMap.end()) {
+        ur_context_handle_t InternalContext;
+        auto Res = getContext()->urDdiTable.Context.pfnCreate(
+            1, &Device, nullptr, &InternalContext);
+        if (Res != UR_RESULT_SUCCESS) {
+            getContext()->logger.error("Failed to create shadow context");
+            return nullptr;
+        }
+        std::shared_ptr<ContextInfo> CI;
+        insertContext(InternalContext, CI);
+        m_ShadowMap[Type] = GetShadowMemory(InternalContext, Device, Type);
+        m_ShadowMap[Type]->Setup();
+    }
+    return m_ShadowMap[Type];
 }
 
 /// Each 8 bytes of application memory are mapped into one byte of shadow memory

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

@@ -56,8 +56,6 @@ struct DeviceInfo {
     // Device handles are special and alive in the whole process lifetime,
     // so we needn't retain&release here.
     explicit DeviceInfo(ur_device_handle_t Device) : Handle(Device) {}
-
-    ur_result_t allocShadowMemory(ur_context_handle_t Context);
 };
 
 struct QueueInfo {
@@ -359,6 +357,9 @@ class AsanInterceptor {
 
     bool isNormalExit() { return m_NormalExit; }
 
+    std::shared_ptr<ShadowMemory>
+    getOrCreateShadowMemory(ur_device_handle_t Device, DeviceType Type);
+
   private:
     ur_result_t updateShadowMemory(std::shared_ptr<ContextInfo> &ContextInfo,
                                    std::shared_ptr<DeviceInfo> &DeviceInfo,
@@ -375,9 +376,6 @@ class AsanInterceptor {
                               ur_kernel_handle_t Kernel,
                               LaunchInfo &LaunchInfo);
 
-    ur_result_t allocShadowMemory(ur_context_handle_t Context,
-                                  std::shared_ptr<DeviceInfo> &DeviceInfo);
-
     ur_result_t registerDeviceGlobals(ur_program_handle_t Program);
     ur_result_t registerSpirKernels(ur_program_handle_t Program);
 
@@ -413,6 +411,9 @@ class AsanInterceptor {
     ur_shared_mutex m_AdaptersMutex;
 
     bool m_NormalExit = true;
+
+    std::unordered_map<DeviceType, std::shared_ptr<ShadowMemory>> m_ShadowMap;
+    ur_shared_mutex m_ShadowMapMutex;
 };
 
 } // namespace asan

source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.cpp

@@ -154,7 +154,10 @@ DeviceType GetDeviceType(ur_context_handle_t Context,
         // FIXME: There's no API querying the address bits of device, so we guess it by the
         // value of device USM pointer (see "USM Allocation Range" in asan_shadow.cpp)
         auto Type = DeviceType::UNKNOWN;
-        if (Ptr >> 48 == 0xff00U) {
+
+        // L0 changes their VA layout.
+        // TODO: update our shadow memory layout/algorithms to accordingly.
+        if (((Ptr >> 52) & 0xff0U) == 0xff0U) {
             Type = DeviceType::GPU_PVC;
         } else {
             Type = DeviceType::GPU_DG2;