[DeviceSanitizer] Disable ASan when the ESIMD kernel is found (#14156)

ESIMD kernel doesn't support noinline functions, so we can't support
sanitizer for it now.

We can't check if the kernel is ESIMD kernel at Unified Runtime, so we
have to disable ASan completely when it found ESIMD kernel in device
image.

Author: AllanZyne

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -1278,7 +1278,8 @@ struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
 static void ExtendSpirKernelArgs(Module &M, FunctionAnalysisManager &FAM) {
   SmallVector<Function *> SpirFixupFuncs;
   for (Function &F : M) {
-    if (F.getCallingConv() == CallingConv::SPIR_KERNEL) {
+    if (F.getCallingConv() == CallingConv::SPIR_KERNEL &&
+        F.hasFnAttribute(Attribute::SanitizeAddress)) {
       SpirFixupFuncs.emplace_back(&F);
     }
   }
@@ -1294,9 +1295,7 @@ static void ExtendSpirKernelArgs(Module &M, FunctionAnalysisManager &FAM) {
       Types.push_back(I->getType());
     }
 
-    // New argument type: uintptr_t as(1)*, as it's allocated in USM buffer, and
-    // it can also be treated as a pointer point to the base address of private
-    // shadow memory
+    // New argument: uintptr_t as(1)*, which is allocated in shared USM buffer
     Types.push_back(IntptrTy->getPointerTo(kSpirOffloadGlobalAS));
 
     FunctionType *NewFTy = FunctionType::get(F->getReturnType(), Types, false);
@@ -1411,6 +1410,21 @@ PreservedAnalyses AddressSanitizerPass::run(Module &M,
       ClUseStackSafety ? &MAM.getResult<StackSafetyGlobalAnalysis>(M) : nullptr;
 
   if (Triple(M.getTargetTriple()).isSPIR()) {
+    bool HasESIMDKernel = false;
+
+    // ESIMD kernel doesn't support noinline functions, so we can't
+    // support sanitizer for it
+    for (Function &F : M)
+      if (F.hasMetadata("sycl_explicit_simd")) {
+        F.removeFnAttr(Attribute::SanitizeAddress);
+        HasESIMDKernel = true;
+      }
+
+    // FIXME: we can't check if the kernel is ESIMD kernel at UR, so we
+    // have to disable ASan completely in this case
+    if (HasESIMDKernel)
+      return PreservedAnalyses::all();
+
     ExtendSpirKernelArgs(M, FAM);
   }
 

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

@@ -0,0 +1,17 @@
+; RUN: opt < %s -passes=asan -asan-instrumentation-with-call-threshold=0 -asan-stack=0 -asan-globals=0 -S | FileCheck %s
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spir64-unknown-unknown"
+
+@__spirv_BuiltInGlobalInvocationId = external addrspace(1) constant <3 x i64>
+
+; Function Attrs: sanitize_address
+define spir_kernel void @esimd_kernel() #0 !sycl_explicit_simd !1 {
+entry:
+  %0 = load i64, ptr addrspace(1) getelementptr inbounds (i8, ptr addrspace(1) @__spirv_BuiltInGlobalInvocationId, i64 8), align 8
+  ret void
+}
+; CHECK-NOT: {{ sanitize_address }}
+
+attributes #0 = { sanitize_address }
+!1 = !{}