[openmp] Annotate tmp variables with omp_thread_mem_alloc

[JonChesterfield]

Summary

Fixes miscompile of calls into ocml. Bug 51445.

The stack variable double __tmp is moved to dynamically allocated shared
memory by CGOpenMPRuntimeGPU. This is usually fine, but when the variable
is passed to a function that is explicitly annotated address_space(5) then
allocating the variable off-stack leads to a miscompile in the back end,
which cannot decide to move the variable back to the stack from shared.

This could be fixed by removing the AS(5) annotation from the math library
or by explicitly marking the variables as thread_mem_alloc. The cast to
AS(5) is still a no-op once IR is reached.

[JonChesterfield]

@ye-luo this fixes the modf and sincos test cases from https://github.com/ye-luo/openmp-target

[jdoerfert]

Assuming this causes us to generate an alloc as(5) for __tmp, LG

[JonChesterfield]

In D107971#2941666, @jdoerfert wrote:

Assuming this causes us to generate an alloc as(5) for __tmp, LG

Yes, with the running example and this patch we get the perfect:

%__tmp = alloca double, align 8, addrspace(5)
%call = call double @__ocml_sincos_f64(double %__x, double addrspace(5)* %__tmp)

Before, we get:

%__tmp = call fastcc i8* @__kmpc_alloc_shared()
%__tmp_on_stack = bitcast i8* %__tmp to double*
%__tmp_on_stack.ascast = addrspacecast double* %__tmp_on_stack to double addrspace(5)*
%call = call double @__ocml_sincos_f64(double %__x, double addrspace(5)* %__tmp_on_stack.ascast)

which interacts badly with the addrspace(5) annotation on the ocml function

[JonChesterfield]

Failed a CI job that builds an openmp test in an environment without omp.h, will revert.

Thoughts on fixing? Putting the omp allocator definition in this header is likely to collide with a real omp.h. Fairly clean fix is to move the definitions into the deviceRTL. Tempting fix is to add a minimal omp.h to the test dir

example failure:

In file included from /home/tcwg-buildslave/worker/clang-aarch64-quick/llvm/clang/test/Headers/amdgcn_openmp_device_math.c:9:
In file included from /home/tcwg-buildslave/worker/clang-aarch64-quick/llvm/clang/test/Headers/../../lib/Headers/openmp_wrappers/math.h:55:
/home/tcwg-buildslave/worker/clang-aarch64-quick/stage1/lib/clang/14.0.0/include/__clang_hip_math.h:23:10: fatal error: 'omp.h' file not found
#include <omp.h>

[jdoerfert]

In D107971#2941739, @JonChesterfield wrote:

Failed a CI job that builds an openmp test in an environment without omp.h, will revert.

Thoughts on fixing? Putting the omp allocator definition in this header is likely to collide with a real omp.h. Fairly clean fix is to move the definitions into the deviceRTL. Tempting fix is to add a minimal omp.h to the test dir

example failure:

In file included from /home/tcwg-buildslave/worker/clang-aarch64-quick/llvm/clang/test/Headers/amdgcn_openmp_device_math.c:9:
In file included from /home/tcwg-buildslave/worker/clang-aarch64-quick/llvm/clang/test/Headers/../../lib/Headers/openmp_wrappers/math.h:55:
/home/tcwg-buildslave/worker/clang-aarch64-quick/stage1/lib/clang/14.0.0/include/__clang_hip_math.h:23:10: fatal error: 'omp.h' file not found
#include <omp.h>

yes add omp.h, in clang/test/Headers/Inputs/include.

[JonChesterfield]

intended content for the new omp.h is

#ifndef __OMP_H
#define __OMP_H

#if _OPENMP
// Follows the pattern in interface.h
// Clang sema checks this type carefully, needs to closely match that from omp.h
typedef enum omp_allocator_handle_t {
  omp_null_allocator = 0,
  omp_default_mem_alloc = 1,
  omp_large_cap_mem_alloc = 2,
  omp_const_mem_alloc = 3,
  omp_high_bw_mem_alloc = 4,
  omp_low_lat_mem_alloc = 5,
  omp_cgroup_mem_alloc = 6,
  omp_pteam_mem_alloc = 7,
  omp_thread_mem_alloc = 8,
  KMP_ALLOCATOR_MAX_HANDLE = ~(0U)
} omp_allocator_handle_t;
#endif

#endif

trying to reproduce the failure locally first to verify that the fix works

edit: confirmed, will let the pre-commit build go through green then re-commit

[JonChesterfield]

• add minimal omp.h to clang test inputs

clang/lib/Headers/__clang_hip_math.h

 #endif
 
 #if !defined(__HIPCC_RTC__)
 #if defined(__cplusplus)
 #include <algorithm>
 #endif
 #include <limits.h>
 #include <stdint.h>
+#ifdef __OPENMP_AMDGCN__
+#include <omp.h>
+#endif
 #endif // !defined(__HIPCC_RTC__)
 
 #pragma push_macro("__DEVICE__")
 
 #ifdef __OPENMP_AMDGCN__
 #define __DEVICE__ static inline __attribute__((always_inline, nothrow))
 #else
 #define __DEVICE__ static __device__ inline __attribute__((always_inline))

 float fminf(float __x, float __y) { return __ocml_fmin_f32(__x, __y); }
 
 __DEVICE__
 float fmodf(float __x, float __y) { return __ocml_fmod_f32(__x, __y); }
 
 __DEVICE__
 float frexpf(float __x, int *__nptr) {
   int __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   float __r =
       __ocml_frexp_f32(__x, (__attribute__((address_space(5))) int *)&__tmp);
   *__nptr = __tmp;
 
   return __r;
 }
 
 __DEVICE__

 long int lrintf(float __x) { return __ocml_rint_f32(__x); }
 
 __DEVICE__
 long int lroundf(float __x) { return __ocml_round_f32(__x); }
 
 __DEVICE__
 float modff(float __x, float *__iptr) {
   float __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   float __r =
       __ocml_modf_f32(__x, (__attribute__((address_space(5))) float *)&__tmp);
   *__iptr = __tmp;
   return __r;
 }
 
 __DEVICE__
 float nanf(const char *__tagp) {

 __DEVICE__
 float remainderf(float __x, float __y) {
   return __ocml_remainder_f32(__x, __y);
 }
 
 __DEVICE__
 float remquof(float __x, float __y, int *__quo) {
   int __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   float __r = __ocml_remquo_f32(
       __x, __y, (__attribute__((address_space(5))) int *)&__tmp);
   *__quo = __tmp;
 
   return __r;
 }
 
 __DEVICE__

 float scalbnf(float __x, int __n) { return __ocml_scalbn_f32(__x, __n); }
 
 __DEVICE__
 __RETURN_TYPE __signbitf(float __x) { return __ocml_signbit_f32(__x); }
 
 __DEVICE__
 void sincosf(float __x, float *__sinptr, float *__cosptr) {
   float __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   *__sinptr =
       __ocml_sincos_f32(__x, (__attribute__((address_space(5))) float *)&__tmp);
   *__cosptr = __tmp;
 }
 
 __DEVICE__
 void sincospif(float __x, float *__sinptr, float *__cosptr) {
   float __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   *__sinptr = __ocml_sincospi_f32(
       __x, (__attribute__((address_space(5))) float *)&__tmp);
   *__cosptr = __tmp;
 }
 
 __DEVICE__
 float sinf(float __x) { return __ocml_sin_f32(__x); }
 

 double fmin(double __x, double __y) { return __ocml_fmin_f64(__x, __y); }
 
 __DEVICE__
 double fmod(double __x, double __y) { return __ocml_fmod_f64(__x, __y); }
 
 __DEVICE__
 double frexp(double __x, int *__nptr) {
   int __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   double __r =
       __ocml_frexp_f64(__x, (__attribute__((address_space(5))) int *)&__tmp);
   *__nptr = __tmp;
   return __r;
 }
 
 __DEVICE__
 double hypot(double __x, double __y) { return __ocml_hypot_f64(__x, __y); }

 long int lrint(double __x) { return __ocml_rint_f64(__x); }
 
 __DEVICE__
 long int lround(double __x) { return __ocml_round_f64(__x); }
 
 __DEVICE__
 double modf(double __x, double *__iptr) {
   double __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   double __r =
       __ocml_modf_f64(__x, (__attribute__((address_space(5))) double *)&__tmp);
   *__iptr = __tmp;
 
   return __r;
 }
 
 __DEVICE__

 __DEVICE__
 double remainder(double __x, double __y) {
   return __ocml_remainder_f64(__x, __y);
 }
 
 __DEVICE__
 double remquo(double __x, double __y, int *__quo) {
   int __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   double __r = __ocml_remquo_f64(
       __x, __y, (__attribute__((address_space(5))) int *)&__tmp);
   *__quo = __tmp;
 
   return __r;
 }
 
 __DEVICE__

 __RETURN_TYPE __signbit(double __x) { return __ocml_signbit_f64(__x); }
 
 __DEVICE__
 double sin(double __x) { return __ocml_sin_f64(__x); }
 
 __DEVICE__
 void sincos(double __x, double *__sinptr, double *__cosptr) {
   double __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   *__sinptr = __ocml_sincos_f64(
       __x, (__attribute__((address_space(5))) double *)&__tmp);
   *__cosptr = __tmp;
 }
 
 __DEVICE__
 void sincospi(double __x, double *__sinptr, double *__cosptr) {
   double __tmp;
+#ifdef __OPENMP_AMDGCN__
+#pragma omp allocate(__tmp) allocator(omp_thread_mem_alloc)
+#endif
   *__sinptr = __ocml_sincospi_f64(
       __x, (__attribute__((address_space(5))) double *)&__tmp);
   *__cosptr = __tmp;
 }
 
 __DEVICE__
 double sinh(double __x) { return __ocml_sinh_f64(__x); }
 

clang/test/Headers/Inputs/include/omp.h

+#ifndef __OMP_H
+#define __OMP_H
+
+#if _OPENMP
+// Follows the pattern in interface.h
+// Clang sema checks this type carefully, needs to closely match that from omp.h
+typedef enum omp_allocator_handle_t {
+  omp_null_allocator = 0,
+  omp_default_mem_alloc = 1,
+  omp_large_cap_mem_alloc = 2,
+  omp_const_mem_alloc = 3,
+  omp_high_bw_mem_alloc = 4,
+  omp_low_lat_mem_alloc = 5,
+  omp_cgroup_mem_alloc = 6,
+  omp_pteam_mem_alloc = 7,
+  omp_thread_mem_alloc = 8,
+  KMP_ALLOCATOR_MAX_HANDLE = ~(0U)
+} omp_allocator_handle_t;
+#endif
+
+#endif