Merge pull request #1879 from AllanZyne/review/yang/fix_dsan_destruction

[DeviceSanitizer] Fix interceptor destruction order

Author: omarahmed1111

source/loader/layers/sanitizer/asan_interceptor.cpp

@@ -186,6 +186,18 @@ SanitizerInterceptor::~SanitizerInterceptor() {
     DestroyShadowMemoryOnCPU();
     DestroyShadowMemoryOnPVC();
     DestroyShadowMemoryOnDG2();
+
+    // We must release these objects before releasing adapters, since
+    // they may use the adapter in their destructor
+    m_Quarantine = nullptr;
+    m_MemBufferMap.clear();
+    m_AllocationMap.clear();
+    m_KernelMap.clear();
+    m_ContextMap.clear();
+
+    for (auto Adapter : m_Adapters) {
+        getContext()->urDdiTable.Global.pfnAdapterRelease(Adapter);
+    }
 }
 
 /// The memory chunk allocated from the underlying allocator looks like this:

source/loader/layers/sanitizer/asan_interceptor.hpp

@@ -22,6 +22,7 @@
 #include <optional>
 #include <queue>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 namespace ur_sanitizer_layer {
@@ -48,9 +49,8 @@ struct DeviceInfo {
     std::queue<std::shared_ptr<AllocInfo>> Quarantine;
     size_t QuarantineSize = 0;
 
-    // TODO: re-enable retaining and releasing device handles in DeviceInfo
-    // constructor/destructor. See PR
-    // https://github.com/oneapi-src/unified-runtime/pull/1883
+    // 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);
@@ -199,6 +199,16 @@ class SanitizerInterceptor {
     ur_result_t eraseMemBuffer(ur_mem_handle_t MemHandle);
     std::shared_ptr<MemBuffer> getMemBuffer(ur_mem_handle_t MemHandle);
 
+    ur_result_t holdAdapter(ur_adapter_handle_t Adapter) {
+        std::scoped_lock<ur_shared_mutex> Guard(m_AdaptersMutex);
+        if (m_Adapters.find(Adapter) != m_Adapters.end()) {
+            return UR_RESULT_SUCCESS;
+        }
+        UR_CALL(getContext()->urDdiTable.Global.pfnAdapterRetain(Adapter));
+        m_Adapters.insert(Adapter);
+        return UR_RESULT_SUCCESS;
+    }
+
     std::optional<AllocationIterator> findAllocInfoByAddress(uptr Address);
 
     std::shared_ptr<ContextInfo> getContextInfo(ur_context_handle_t Context) {
@@ -260,6 +270,9 @@ class SanitizerInterceptor {
 
     std::unique_ptr<Quarantine> m_Quarantine;
     logger::Logger &logger;
+
+    std::unordered_set<ur_adapter_handle_t> m_Adapters;
+    ur_shared_mutex m_AdaptersMutex;
 };
 
 } // namespace ur_sanitizer_layer

source/loader/layers/sanitizer/ur_sanddi.cpp

@@ -52,6 +52,38 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices,
 
 } // namespace
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urAdapterGet
+__urdlllocal ur_result_t UR_APICALL urAdapterGet(
+    uint32_t
+        NumEntries, ///< [in] the number of adapters to be added to phAdapters.
+    ///< If phAdapters is not NULL, then NumEntries should be greater than
+    ///< zero, otherwise ::UR_RESULT_ERROR_INVALID_SIZE,
+    ///< will be returned.
+    ur_adapter_handle_t *
+        phAdapters, ///< [out][optional][range(0, NumEntries)] array of handle of adapters.
+    ///< If NumEntries is less than the number of adapters available, then
+    ///< ::urAdapterGet shall only retrieve that number of platforms.
+    uint32_t *
+        pNumAdapters ///< [out][optional] returns the total number of adapters available.
+) {
+    auto pfnAdapterGet = getContext()->urDdiTable.Global.pfnAdapterGet;
+
+    if (nullptr == pfnAdapterGet) {
+        return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    }
+
+    ur_result_t result = pfnAdapterGet(NumEntries, phAdapters, pNumAdapters);
+    if (result == UR_RESULT_SUCCESS && phAdapters) {
+        const uint32_t NumAdapters = pNumAdapters ? *pNumAdapters : NumEntries;
+        for (uint32_t i = 0; i < NumAdapters; ++i) {
+            UR_CALL(getContext()->interceptor->holdAdapter(phAdapters[i]));
+        }
+    }
+
+    return result;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Intercept function for urUSMHostAlloc
 __urdlllocal ur_result_t UR_APICALL urUSMHostAlloc(
@@ -1324,6 +1356,36 @@ __urdlllocal ur_result_t UR_APICALL urKernelSetArgPointer(
     return result;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Exported function for filling application's Global table
+///        with current process' addresses
+///
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+///     - ::UR_RESULT_ERROR_UNSUPPORTED_VERSION
+__urdlllocal ur_result_t UR_APICALL urGetGlobalProcAddrTable(
+    ur_api_version_t version, ///< [in] API version requested
+    ur_global_dditable_t
+        *pDdiTable ///< [in,out] pointer to table of DDI function pointers
+) {
+    if (nullptr == pDdiTable) {
+        return UR_RESULT_ERROR_INVALID_NULL_POINTER;
+    }
+
+    if (UR_MAJOR_VERSION(ur_sanitizer_layer::getContext()->version) !=
+            UR_MAJOR_VERSION(version) ||
+        UR_MINOR_VERSION(ur_sanitizer_layer::getContext()->version) >
+            UR_MINOR_VERSION(version)) {
+        return UR_RESULT_ERROR_UNSUPPORTED_VERSION;
+    }
+
+    ur_result_t result = UR_RESULT_SUCCESS;
+
+    pDdiTable->pfnAdapterGet = ur_sanitizer_layer::urAdapterGet;
+
+    return result;
+}
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Exported function for filling application's Context table
 ///        with current process' addresses
@@ -1597,6 +1659,11 @@ ur_result_t context_t::init(ur_dditable_t *dditable,
 
     urDdiTable = *dditable;
 
+    if (UR_RESULT_SUCCESS == result) {
+        result = ur_sanitizer_layer::urGetGlobalProcAddrTable(
+            UR_API_VERSION_CURRENT, &dditable->Global);
+    }
+
     if (UR_RESULT_SUCCESS == result) {
         result = ur_sanitizer_layer::urGetContextProcAddrTable(
             UR_API_VERSION_CURRENT, &dditable->Context);

source/loader/ur_lib.cpp

@@ -57,7 +57,8 @@ void context_t::initLayers() const {
 }
 
 void context_t::tearDownLayers() const {
-    for (auto &[layer, destroy] : layers) {
+    for (auto it = layers.rbegin(); it != layers.rend(); ++it) {
+        auto [layer, destroy] = *it;
         layer->tearDown();
         destroy();
     }

test/layers/CMakeLists.txt

@@ -8,3 +8,7 @@ add_subdirectory(validation)
 if(UR_ENABLE_TRACING)
     add_subdirectory(tracing)
 endif()
+
+if(UR_ENABLE_SANITIZER)
+    add_subdirectory(sanitizer)
+endif()

test/layers/sanitizer/CMakeLists.txt

@@ -0,0 +1,37 @@
+# Copyright (C) 2023-2024 Intel Corporation
+# Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
+# See LICENSE.TXT
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+set(UR_SANITIZER_TEST_DIR ${CMAKE_CURRENT_SOURCE_DIR})
+set(SAN_TEST_PREFIX sanitizer_test)
+
+function(add_sanitizer_test_executable name)
+    add_ur_executable(${SAN_TEST_PREFIX}-${name}
+        ${ARGN})
+    target_link_libraries(${SAN_TEST_PREFIX}-${name}
+        PRIVATE
+        ${PROJECT_NAME}::loader
+        ${PROJECT_NAME}::headers
+        ${PROJECT_NAME}::testing
+        ${PROJECT_NAME}::mock
+        GTest::gtest_main)
+endfunction()
+
+function(set_sanitizer_test_properties name)
+    set_tests_properties(${name} PROPERTIES LABELS "sanitizer")
+    set_property(TEST ${name} PROPERTY ENVIRONMENT
+        "UR_LOG_SANITIZER=level:debug\;flush:debug\;output:stdout")
+endfunction()
+
+function(add_sanitizer_test name)
+    add_sanitizer_test_executable(${name} ${ARGN})
+
+    add_test(NAME ${name}
+        COMMAND ${SAN_TEST_PREFIX}-${name}
+        WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
+
+    set_sanitizer_test_properties(${name})
+endfunction()
+
+add_sanitizer_test(asan asan.cpp)

test/layers/sanitizer/asan.cpp

@@ -0,0 +1,58 @@
+/*
+ *
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
+ * See LICENSE.TXT
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ * @file asan.cpp
+ *
+ */
+
+#include <gtest/gtest.h>
+#include <ur_api.h>
+
+TEST(DeviceAsan, Initialization) {
+    ur_result_t status;
+
+    ur_loader_config_handle_t loaderConfig;
+    status = urLoaderConfigCreate(&loaderConfig);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+    status = urLoaderConfigEnableLayer(loaderConfig, "UR_LAYER_ASAN");
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    status = urLoaderInit(0, loaderConfig);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    ur_adapter_handle_t adapter;
+    status = urAdapterGet(1, &adapter, nullptr);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    ur_platform_handle_t platform;
+    status = urPlatformGet(&adapter, 1, 1, &platform, nullptr);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    ur_device_handle_t device;
+    status = urDeviceGet(platform, UR_DEVICE_TYPE_DEFAULT, 1, &device, nullptr);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    ur_context_handle_t context;
+    status = urContextCreate(1, &device, nullptr, &context);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    status = urContextRelease(context);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    status = urDeviceRelease(device);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    status = urAdapterRelease(adapter);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    status = urLoaderTearDown();
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+
+    status = urLoaderConfigRelease(loaderConfig);
+    ASSERT_EQ(status, UR_RESULT_SUCCESS);
+}