[HIP] Preserve ASAN bitcode library functions

[yaxunl]

Summary

Address sanitizer passes may generate call of ASAN bitcode library
functions after bitcode linking in lld, therefore lld cannot add
those symbols since it does not know they will be used later.

To solve this issue, clang emits a reference to a bicode library
function which calls all ASAN functions which need to be
preserved. This basically force all ASAN functions to be
linked in.

[tra]

LGTM in general.

One question -- does it have to be a function calling other functions just for the sake of preserving them?
Can it be a flat array of pointers to the functions you need to keep around?

[yaxunl]

In D106315#2898536, @tra wrote:

LGTM in general.

One question -- does it have to be a function calling other functions just for the sake of preserving them?
Can it be a flat array of pointers to the functions you need to keep around?

Yes that's possible. However that would require FE to know these functions and declare them, whereas the current approach leave the concern to the device library.

[tra]

In D106315#2899928, @yaxunl wrote:

Yes that's possible. However that would require FE to know these functions and declare them, whereas the current approach leave the concern to the device library.

I was thinking of just adding a pointer to an array of pointers to @llvm.compiler.used.
The array itself would come from the bitcode library and would be populated there. I'm not sure if it's doable without knowing the array size, though.

e.g

; Added by compiler
@llvm.compiler.used = appending global [1 x i8*] [i8* bitcast ([16 x i8*]* @_ZL5funcs to i8*)], section "llvm.metadata"

; comes from the bitcode library, initialized with the pointers to functions you need to keep.
@_ZL5funcs = internal global [16 x i8*] zeroinitializer, align 4

It should be possible to make the pointer to the array opaque with an intermadiate variable in the library.

[yaxunl]

In D106315#2900882, @tra wrote:

In D106315#2899928, @yaxunl wrote:

Yes that's possible. However that would require FE to know these functions and declare them, whereas the current approach leave the concern to the device library.

I was thinking of just adding a pointer to an array of pointers to @llvm.compiler.used.
The array itself would come from the bitcode library and would be populated there. I'm not sure if it's doable without knowing the array size, though.

e.g

; Added by compiler
@llvm.compiler.used = appending global [1 x i8*] [i8* bitcast ([16 x i8*]* @_ZL5funcs to i8*)], section "llvm.metadata"

; comes from the bitcode library, initialized with the pointers to functions you need to keep.
@_ZL5funcs = internal global [16 x i8*] zeroinitializer, align 4

It should be possible to make the pointer to the array opaque with an intermadiate variable in the library.

Yes that is possible. Actually it can be simpler. We can add __attribute__((used)) to the dummy function in the bitcode lib, and it will be kept when linked by -mlink-builtin-bitcode.

However, it has a drawback compared to the current approach. It has no control when to keep the dummy function, whereas the current approach only keep it for -fsanitize=asan.

clang/lib/CodeGen/CodeGenModule.cpp

if (CodeGenOpts.SanitizeCfiCrossDso) {
     CodeGenFunction(*this).EmitCfiCheckFail();
     CodeGenFunction(*this).EmitCfiCheckStub();
   }
   emitAtAvailableLinkGuard();
   if (Context.getTargetInfo().getTriple().isWasm() &&
       !Context.getTargetInfo().getTriple().isOSEmscripten()) {
     EmitMainVoidAlias();
   }
+
+  // Emit reference of __amdgpu_device_library_preserve_asan_functions to
+  // preserve ASAN functions in bitcode libraries.
+  if (LangOpts.Sanitize.has(SanitizerKind::Address) && getTriple().isAMDGPU()) {
+    auto *FT = llvm::FunctionType::get(VoidTy, {});
+    auto *F = llvm::Function::Create(
+        FT, llvm::GlobalValue::ExternalLinkage,
+        "__amdgpu_device_library_preserve_asan_functions", &getModule());
+    auto *Var = new llvm::GlobalVariable(
+        getModule(), FT->getPointerTo(),
+        /*isConstant=*/true, llvm::GlobalValue::WeakAnyLinkage, F,
+        "__amdgpu_device_library_preserve_asan_functions_ptr", nullptr,
+        llvm::GlobalVariable::NotThreadLocal);
+    addCompilerUsedGlobal(Var);
+  }
+
   emitLLVMUsed();
   if (SanStats)
     SanStats->finish();
 
   if (CodeGenOpts.Autolink &&
       (Context.getLangOpts().Modules || !LinkerOptionsMetadata.empty())) {
     EmitModuleLinkOptions();
   }

clang/lib/Driver/Driver.cpp

getDeviceDependences(OffloadAction::DeviceDependences &DA,
       if (!Relocatable && CurPhase == phases::Backend && !EmitLLVM &&
           !EmitAsm) {
         // If we are in backend phase, we attempt to generate the fat binary.
         // We compile each arch to IR and use a link action to generate code
         // object containing ISA. Then we use a special "link" action to create
         // a fat binary containing all the code objects for different GPU's.
         // The fat binary is then an input to the host action.
         for (unsigned I = 0, E = GpuArchList.size(); I != E; ++I) {
-          if (GPUSanitize || C.getDriver().isUsingLTO(/*IsOffload=*/true)) {
-            // When GPU sanitizer is enabled, since we need to link in the
-            // the sanitizer runtime library after the sanitize pass, we have
-            // to skip the backend and assemble phases and use lld to link
-            // the bitcode. The same happens if users request to use LTO
-            // explicitly.
+          if (C.getDriver().isUsingLTO(/*IsOffload=*/true)) {
+            // When LTO is enabled, skip the backend and assemble phases and
+            // use lld to link the bitcode.
             ActionList AL;
             AL.push_back(CudaDeviceActions[I]);
             // Create a link action to link device IR with device library
             // and generate ISA.
             CudaDeviceActions[I] =
                 C.MakeAction<LinkJobAction>(AL, types::TY_Image);
           } else {
-            // When GPU sanitizer is not enabled, we follow the conventional
+            // When LTO is not enabled, we follow the conventional
             // compiler phases, including backend and assemble phases.
             ActionList AL;
             auto BackendAction = C.getDriver().ConstructPhaseAction(
                 C, Args, phases::Backend, CudaDeviceActions[I],
                 AssociatedOffloadKind);
             auto AssembleAction = C.getDriver().ConstructPhaseAction(
                 C, Args, phases::Assemble, BackendAction,
                 AssociatedOffloadKind);

clang/test/CodeGenCUDA/amdgpu-asan.cu

+// RUN: %clang_cc1 %s -emit-llvm -o - -triple=amdgcn-amd-amdhsa \
+// RUN:   -fcuda-is-device -target-cpu gfx906 -fsanitize=address \
+// RUN:   -x hip | FileCheck -check-prefix=ASAN %s
+
+// RUN: %clang_cc1 %s -emit-llvm -o - -triple=amdgcn-amd-amdhsa \
+// RUN:   -fcuda-is-device -target-cpu gfx906 -x hip \
+// RUN:   | FileCheck %s
+
+// REQUIRES: amdgpu-registered-target
+
+// ASAN-DAG: declare void @__amdgpu_device_library_preserve_asan_functions()
+// ASAN-DAG: @__amdgpu_device_library_preserve_asan_functions_ptr = weak addrspace(1) constant void ()* @__amdgpu_device_library_preserve_asan_functions
+// ASAN-DAG: @llvm.compiler.used = {{.*}}@__amdgpu_device_library_preserve_asan_functions_ptr
+
+// CHECK-NOT: @__amdgpu_device_library_preserve_asan_functions_ptr

clang/test/Driver/hip-sanitize-options.hip

 // RUN:   -fsanitize=address -fgpu-sanitize \
 // RUN:   -nogpuinc --rocm-path=%S/Inputs/rocm \
 // RUN:   %s 2>&1 | FileCheck -check-prefix=XNACK %s
 
 // CHECK-NOT: {{"[^"]*clang[^"]*".* "-fcuda-is-device".* "-fsanitize=address"}}
 // CHECK-NOT: {{"[^"]*lld(\.exe){0,1}".* ".*hip.bc"}}
 // CHECK: {{"[^"]*clang[^"]*".* "-triple" "x86_64-unknown-linux-gnu".* "-fsanitize=address"}}
 
-// NORDC: {{"[^"]*clang[^"]*".* "-fcuda-is-device".* "-fsanitize=address".*}} "-o" "[[OUT:[^"]*.bc]]"
+// NORDC: {{"[^"]*clang[^"]*".* "-emit-obj".* "-fcuda-is-device".* "-fsanitize=address".*}} "-o" "[[OUT:[^"]*.o]]"
 // NORDC: {{"[^"]*lld(\.exe){0,1}".*}} "[[OUT]]" {{".*asanrtl.bc" ".*hip.bc"}}
 // NORDC: {{"[^"]*clang[^"]*".* "-triple" "x86_64-unknown-linux-gnu".* "-fsanitize=address"}}
 
 // RDC: {{"[^"]*clang[^"]*".* "-triple" "x86_64-unknown-linux-gnu".* "-fsanitize=address"}}
 // RDC: {{"[^"]*clang[^"]*".* "-emit-llvm-bc".* "-fcuda-is-device".* "-fsanitize=address".*}} "-o" "[[OUT:[^"]*.bc]]"
 // RDC: {{"[^"]*lld(\.exe){0,1}".*}} "[[OUT]]" {{".*asanrtl.bc" ".*hip.bc"}}
 
 // FAIL: AMDGPU address sanitizer runtime library (asanrtl) is not found. Please install ROCm device library which supports address sanitizer
 
 // XNACK: error: '-fgpu-sanitize' is not compatible with offload arch 'gfx900:xnack-'. Use an offload arch without 'xnack-' instead