[ASan][Offload] Add support to utlize device malloc/free.

Dynamic memory allocation/deallocation support using api's malloc/new or
free/delete must inherently invoke __asan* friendly api's when ASan is enabled.

Change-Id: I5732428319194b458f3fb36daa05709317990e63

Author: ampandey-AMD

offload/DeviceRTL/CMakeLists.txt

@@ -149,8 +149,9 @@ function(compileDeviceRTLLibrary target_name target_triple)
       VERBATIM
     )
 
-    if(${outfile} MATCHES "State.cpp")
+    if(${outfile} MATCHES "State.cpp" OR SANITIZER_AMDGPU)
       # Run the prep tool on the library to replace internal attribute with linkonce_odr for dm_alloc only.
+      # If SANITIZER_AMDGPU is enabled then run prep tool for all device functions.
       set(outfile_prep "${outfile}-${target_name}-prep.bc")
       add_custom_target(${outfile_prep}
           COMMAND ${PREP_TOOL} -dm ${outfile}
@@ -346,6 +347,13 @@ set(amd_options -Xclang -mcode-object-version=none
                 -Xclang -mcode-object-version=none
 )
 
+if(SANITIZER_AMDGPU)
+get_target_property(asanrtl_path asanrtl IMPORTED_LOCATION)
+list(APPEND amd_options -DSANITIZER_AMDGPU=1)
+list(APPEND amd_options -Xclang -mlink-bitcode-file -Xclang ${asanrtl_path})
+endif()
+
+
 compileDeviceRTLLibrary(amdgpu amdgcn-amd-amdhsa ${amd_options})
 compileDeviceRTLLibrary(nvptx nvptx64-nvidia-cuda --cuda-feature=+ptx63)
 

offload/DeviceRTL/src/State.cpp

@@ -24,8 +24,8 @@
 using namespace ompx;
 
 #pragma omp begin declare target device_type(host)
-void *internal_malloc(uint64_t Size);
-void internal_free(void *Ptr);
+__attribute__((noinline)) void *internal_malloc(uint64_t Size);
+__attribute__((noinline)) void internal_free(void *Ptr);
 #pragma omp end declare target
 
 #pragma omp begin declare target device_type(nohost)
@@ -59,18 +59,35 @@ namespace {
 ///
 ///{
 
-// global_allocate uses ockl_dm_alloc to manage a global memory heap
+// global_allocate uses ockl_dm_alloc/asan_malloc_impl to manage a global memory
+// heap
 __attribute__((noinline)) extern "C" uint64_t __ockl_dm_alloc(uint64_t bufsz);
 __attribute__((noinline)) extern "C" void __ockl_dm_dealloc(uint64_t ptr);
-
+#if SANITIZER_AMDGPU
+__attribute__((noinline)) extern "C" uint64_t __asan_malloc_impl(uint64_t bufsz,
+                                                                 uint64_t pc);
+__attribute__((noinline)) extern "C" void __asan_free_impl(uint64_t ptr,
+                                                           uint64_t pc);
+#endif
 #pragma omp begin declare variant match(device = {arch(amdgcn)})
 extern "C" {
-void *internal_malloc(uint64_t Size) {
+__attribute__((noinline)) void *internal_malloc(uint64_t Size) {
+#if SANITIZER_AMDGPU
+  uint64_t ptr =
+      __asan_malloc_impl(Size, (uint64_t)__builtin_return_address(0));
+#else
   uint64_t ptr = __ockl_dm_alloc(Size);
+#endif
   return (void *)ptr;
 }
 
-void internal_free(void *Ptr) { __ockl_dm_dealloc((uint64_t)Ptr); }
+__attribute__((noinline)) void internal_free(void *Ptr) {
+#if SANITIZER_AMDGPU
+  __asan_free_impl((uint64_t)Ptr, (uint64_t)__builtin_return_address(0));
+#else
+  __ockl_dm_dealloc((uint64_t)Ptr);
+#endif
+}
 }
 #pragma omp end declare variant
 

offload/DeviceRTL/src/exports

@@ -24,6 +24,8 @@ f90print*
 fprintf_*
 
 hostexec*
+__ockl*
+__asan*
 __strlen_max
 __tgt_fort_ptr_assn_i8
 global_allocate
@@ -36,4 +38,4 @@ sin
 cos
 atan2f
 
-abs
+abs

offload/hostexec/CMakeLists.txt

@@ -179,6 +179,12 @@ foreach(archname ${HOSTRPC_ARCHS})
           -Xclang -mlink-builtin-bitcode -Xclang ${ockl_path}
           -Wno-linker-warnings # Silence the empty host compilation.
           -c -emit-llvm --offload-device-only -nogpulib)
+
+      if(SANITIZER_AMDGPU)
+        get_target_property(asanrtl_path asanrtl IMPORTED_LOCATION)
+        list(APPEND openmp_device_args -DSANITIZER_AMDGPU=1)
+        list(APPEND openmp_device_args -Xclang -mlink-bitcode-file -Xclang ${asanrtl_path})
+      endif()
     elseif(${archname} STREQUAL "nvptx")
        set(openmp_device_args 
       -I../../runtime/src  # to pickup omp.h, we may need a dependency

offload/hostexec/src/hostexec_stubs.cpp

@@ -224,13 +224,28 @@ long hostexec_long_execute(char *print_buffer, uint32_t bufsz) {
 // See https://github.com/llvm/llvm-project/issues/63597
 __attribute__((noinline)) extern "C" uint64_t __ockl_dm_alloc(uint64_t bufsz);
 __attribute__((noinline)) extern "C" void __ockl_dm_dealloc(uint64_t ptr);
+#if SANITIZER_AMDGPU
+__attribute__((noinline)) extern "C" uint64_t __asan_malloc_impl(uint64_t bufsz,
+                                                                 uint64_t pc);
+__attribute__((noinline)) extern "C" void __asan_free_impl(uint64_t ptr,
+                                                           uint64_t pc);
+#endif
 
 // FIXME: Deprecate upstream, change test cases to use malloc & free directly
 __attribute__((flatten, always_inline)) char *global_allocate(uint32_t bufsz) {
+#if SANITIZER_AMDGPU
+  return (char *)__asan_malloc_impl(bufsz,
+                                    (uint64_t)__builtin_return_address(0));
+#else
   return (char *)__ockl_dm_alloc(bufsz);
+#endif
 }
 __attribute__((flatten, always_inline)) int global_free(void *ptr) {
+#if SANITIZER_AMDGPU
+  __asan_free_impl((uint64_t)ptr, (uint64_t)__builtin_return_address(0));
+#else
   __ockl_dm_dealloc((uint64_t)ptr);
+#endif
   return 0;
 }
 

offload/tools/prep-libomptarget-bc/prep-libomptarget-bc.cpp

@@ -152,6 +152,25 @@ static bool convertExternsToLinkOnce(Module *MOUT, LLVMContext &Ctx) {
         if (!strncmp(F->getName().str().c_str(), "hostexec_invoke",
                      strlen("hostexec_invoke")))
           continue;
+#if SANITIZER_AMDGPU
+        if (!strncmp(F->getName().str().c_str(), "internal_malloc",
+                     strlen("internal_malloc")) ||
+            !strncmp(F->getName().str().c_str(), "internal_free",
+                     strlen("internal_free")))
+          continue;
+        if (!strncmp(F->getName().str().c_str(), "global_allocate",
+                     strlen("global_allocate")) ||
+            !strncmp(F->getName().str().c_str(), "global_free",
+                     strlen("global_free"))) {
+          F->removeFnAttr(llvm::Attribute::AlwaysInline);
+          F->addFnAttr(llvm::Attribute::NoInline);
+          continue;
+        }
+        // Tag every '__kmpc*' function with 'sanitize_address' attribute by
+        // default when sanitizer is enabled for amdgpu. This is mainly required
+        // for LDS overflow detection as it will instrument all LDS variables.
+        F->addFnAttr(llvm::Attribute::SanitizeAddress);
+#endif
         // all other functions
         F->removeFnAttr(llvm::Attribute::OptimizeNone);
         F->removeFnAttr(llvm::Attribute::NoInline);
@@ -185,7 +204,14 @@ static bool convertDmAllocToLinkOnce(Module *MOUT, LLVMContext &Ctx) {
           F->setLinkage(GlobalValue::LinkOnceODRLinkage);
           F->setVisibility(GlobalValue::ProtectedVisibility);
         }
-          continue;
+#if SANITIZER_AMDGPU
+        // Tag every '__kmpc*' function with 'sanitize_address' attribute by
+        // default when sanitizer is enabled for amdgpu. This is mainly required
+        // for LDS overflow detection as it will instrument all LDS variables.
+        // if (regex.match(F->getName().str()))
+        F->addFnAttr(llvm::Attribute::SanitizeAddress);
+#endif
+        continue;
       }
     }
   }