[DeviceASAN] Fix ASAN with kernel assert (#16256)

UR: oneapi-src/unified-runtime#2415

---------

Co-authored-by: Martin Morrison-Grant <[email protected]>

Author: AllanZyne

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -1333,13 +1333,27 @@ static void ExtendSpirKernelArgs(Module &M, FunctionAnalysisManager &FAM,
 
   if (!HasESIMD)
     for (Function &F : M) {
-      if (F.getCallingConv() != CallingConv::SPIR_KERNEL)
-        continue;
-
       if (!F.hasFnAttribute(Attribute::SanitizeAddress) ||
           F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation))
         continue;
 
+      if (F.getName().contains("__sycl_service_kernel__")) {
+        F.addFnAttr(Attribute::DisableSanitizerInstrumentation);
+        continue;
+      }
+
+      // Skip referenced-indirectly function as we insert access to shared
+      // local memory (SLM) __AsanLaunchInfo and access to SLM in
+      // referenced-indirectly function isn't supported yet in
+      // intel-graphics-compiler.
+      if (F.hasFnAttribute("referenced-indirectly")) {
+        F.addFnAttr(Attribute::DisableSanitizerInstrumentation);
+        continue;
+      }
+
+      if (F.getCallingConv() != CallingConv::SPIR_KERNEL)
+        continue;
+
       SpirFixupKernels.emplace_back(&F);
 
       auto KernelName = F.getName();
@@ -3685,16 +3699,6 @@ bool AddressSanitizer::instrumentFunction(Function &F,
   if (F.isPresplitCoroutine())
     return false;
 
-  if (TargetTriple.isSPIROrSPIRV()) {
-    if (F.getName().contains("__sycl_service_kernel__"))
-      return false;
-    // Skip referenced-indirectly function as we insert access to shared local
-    // memory (SLM) __AsanLaunchInfo and access to SLM in referenced-indirectly
-    // function isn't supported yet in intel-graphics-compiler.
-    if (F.hasFnAttribute("referenced-indirectly"))
-      return false;
-  }
-
   bool FunctionModified = false;
 
   // Do not apply any instrumentation for naked functions.

llvm/test/Instrumentation/AddressSanitizer/SPIRV/skip_referenced_indirectly.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes=asan -asan-instrumentation-with-call-threshold=0 | FileCheck %s
+; RUN: opt < %s -passes=asan -asan-instrumentation-with-call-threshold=0 -S | FileCheck %s
 
 ; Check referenced-indirectly function isn't instrumented.
 
@@ -9,6 +9,7 @@ target triple = "spir64-unknown-unknown"
 @_ZTV8Derived1 = linkonce_odr addrspace(1) constant %structtype { [3 x ptr addrspace(4)] [ptr addrspace(4) null, ptr addrspace(4) null, ptr addrspace(4) addrspacecast (ptr @_ZN8Derived17displayEv to ptr addrspace(4))] }, align 8, !spirv.Decorations !0
 
 define linkonce_odr spir_func i32 @_ZN8Derived17displayEv(ptr addrspace(4) align 8 %this) sanitize_address "referenced-indirectly" {
+; CHECK: @_ZN8Derived17displayEv{{.*}}#1
 entry:
 ; CHECK-NOT: call void @__asan_load
 
@@ -17,6 +18,8 @@ entry:
   ret i32 %1
 }
 
+; CHECK: #1 {{.*}} disable_sanitizer_instrumentation
+
 !0 = !{!1, !2, !3}
 !1 = !{i32 22}
 !2 = !{i32 41, !"_ZTV8Derived1", i32 2}

llvm/test/Instrumentation/AddressSanitizer/SPIRV/skip_sycl_service_kernel.ll

@@ -0,0 +1,19 @@
+; RUN: opt < %s -passes=asan -asan-instrumentation-with-call-threshold=0 -S | FileCheck %s
+
+; Check "sycl_service_kernel" function isn't instrumented.
+
+target triple = "spir64-unknown-unknown"
+
+%structtype = type { [3 x ptr addrspace(4)] }
+%class.Base = type <{ ptr addrspace(4), i32, [4 x i8] }>
+
+define linkonce_odr spir_func i32 @_ZTSN4sycl3_V16detail23__sycl_service_kernel__16AssertInfoCopierE(ptr addrspace(4) align 8 %this) sanitize_address "referenced-indirectly" {
+; CHECK: @_ZTSN4sycl3_V16detail23__sycl_service_kernel__16AssertInfoCopierE{{.*}}#1
+entry:
+; CHECK-NOT: call void @__asan_load
+  %base_data = getelementptr inbounds %class.Base, ptr addrspace(4) %this, i64 0, i32 1
+  %1 = load i32, ptr addrspace(4) %base_data, align 8
+  ret i32 %1
+}
+
+; CHECK: #1 {{.*}} disable_sanitizer_instrumentation

sycl/cmake/modules/UnifiedRuntimeTag.cmake

@@ -1,7 +1,7 @@
-# commit 39df0317814c164f5242eda8d6f08550f6268492
-# Merge: 68d93efd be27d8f0
-# Author: Kenneth Benzie (Benie) <k.benzie@codeplay.com>
-# Date:   Mon Dec 16 13:53:13 2024 +0000
-#     Merge pull request #2467 from nrspruit/fix_external_import_function_call
-#     [L0] Fix external semaphore import function calls to match the header
-set(UNIFIED_RUNTIME_TAG 39df0317814c164f5242eda8d6f08550f6268492)
+# commit d18d52393aadf0083a32912096baaac558378a99
+# Merge: c45de9a5f7bf 05f94a8ab2a9
+# Author: Martin Grant <martin.morrisongrant@codeplay.com>
+# Date:   Wed Dec 18 15:01:30 2024 +0000
+#     Merge pull request #2415 from AllanZyne/review/yang/fix_metadata_assert
+#     [DeviceASAN] Fix ASAN with kernel assert
+set(UNIFIED_RUNTIME_TAG d18d52393aadf0083a32912096baaac558378a99)

sycl/test-e2e/AddressSanitizer/common/ignorelist.txt

@@ -1 +1 @@
-fun:*MyKernel*
+fun:*NoSanitized*

sycl/test-e2e/AddressSanitizer/common/kernel-filter.cpp

@@ -1,27 +1,51 @@
-// REQUIRES: linux
-// RUN: %{build} %device_asan_flags -O2 -fsanitize-ignorelist=%p/ignorelist.txt -o %t
-// RUN: %{run} %t 2>&1 | FileCheck %s
-// RUN: %{build} %device_asan_flags %if cpu %{ -fsycl-targets=spir64_x86_64 %} %if gpu %{ -fsycl-targets=spir64_gen -Xsycl-target-backend=spir64_gen %gpu_aot_target_opts %} -O2 -fsanitize-ignorelist=%p/ignorelist.txt -o %t2
+// REQUIRES: linux, cpu || (gpu && level_zero)
+// RUN: %{build} %device_asan_flags -O2 -fsanitize-ignorelist=%p/ignorelist.txt -o %t1
+// RUN: %{run} %t1 2>&1 | FileCheck %s
+// RUN: %{build} %device_asan_aot_flags -O2 -fsanitize-ignorelist=%p/ignorelist.txt -o %t2
 // RUN: %{run} %t2 2>&1 | FileCheck %s
 
 #include <sycl/detail/core.hpp>
 #include <sycl/usm.hpp>
 
 int main() {
+  constexpr std::size_t N = 8;
+  constexpr std::size_t group_size = 4;
+
   sycl::queue Q;
-  constexpr std::size_t N = 16;
+
   auto *array = sycl::malloc_device<int>(N, Q);
 
+  std::vector<int> v(N);
+  sycl::buffer<int, 1> buf(v.data(), v.size());
+
   Q.submit([&](sycl::handler &h) {
-    h.parallel_for<class MyKernel>(
-        sycl::nd_range<1>(N + 1, 1),
-        [=](sycl::nd_item<1> item) { ++array[item.get_global_id(0)]; });
+    auto buf_acc = buf.get_access<sycl::access::mode::read_write>(h);
+    auto loc_acc = sycl::local_accessor<int>(group_size, h);
+    h.parallel_for<class NoSanitized>(
+        sycl::nd_range<1>(N, group_size), [=](sycl::nd_item<1> item) {
+          auto gid = item.get_global_id(0);
+          auto lid = item.get_local_id(0);
+          array[gid] = buf_acc[gid] + loc_acc[lid];
+        });
+  });
+  Q.wait();
+  // CHECK-NOT: ERROR: DeviceSanitizer: out-of-bounds-access
+
+  Q.submit([&](sycl::handler &h) {
+    auto buf_acc = buf.get_access<sycl::access::mode::read_write>(h);
+    auto loc_acc = sycl::local_accessor<int>(group_size, h);
+    h.parallel_for<class Sanitized>(sycl::nd_range<1>(N, group_size),
+                                    [=](sycl::nd_item<1> item) {
+                                      auto gid = item.get_global_id(0);
+                                      auto lid = item.get_local_id(0);
+                                      array[gid] = buf_acc[gid] + loc_acc[lid];
+                                    });
   });
   Q.wait();
 
   sycl::free(array, Q);
   std::cout << "PASS" << std::endl;
+  // CHECK: PASS
+
   return 0;
 }
-
-// CHECK: PASS

sycl/test-e2e/AddressSanitizer/lit.local.cfg

@@ -1,14 +1,12 @@
-# This test assumes it can only run in CPU/PVC/DG2 devices, which support usm_device_allocations aspect
-
 config.substitutions.append(
     ("%device_asan_flags", "-Xarch_device -fsanitize=address")
 )
-
+config.substitutions.append(
+    ("%device_asan_aot_flags", "-Xarch_device -fsanitize=address %if cpu %{ -fsycl-targets=spir64_x86_64 %} %if gpu %{ -fsycl-targets=spir64_gen -Xsycl-target-backend=spir64_gen %gpu_aot_target_opts %}")
+)
 config.substitutions.append(
     ("%force_device_asan_rt", "env UR_ENABLE_LAYERS=UR_LAYER_ASAN")
 )
 
-config.unsupported_features += ['cuda', 'hip']
-
 # https://github.com/intel/llvm/issues/15953
 config.unsupported_features += ['gpu-intel-gen12']