Merge branch 'main' into do-alloc-use-pool

Author: zhaomaosu

CMakeLists.txt

@@ -27,7 +27,7 @@ option(UR_USE_ASAN "enable AddressSanitizer" OFF)
 option(UR_USE_UBSAN "enable UndefinedBehaviorSanitizer" OFF)
 option(UR_USE_MSAN "enable MemorySanitizer" OFF)
 option(UR_USE_TSAN "enable ThreadSanitizer" OFF)
-option(UR_USE_CFI "enable Control Flow Integrity checks (requires clang and implies -flto)" ON)
+option(UR_USE_CFI "enable Control Flow Integrity checks (requires clang and implies -flto)" OFF)
 option(UR_ENABLE_TRACING "enable api tracing through xpti" OFF)
 option(UR_ENABLE_SANITIZER "enable device sanitizer" ON)
 option(UR_ENABLE_SYMBOLIZER "enable symoblizer for sanitizer" OFF)

README.md

@@ -130,7 +130,7 @@ List of options provided by CMake:
 | UR_USE_TSAN | Enable ThreadSanitizer | ON/OFF | OFF |
 | UR_USE_UBSAN | Enable UndefinedBehavior Sanitizer | ON/OFF | OFF |
 | UR_USE_MSAN | Enable MemorySanitizer (clang only) | ON/OFF | OFF |
-| UR_USE_CFI | Enable Control Flow Integrity checks (clang only, also enables lto) | ON/OFF | ON |
+| UR_USE_CFI | Enable Control Flow Integrity checks (clang only, also enables lto) | ON/OFF | OFF |
 | UR_ENABLE_TRACING | Enable XPTI-based tracing layer | ON/OFF | OFF |
 | UR_ENABLE_SANITIZER | Enable device sanitizer layer | ON/OFF | ON |
 | UR_CONFORMANCE_TARGET_TRIPLES | SYCL triples to build CTS device binaries for | Comma-separated list | spir64 |

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

@@ -55,6 +55,9 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices,
 bool isInstrumentedKernel(ur_kernel_handle_t hKernel) {
     auto hProgram = GetProgram(hKernel);
     auto PI = getAsanInterceptor()->getProgramInfo(hProgram);
+    if (PI == nullptr) {
+        return false;
+    }
     return PI->isKernelInstrumented(hKernel);
 }
 
@@ -290,8 +293,9 @@ __urdlllocal ur_result_t UR_APICALL urProgramRetain(
     UR_CALL(pfnRetain(hProgram));
 
     auto ProgramInfo = getAsanInterceptor()->getProgramInfo(hProgram);
-    UR_ASSERT(ProgramInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
-    ProgramInfo->RefCount++;
+    if (ProgramInfo != nullptr) {
+        ProgramInfo->RefCount++;
+    }
 
     return UR_RESULT_SUCCESS;
 }
@@ -364,6 +368,7 @@ __urdlllocal ur_result_t UR_APICALL urProgramLink(
 
     UR_CALL(pfnProgramLink(hContext, count, phPrograms, pOptions, phProgram));
 
+    UR_CALL(getAsanInterceptor()->insertProgram(*phProgram));
     UR_CALL(getAsanInterceptor()->registerProgram(*phProgram));
 
     return UR_RESULT_SUCCESS;
@@ -395,6 +400,7 @@ ur_result_t UR_APICALL urProgramLinkExp(
     UR_CALL(pfnProgramLinkExp(hContext, numDevices, phDevices, count,
                               phPrograms, pOptions, phProgram));
 
+    UR_CALL(getAsanInterceptor()->insertProgram(*phProgram));
     UR_CALL(getAsanInterceptor()->registerProgram(*phProgram));
 
     return UR_RESULT_SUCCESS;
@@ -417,8 +423,7 @@ ur_result_t UR_APICALL urProgramRelease(
     UR_CALL(pfnProgramRelease(hProgram));
 
     auto ProgramInfo = getAsanInterceptor()->getProgramInfo(hProgram);
-    UR_ASSERT(ProgramInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
-    if (--ProgramInfo->RefCount == 0) {
+    if (ProgramInfo != nullptr && --ProgramInfo->RefCount == 0) {
         UR_CALL(getAsanInterceptor()->unregisterProgram(hProgram));
         UR_CALL(getAsanInterceptor()->eraseProgram(hProgram));
     }

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

@@ -226,15 +226,18 @@ ur_result_t AsanInterceptor::releaseMemory(ur_context_handle_t Context,
     if (ReleaseList.size()) {
         std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
         for (auto &It : ReleaseList) {
+            auto ToFreeAllocInfo = It->second;
             getContext()->logger.info("Quarantine Free: {}",
-                                      (void *)It->second->AllocBegin);
+                                      (void *)ToFreeAllocInfo->AllocBegin);
 
-            ContextInfo->Stats.UpdateUSMRealFreed(AllocInfo->AllocSize,
-                                                  AllocInfo->getRedzoneSize());
+            ContextInfo->Stats.UpdateUSMRealFreed(
+                ToFreeAllocInfo->AllocSize, ToFreeAllocInfo->getRedzoneSize());
 
-            m_AllocationMap.erase(It);
             UR_CALL(getContext()->urDdiTable.USM.pfnFree(
-                Context, (void *)(It->second->AllocBegin)));
+                Context, (void *)(ToFreeAllocInfo->AllocBegin)));
+
+            // Erase it at last to avoid use-after-free.
+            m_AllocationMap.erase(It);
         }
     }
     ContextInfo->Stats.UpdateUSMFreed(AllocInfo->AllocSize);
@@ -426,6 +429,7 @@ ur_result_t AsanInterceptor::registerProgram(ur_program_handle_t Program) {
 
 ur_result_t AsanInterceptor::unregisterProgram(ur_program_handle_t Program) {
     auto ProgramInfo = getProgramInfo(Program);
+    assert(ProgramInfo != nullptr && "unregistered program!");
 
     ProgramInfo->InstrumentedKernels.clear();
 
@@ -464,6 +468,7 @@ ur_result_t AsanInterceptor::registerSpirKernels(ur_program_handle_t Program) {
         }
 
         auto PI = getProgramInfo(Program);
+        assert(PI != nullptr && "unregistered program!");
         for (const auto &SKI : SKInfo) {
             if (SKI.Size == 0) {
                 continue;
@@ -500,6 +505,7 @@ AsanInterceptor::registerDeviceGlobals(ur_program_handle_t Program) {
     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);

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

@@ -339,8 +339,10 @@ class AsanInterceptor {
 
     std::shared_ptr<ProgramInfo> getProgramInfo(ur_program_handle_t Program) {
         std::shared_lock<ur_shared_mutex> Guard(m_ProgramMapMutex);
-        assert(m_ProgramMap.find(Program) != m_ProgramMap.end());
-        return m_ProgramMap[Program];
+        if (m_ProgramMap.find(Program) != m_ProgramMap.end()) {
+            return m_ProgramMap[Program];
+        }
+        return nullptr;
     }
 
     std::shared_ptr<KernelInfo> getKernelInfo(ur_kernel_handle_t Kernel) {