[OpenMP] Introduce ompx.h and 3D wrappers (threadId, threadDim, ...)

[jdoerfert]

Summary

The new ompx.h header will give us a place to put extensions. The first
are 3D getters for the common cuda values:

`{threadId,threadDim,blockId,blockDim}.{x,y,z}`

[jdoerfert]

add missing file

[tianshilei1992]

openmp/runtime/src/include/ompx.h.var

49

Since the patch contains host support as well, we need a test for host.

[jdoerfert]

openmp/runtime/src/include/ompx.h.var

49

The tests are for device and host, no?

openmp/libomptarget/DeviceRTL/src/Mapping.cpp

 }
 
 __attribute__((noinline)) uint32_t __kmpc_get_warp_size() {
   FunctionTracingRAII();
   return impl::getWarpSize();
 }
 }
 
+#define _TGT_KERNEL_LANGUAGE(NAME, MAPPER_NAME)                                \
+  extern "C" int ompx_##NAME(int Dim) { return mapping::MAPPER_NAME(Dim); }
+
+_TGT_KERNEL_LANGUAGE(thread_id, getThreadIdInBlock)
+_TGT_KERNEL_LANGUAGE(thread_dim, getNumberOfThreadsInBlock)
+_TGT_KERNEL_LANGUAGE(block_id, getBlockIdInKernel)
+_TGT_KERNEL_LANGUAGE(block_dim, getNumberOfBlocksInKernel)
+
 #pragma omp end declare target

openmp/libomptarget/DeviceRTL/src/exports

 omp_*
+ompx_*
 *llvm_*
 __kmpc_*
 
 _ZN4ompx*
 
 IsSPMDMode
 
 memcmp
 printf
 __assert_fail

openmp/libomptarget/test/api/ompx_3d.c

+// RUN: %libomptarget-compile-run-and-check-generic
+
+#include <omp.h>
+#include <ompx.h>
+#include <stdio.h>
+
+void foo(int device) {
+  int tid = 0, bid = 0, bdim = 0;
+#pragma omp target teams distribute parallel for map(from                      \
+                                                     : tid, bid, bdim)         \
+    device(device) thread_limit(2) num_teams(5)
+  for (int i = 0; i < 1000; ++i) {
+    if (i == 42) {
+      tid = ompx_thread_dim_x();
+      bid = ompx_block_id_x();
+      bdim = ompx_block_dim_x();
+    }
+  }
+  // CHECK: tid: 2, bid: 1, bdim: 5
+  // CHECK: tid: 2, bid: 0, bdim: 1
+  printf("tid: %i, bid: %i, bdim: %i\n", tid, bid, bdim);
+}
+
+int main() {
+  foo(omp_get_default_device());
+  foo(omp_get_initial_device());
+}

openmp/libomptarget/test/api/ompx_3d.cpp

+// RUN: %libomptarget-compilexx-run-and-check-generic
+
+#include <omp.h>
+#include <ompx.h>
+#include <stdio.h>
+
+void foo(int device) {
+  int tid = 0, bid = 0, bdim = 0;
+#pragma omp target teams distribute parallel for map(from                      \
+                                                     : tid, bid, bdim)         \
+    device(device) thread_limit(2) num_teams(5)
+  for (int i = 0; i < 1000; ++i) {
+    if (i == 42) {
+      tid = ompx::thread_dim_x();
+      bid = ompx::block_id_x();
+      bdim = ompx::block_dim_x();
+    }
+  }
+  // CHECK: tid: 2, bid: 1, bdim: 5
+  // CHECK: tid: 2, bid: 0, bdim: 1
+  printf("tid: %i, bid: %i, bdim: %i\n", tid, bid, bdim);
+}
+
+int main() {
+  foo(omp_get_default_device());
+  foo(omp_get_initial_device());
+}

openmp/runtime/cmake/LibompExports.cmake

 set(LIBOMP_EXPORTS_INC_DIR "${LIBOMP_EXPORTS_PLATFORM_DIR}/include")
 set(LIBOMP_EXPORTS_MOD_DIR "${LIBOMP_EXPORTS_PLATFORM_DIR}/include_compat")
 set(LIBOMP_EXPORTS_LIB_DIR "${LIBOMP_EXPORTS_DIR}/${libomp_platform}${libomp_suffix}/lib")
 
 # Put headers in exports/ directory post build
 add_custom_command(TARGET omp POST_BUILD
   COMMAND ${CMAKE_COMMAND} -E make_directory ${LIBOMP_EXPORTS_CMN_DIR}
   COMMAND ${CMAKE_COMMAND} -E copy omp.h ${LIBOMP_EXPORTS_CMN_DIR}
+  COMMAND ${CMAKE_COMMAND} -E copy ompx.h ${LIBOMP_EXPORTS_CMN_DIR}
 )
 if(${LIBOMP_OMPT_SUPPORT})
   add_custom_command(TARGET omp POST_BUILD
     COMMAND ${CMAKE_COMMAND} -E copy omp-tools.h ${LIBOMP_EXPORTS_CMN_DIR}
   )
 endif()
 if(${LIBOMP_FORTRAN_MODULES})
   add_custom_command(TARGET libomp-mod POST_BUILD

openmp/runtime/src/CMakeLists.txt

 #
 #//===----------------------------------------------------------------------===//
 #//
 #// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #// See https://llvm.org/LICENSE.txt for license information.
 #// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 #//
 #//===----------------------------------------------------------------------===//
 #
 
 include(ExtendPath)
 
 # Configure omp.h, kmp_config.h and omp-tools.h if necessary
 configure_file(${LIBOMP_INC_DIR}/omp.h.var omp.h @ONLY)
+configure_file(${LIBOMP_INC_DIR}/ompx.h.var ompx.h @ONLY)
 configure_file(kmp_config.h.cmake kmp_config.h @ONLY)
 if(${LIBOMP_OMPT_SUPPORT})
   configure_file(${LIBOMP_INC_DIR}/omp-tools.h.var omp-tools.h @ONLY)
 endif()
 
 # Generate message catalog files: kmp_i18n_id.inc and kmp_i18n_default.inc
 add_custom_command(
   OUTPUT  kmp_i18n_id.inc
foreach(alias IN LISTS LIBOMP_ALIASES)
         \"${alias}${LIBOMP_LIBRARY_SUFFIX}\" WORKING_DIRECTORY
         \"\$ENV{DESTDIR}${outdir}\")")
     endforeach()
   endif()
 endif()
 install(
   FILES
   ${CMAKE_CURRENT_BINARY_DIR}/omp.h
+  ${CMAKE_CURRENT_BINARY_DIR}/ompx.h
   DESTINATION ${LIBOMP_HEADERS_INSTALL_PATH}
 )
 if(${LIBOMP_OMPT_SUPPORT})
   install(FILES ${CMAKE_CURRENT_BINARY_DIR}/omp-tools.h DESTINATION ${LIBOMP_HEADERS_INSTALL_PATH})
   # install under legacy name ompt.h
   install(FILES ${CMAKE_CURRENT_BINARY_DIR}/omp-tools.h DESTINATION ${LIBOMP_HEADERS_INSTALL_PATH} RENAME ompt.h)
   set(LIBOMP_OMP_TOOLS_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR} PARENT_SCOPE)
 endif()
 if(${LIBOMP_FORTRAN_MODULES})
   install(FILES
     ${CMAKE_CURRENT_BINARY_DIR}/omp_lib.h
     ${CMAKE_CURRENT_BINARY_DIR}/omp_lib.mod
     ${CMAKE_CURRENT_BINARY_DIR}/omp_lib_kinds.mod
     DESTINATION ${LIBOMP_HEADERS_INSTALL_PATH}
   )
 endif()

openmp/runtime/src/include/ompx.h.var

+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef __OMPX_H
+#define __OMPX_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+int omp_get_ancestor_thread_num(int);
+int omp_get_team_size(int);
+
+#ifdef __cplusplus
+}
+#endif
+
+/// Target kernel language extensions
+///
+/// These extensions exist for the host to allow fallback implementations,
+/// however, they cannot be arbitrarily composed with OpenMP. If the rules of
+/// the kernel language are followed, the host fallbacks should behave as
+/// expected since the kernel is represented as 3 sequential outer loops, one
+/// for each grid dimension, and three (nested) parallel loops, one for each
+/// block dimension. This fallback is not supposed to be optimal and should be
+/// configurable by the user.
+///
+///{
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+enum {
+  ompx_dim_x = 0,
+  ompx_dim_y = 1,
+  ompx_dim_z = 2,
+};
+
+#pragma omp begin declare variant match(device = {kind(host)})
+#define _TGT_KERNEL_LANGUAGE_HOST_IMPL_C(NAME, VALUE)                          \
+  static inline int ompx_##NAME(int Dim) { return VALUE; }
+
+_TGT_KERNEL_LANGUAGE_HOST_IMPL_C(thread_id,
+                                 omp_get_ancestor_thread_num(Dim + 1))
+_TGT_KERNEL_LANGUAGE_HOST_IMPL_C(thread_dim, omp_get_team_size(Dim + 1))
+_TGT_KERNEL_LANGUAGE_HOST_IMPL_C(block_id, 0)
+_TGT_KERNEL_LANGUAGE_HOST_IMPL_C(block_dim, 1)
+#undef _TGT_KERNEL_LANGUAGE_HOST_IMPL_C
+
+#pragma omp end declare variant
+
+#define _TGT_KERNEL_LANGUAGE_DECL_C(NAME)                                      \
+  int ompx_##NAME(int Dim);                                                    \
+  static inline int ompx_##NAME##_x() { return ompx_##NAME(ompx_dim_x); }      \
+  static inline int ompx_##NAME##_y() { return ompx_##NAME(ompx_dim_y); }      \
+  static inline int ompx_##NAME##_z() { return ompx_##NAME(ompx_dim_z); }
+
+_TGT_KERNEL_LANGUAGE_DECL_C(thread_id)
+_TGT_KERNEL_LANGUAGE_DECL_C(thread_dim)
+_TGT_KERNEL_LANGUAGE_DECL_C(block_id)
+_TGT_KERNEL_LANGUAGE_DECL_C(block_dim)
+#undef _TGT_KERNEL_LANGUAGE_DECL_C
+
+#ifdef __cplusplus
+}
+#endif
+
+#ifdef __cplusplus
+
+namespace ompx {
+
+enum {
+  dim_x = ompx_dim_x,
+  dim_y = ompx_dim_y,
+  dim_z = ompx_dim_z,
+};
+
+#define _TGT_KERNEL_LANGUAGE_HOST_IMPL_CXX(NAME)                               \
+  static inline int NAME(int Dim) noexcept { return ompx_##NAME(Dim); }        \
+  static inline int NAME##_x() noexcept { return NAME(ompx_dim_x); }           \
+  static inline int NAME##_y() noexcept { return NAME(ompx_dim_y); }           \
+  static inline int NAME##_z() noexcept { return NAME(ompx_dim_z); }
+
+_TGT_KERNEL_LANGUAGE_HOST_IMPL_CXX(thread_id)
+_TGT_KERNEL_LANGUAGE_HOST_IMPL_CXX(thread_dim)
+_TGT_KERNEL_LANGUAGE_HOST_IMPL_CXX(block_id)
+_TGT_KERNEL_LANGUAGE_HOST_IMPL_CXX(block_dim)
+#undef _TGT_KERNEL_LANGUAGE_HOST_IMPL_CXX
+
+} // namespace ompx
+#endif
+
+///}
+
+#endif /* __OMPX_H */