[HIP] Move std headers after device malloc/free

[ashi1]

Summary

Set the device malloc and free functions as weak,
and move the std headers after device malloc/free
to avoid issues with std malloc/free.

[yaxunl]

need a test in https://github.com/llvm/llvm-project/blob/main/clang/test/Headers/hip-header.hip

[ashi1]

Thanks, added a test.

[yaxunl]

clang/test/Headers/hip-header.hip

29–34

can we reuse the first RUN line instead of adding a new one?

123

need to check definition of malloc for linkage and name mangling. same for free. we do not need to check test_malloc and test_free.

[ashi1]

Updated the checks in the hip test.

[yaxunl]

LGTM. Thanks!

clang/lib/Headers/__clang_hip_runtime_wrapper.h

  *
  */
 
 #ifndef __CLANG_HIP_RUNTIME_WRAPPER_H__
 #define __CLANG_HIP_RUNTIME_WRAPPER_H__
 
 #if __HIP__
 
-#if !defined(__HIPCC_RTC__)
-#include <cmath>
-#include <cstdlib>
-#include <stdlib.h>
-#else
-typedef __SIZE_TYPE__ size_t;
-// Define macros which are needed to declare HIP device API's without standard
-// C/C++ headers. This is for readability so that these API's can be written
-// the same way as non-hipRTC use case. These macros need to be popped so that
-// they do not pollute users' name space.
-#pragma push_macro("NULL")
-#pragma push_macro("uint32_t")
-#pragma push_macro("uint64_t")
-#pragma push_macro("CHAR_BIT")
-#pragma push_macro("INT_MAX")
-#define NULL (void *)0
-#define uint32_t __UINT32_TYPE__
-#define uint64_t __UINT64_TYPE__
-#define CHAR_BIT __CHAR_BIT__
-#define INT_MAX __INTMAX_MAX__
-#endif // __HIPCC_RTC__
-
 #define __host__ __attribute__((host))
 #define __device__ __attribute__((device))
 #define __global__ __attribute__((global))
 #define __shared__ __attribute__((shared))
 #define __constant__ __attribute__((constant))
 #define __managed__ __attribute__((managed))
 
 #if !defined(__cplusplus) || __cplusplus < 201103L
extern "C" {
   __attribute__((weak))
   __attribute__((noreturn))
   __device__ void __cxa_deleted_virtual(void) {
     __builtin_trap();
   }
 }
 #endif //__cplusplus
 
+typedef __SIZE_TYPE__ __hip_size_t;
+
+#ifdef __cplusplus
+extern "C" {
+#endif //__cplusplus
+
 #if __HIP_ENABLE_DEVICE_MALLOC__
-extern "C" __device__ void *__hip_malloc(size_t __size);
-extern "C" __device__ void *__hip_free(void *__ptr);
-static inline __device__ void *malloc(size_t __size) {
+__device__ void *__hip_malloc(__hip_size_t __size);
+__device__ void *__hip_free(void *__ptr);
+__attribute__((weak)) inline __device__ void *malloc(__hip_size_t __size) {
   return __hip_malloc(__size);
 }
-static inline __device__ void *free(void *__ptr) { return __hip_free(__ptr); }
+__attribute__((weak)) inline __device__ void *free(void *__ptr) {
+  return __hip_free(__ptr);
+}
 #else
-static inline __device__ void *malloc(size_t __size) {
+__attribute__((weak)) inline __device__ void *malloc(__hip_size_t __size) {
   __builtin_trap();
   return nullptr;
 }
-static inline __device__ void *free(void *__ptr) {
+__attribute__((weak)) inline __device__ void *free(void *__ptr) {
   __builtin_trap();
   return nullptr;
 }
 #endif
 
+#ifdef __cplusplus
+} // extern "C"
+#endif //__cplusplus
+
+#if !defined(__HIPCC_RTC__)
+#include <cmath>
+#include <cstdlib>
+#include <stdlib.h>
+#else
+typedef __SIZE_TYPE__ size_t;
+// Define macros which are needed to declare HIP device API's without standard
+// C/C++ headers. This is for readability so that these API's can be written
+// the same way as non-hipRTC use case. These macros need to be popped so that
+// they do not pollute users' name space.
+#pragma push_macro("NULL")
+#pragma push_macro("uint32_t")
+#pragma push_macro("uint64_t")
+#pragma push_macro("CHAR_BIT")
+#pragma push_macro("INT_MAX")
+#define NULL (void *)0
+#define uint32_t __UINT32_TYPE__
+#define uint64_t __UINT64_TYPE__
+#define CHAR_BIT __CHAR_BIT__
+#define INT_MAX __INTMAX_MAX__
+#endif // __HIPCC_RTC__
+
 #include <__clang_hip_libdevice_declares.h>
 #include <__clang_hip_math.h>
 
 #if !_OPENMP || __HIP_ENABLE_CUDA_WRAPPER_FOR_OPENMP__
 #if defined(__HIPCC_RTC__)
 #include <__clang_hip_cmath.h>
 #else
 #include <__clang_cuda_math_forward_declares.h>

clang/test/Headers/hip-header.hip

 // RUN:   -D__HIPCC_RTC__ -DUSE_ISNAN_WITH_INT_RETURN | FileCheck %s -check-prefixes=AMD_INT_RETURN
 // RUN: %clang_cc1 -include __clang_hip_runtime_wrapper.h \
 // RUN:   -internal-isystem %S/../../lib/Headers/cuda_wrappers \
 // RUN:   -internal-isystem %S/Inputs/include \
 // RUN:   -triple amdgcn-amd-amdhsa -aux-triple x86_64-unknown-unknown \
 // RUN:   -target-cpu gfx906 -emit-llvm %s -fcuda-is-device -o - \
 // RUN:   -D__HIPCC_RTC__ -std=c++14 | FileCheck -check-prefixes=CHECK,CXX14 %s
 
 // expected-no-diagnostics
 
 // Check support for pure and deleted virtual functions
 struct base {
   __host__
   __device__
   virtual void pv() = 0;
   __host__
   __device__
   virtual void dv() = delete;
 };
 struct derived:base {
   __host__
   __device__
__device__ double test_isnan() {
   r += isnan(f);
 
   // AMD_INT_RETURN: call i32 @__ocml_isnan_f64(double
   // AMD_BOOL_RETURN: call i32 @__ocml_isnan_f64(double
   r += isnan(d);
 
   return r ;
 }
+
+// Check that device malloc and free do not conflict with std headers.
+#include <cstdlib>
+// CHECK-LABEL: define{{.*}}@_Z11test_malloc
+// CHECK: call {{.*}}i8* @malloc(i64
+// CHECK: define weak {{.*}}i8* @malloc(i64
+__device__ void test_malloc(void *a) {
+  a = malloc(42);
+}
+
+// CHECK-LABEL: define{{.*}}@_Z9test_free
+// CHECK: call {{.*}}i8* @free(i8*
+// CHECK: define weak {{.*}}i8* @free(i8*
+__device__ void test_free(void *a) {
+  free(a);
+}