[OpenMP][DeviceRTL] Extract shuffle idiom and port it to declare variant

[jdoerfert]

Summary

The shuffle idiom is differently implemented in our supported targets.
To reduce the "target_impl" file we now move the shuffle idiom in it's
own self-contained header that provides the implementation for AMDGPU
and NVPTX. A fallback can be added later on.

[tianshilei1992]

LGTM. This will be the first example to merge different implementations into one file.

openmp/libomptarget/deviceRTLs/common/src/data_sharing.cu

17

One nit: you might have shuffle.h before target_impl.h if using clang-format.

[tianshilei1992]

I can see many warnings emitted, like the following one:

In file included from /home/shiltian/Documents/vscode/llvm-project/openmp/libomptarget/deviceRTLs/common/src/loop.cu:17:
/home/shiltian/Documents/vscode/llvm-project/openmp/libomptarget/deviceRTLs/common/include/shuffle.h:96:16: warning: inline function '__kmpc_impl_shfl_d
own_sync' is not defined [-Wundefined-inline]
inline int32_t __kmpc_impl_shfl_down_sync(int64_t Mask, int32_t Var,                                                                                                   ^

[JonChesterfield]

Comments inline. Not totally sure this is better, code seems longer and more complicated than it was before. Aim is merging the two GPU implementations?

openmp/libomptarget/deviceRTLs/common/include/shuffle.h

27 ↗

(On Diff #320332)

Extern C? Also, why forward declare instead of include header?

40 ↗

(On Diff #320332)

uint64_t for mask. Sign bit is just one of the lanes.

88 ↗

(On Diff #320332)

Seems bad, both because it's a macro instead of variant, and because I thought we'd already got rid of that macro

117 ↗

(On Diff #320332)

Not sure about int constants near places where the 32/64 bit distinction is important

[JonChesterfield]

openmp/libomptarget/deviceRTLs/common/include/shuffle.h

66 ↗

(On Diff #320332)

Wonder if this compiles without a forward declare of the intrinsic

[jdoerfert]

In D95752#2532552, @tianshilei1992 wrote:

I can see many warnings emitted, like the following one:

In file included from /home/shiltian/Documents/vscode/llvm-project/openmp/libomptarget/deviceRTLs/common/src/loop.cu:17:
/home/shiltian/Documents/vscode/llvm-project/openmp/libomptarget/deviceRTLs/common/include/shuffle.h:96:16: warning: inline function '__kmpc_impl_shfl_d
own_sync' is not defined [-Wundefined-inline]
inline int32_t __kmpc_impl_shfl_down_sync(int64_t Mask, int32_t Var,                                                                                                   ^

I will fix this in clang though before.

[jdoerfert]

openmp/libomptarget/deviceRTLs/common/include/shuffle.h

27 ↗

(On Diff #320332)

because our headers are a mess. These are C++ functions, the interface is a mess too

117 ↗

(On Diff #320332)

That's why it's int64_t above, -1 sign extended stays -1.

[jdoerfert]

Addressed comments. uint64, and static to avoid warning, once we have a default
impl we can actually remove the first decl variant and the static if we want.

[jdoerfert]

I'll update this tomorrow, certain parts are not great.

openmp/libomptarget/deviceRTLs/common/include/shuffle.h

88 ↗

(On Diff #320332)

we will, ptx selection. I'll actually update this patch.

[jdoerfert]

Replace the CUDA_VERSION macro, provide a mock default impl for shuffle

[tianshilei1992]

openmp/libomptarget/deviceRTLs/nvptx/CMakeLists.txt

131

There is no shuffle.cpp

[tianshilei1992]

^ shuffle.cpp doesn't exist. Except that, I tested locally and it works properly. LGTM.

[jdoerfert]

Add shuffle.cpp

[JonChesterfield]

Not too keen on mixing code for different targets in the same file. Would prefer the prototype declared in a header and a nvptx.cpp, amdgcn.cpp, other.cpp implementing that interface, preferably via variant so that the whole thing compiles out for some arch.

If the 'fallback/default' works as I suspect, it makes a really nice way to implement generic versions. Ptx uses asm for an operation that amdgpu uses a shift for, would be nice to use the pure c++ one as the default and substitute in the specialised one for only the targets that use it.

openmp/libomptarget/deviceRTLs/common/include/shuffle.h

105 ↗

(On Diff #320368)

I like this a lot. Way better that CUDA_VERSION macros.

[jdoerfert]

In D95752#2534578, @JonChesterfield wrote:

Not too keen on mixing code for different targets in the same file. Would prefer the prototype declared in a header and a nvptx.cpp, amdgcn.cpp, other.cpp implementing that interface, preferably via variant so that the whole thing compiles out for some arch.

I'm hoping for a structure where we put a feature into a file, rather than an architecture per file. If we do both, we get 4 files per feature. This is not intrinsically bad but makes updates less convenient. At the end of the day, it doesn't matter much what we do since we perform LTO on this. I was hoping the approach with one header per feature and one cpp file for externally used functions would spur reuse. That is, helpers are defined only once, opportunities to share code are easier to spot during updates, etc.

If the 'fallback/default' works as I suspect, it makes a really nice way to implement generic versions. Ptx uses asm for an operation that amdgpu uses a shift for, would be nice to use the pure c++ one as the default and substitute in the specialised one for only the targets that use it.

So this is what's happening, at least once we have a "fallback" target model for which we can define what a "shuffle" means.

[JonChesterfield]

Header per target-specific-function as opposed to a list of things like __kmpc_impl_syncwarp seems reasonable. Not the way I'd have sliced it but clearly equivalent.

Minor request, let it be target/shuffle.h, so that we build up a list of functions (stuff in that directory) that new architectures need to implement in one place, distinct from code that will hopefully work out of the box for a new target.

This starts us down a path towards:

devicertl
  - src
     - parallel.cpp
     - target
        - shuffle.h

as opposed to the current nvptx/amdgcn/common split. That has the attraction of being a more conventional build system (it's weird treating one source file as openmp for one target and hip for another, the semantics don't really match) and, given we're building with clang, even using a single cmake file to build for all targets.

[ronlieb]

openmp/libomptarget/deviceRTLs/nvptx/CMakeLists.txt

131

is there an amdgcn CMakeLists.txt equivalent change ? should there be ?

[JonChesterfield]

We could presumably replace the #if CUDA_VERSION >= 9000 in the target_impl.cu file (we should rename these!) with variant, orthogonal to this change. Doing that for the five instances, even just within that file, would let us significantly reduce the number of devicertl libraries compiled.

[jdoerfert]

openmp/libomptarget/deviceRTLs/nvptx/CMakeLists.txt

131

probably, I'll look and add it. FWIW, if we had tests and CI for this, e.g., AMD CI that builds the runtime for AMDGPU, that would expose such a mistake right away ;)

[jdoerfert]

In D95752#2535155, @JonChesterfield wrote:

We could presumably replace the #if CUDA_VERSION >= 9000 in the target_impl.cu file (we should rename these!) with variant, orthogonal to this change. Doing that for the five instances, even just within that file, would let us significantly reduce the number of devicertl libraries compiled.

yes, it's case by case though. We should check what ptx version, or other criterion, is a good selector and replace them. Here it was rather easy in the end.

[ronlieb]

openmp/libomptarget/deviceRTLs/nvptx/CMakeLists.txt

131

i agree we really do need an AMD CI, and to get there we also need to be upstreaming our clang support. so in the spirit of making more progress on this, could you do another review of Singh's patch https://reviews.llvm.org/D94961

[jdoerfert]

Address comments

[JonChesterfield]

openmp/libomptarget/deviceRTLs/nvptx/CMakeLists.txt

131

Note to self, amdgcn does indeed need shuffle.cpp added to the cmake list (plus include path)

openmp/libomptarget/deviceRTLs/amdgcn/CMakeLists.txt

set(cuda_sources
   ${devicertl_base_directory}/common/src/data_sharing.cu
   ${devicertl_base_directory}/common/src/libcall.cu
   ${devicertl_base_directory}/common/src/loop.cu
   ${devicertl_base_directory}/common/src/omp_data.cu
   ${devicertl_base_directory}/common/src/omptarget.cu
   ${devicertl_base_directory}/common/src/parallel.cu
   ${devicertl_base_directory}/common/src/reduction.cu
   ${devicertl_base_directory}/common/src/support.cu
+  ${devicertl_base_directory}/common/src/shuffle.cpp
   ${devicertl_base_directory}/common/src/sync.cu
   ${devicertl_base_directory}/common/src/task.cu)
 
 set(h_files
   ${CMAKE_CURRENT_SOURCE_DIR}/src/amdgcn_interface.h
   ${CMAKE_CURRENT_SOURCE_DIR}/src/target_impl.h
   ${devicertl_base_directory}/common/debug.h
   ${devicertl_base_directory}/common/device_environment.h
set(cu_cmd ${AOMP_BINDIR}/clang++
     -D__AMDGCN__
     -Xclang -target-cpu -Xclang ${mcpu}
     -fvisibility=default
     -Wno-unused-value
     -nogpulib
     -O${optimization_level}
     ${CUDA_DEBUG}
     -I${CMAKE_CURRENT_SOURCE_DIR}/src
+    -I${devicertl_base_directory}/common/include
     -I${devicertl_base_directory})
 
   set(bc1_files)
 
   foreach(file ${ARGN})
     get_filename_component(fname ${file} NAME_WE)
     set(bc1_filename ${fname}.${mcpu}.bc)
 

openmp/libomptarget/deviceRTLs/amdgcn/src/target_impl.hip

EXTERN double __kmpc_impl_get_wtime() {
   return 0;
 }
 
 // Warp vote function
 EXTERN __kmpc_impl_lanemask_t __kmpc_impl_activemask() {
   return __builtin_amdgcn_read_exec();
 }
 
-EXTERN int32_t __kmpc_impl_shfl_sync(__kmpc_impl_lanemask_t, int32_t var,
-                                     int32_t srcLane) {
-  int width = WARPSIZE;
-  int self = GetLaneId();
-  int index = srcLane + (self & ~(width - 1));
-  return __builtin_amdgcn_ds_bpermute(index << 2, var);
-}
-
-EXTERN int32_t __kmpc_impl_shfl_down_sync(__kmpc_impl_lanemask_t, int32_t var,
-                                          uint32_t laneDelta, int32_t width) {
-  int self = GetLaneId();
-  int index = self + laneDelta;
-  index = (int)(laneDelta + (self & (width - 1))) >= width ? self : index;
-  return __builtin_amdgcn_ds_bpermute(index << 2, var);
-}
-
 uint32_t __kmpc_L1_Barrier [[clang::loader_uninitialized]];
 #pragma allocate(__kmpc_L1_Barrier) allocator(omp_pteam_mem_alloc)
 
 EXTERN void __kmpc_impl_target_init() {
   // Don't have global ctors, and shared memory is not zero init
   __atomic_store_n(&__kmpc_L1_Barrier, 0u, __ATOMIC_RELEASE);
 }
 

openmp/libomptarget/deviceRTLs/common/include/target/shuffle.h

+//===- shuffle.h - OpenMP variants of the shuffle idiom for all targets -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Shuffle function implementations for all supported targets.
+//
+// Note: We unify the mask type to uint64_t instead of __kmpc_impl_lanemask_t.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LIBOMPTARGET_DEVICERTL_SHUFFLE_H
+#define LIBOMPTARGET_DEVICERTL_SHUFFLE_H
+
+#include <assert.h>
+#include <inttypes.h>
+
+#pragma omp declare target
+
+/// External shuffle API
+///
+///{
+
+extern "C" {
+int32_t __kmpc_shuffle_int32(int32_t val, int16_t delta, int16_t size);
+int64_t __kmpc_shuffle_int64(int64_t val, int16_t delta, int16_t size);
+}
+
+///}
+
+/// Forward declarations
+///
+///{
+unsigned GetLaneId();
+unsigned GetWarpSize();
+void __kmpc_impl_unpack(uint64_t val, uint32_t &lo, uint32_t &hi);
+uint64_t __kmpc_impl_pack(uint32_t lo, uint32_t hi);
+///}
+
+/// Fallback implementations of the shuffle sync idiom.
+///
+///{
+
+inline int32_t __kmpc_impl_shfl_sync(uint64_t Mask, int32_t Var,
+                                     int32_t SrcLane) {
+  assert(false &&
+         "Fallback version of __kmpc_impl_shfl_sync is not available!");
+}
+
+inline int32_t __kmpc_impl_shfl_down_sync(uint64_t Mask, int32_t Var,
+                                          uint32_t Delta, int32_t Width) {
+  assert(false &&
+         "Fallback version of __kmpc_impl_shfl_down_sync is not available!");
+}
+
+///}
+
+/// AMDGCN implementations of the shuffle sync idiom.
+///
+///{
+#pragma omp begin declare variant match(device = {arch(amdgcn)})
+
+inline int32_t __kmpc_impl_shfl_sync(uint64_t Mask, int32_t Var,
+                                     int32_t SrcLane) {
+  int Width = GetWarpSize();
+  int Self = GetLaneId();
+  int Index = SrcLane + (Self & ~(Width - 1));
+  return __builtin_amdgcn_ds_bpermute(Index << 2, Var);
+}
+
+inline int32_t __kmpc_impl_shfl_down_sync(uint64_t Mask, int32_t Var,
+                                          uint32_t LaneDelta, int32_t Width) {
+  int Self = GetLaneId();
+  int Index = Self + LaneDelta;
+  Index = (int)(LaneDelta + (Self & (Width - 1))) >= Width ? Self : Index;
+  return __builtin_amdgcn_ds_bpermute(Index << 2, Var);
+}
+
+#pragma omp end declare variant
+///}
+
+/// NVPTX implementations of the shuffle and shuffle sync idiom.
+///
+///{
+#pragma omp begin declare variant match(                                       \
+    device = {arch(nvptx, nvptx64)}, implementation = {extension(match_any)})
+
+inline int32_t __kmpc_impl_shfl_sync(uint64_t Mask, int32_t Var,
+                                     int32_t SrcLane) {
+  return __nvvm_shfl_idx_i32(Var, SrcLane, 0x1f);
+}
+
+inline int32_t __kmpc_impl_shfl_down_sync(uint64_t Mask, int32_t Var,
+                                          uint32_t Delta, int32_t Width) {
+  int32_t T = ((GetWarpSize() - Width) << 8) | 0x1f;
+  return __nvvm_shfl_down_i32(Var, Delta, T);
+}
+
+#pragma omp end declare variant
+///}
+
+#pragma omp end declare target
+
+#endif

openmp/libomptarget/deviceRTLs/common/src/data_sharing.cu

 //===----- data_sharing.cu - OpenMP GPU data sharing ------------- CUDA -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the implementation of data sharing environments
 //
 //===----------------------------------------------------------------------===//
 #pragma omp declare target
 
 #include "common/omptarget.h"
+#include "target/shuffle.h"
 #include "target_impl.h"
 
 // Return true if this is the master thread.
 INLINE static bool IsMasterThread(bool isSPMDExecutionMode) {
   return !isSPMDExecutionMode && GetMasterThreadID() == GetThreadIdInBlock();
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Runtime functions for trunk data sharing scheme.
 ////////////////////////////////////////////////////////////////////////////////

openmp/libomptarget/deviceRTLs/common/src/loop.cu

 //===------------ loop.cu - NVPTX OpenMP loop constructs --------- CUDA -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the implementation of the KMPC interface
 // for the loop construct plus other worksharing constructs that use the same
 // interface as loops.
 //
 //===----------------------------------------------------------------------===//
 #pragma omp declare target
 
 #include "common/omptarget.h"
+#include "target/shuffle.h"
 #include "target_impl.h"
 
 ////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////
 // template class that encapsulate all the helper functions
 //
 // T is loop iteration type (32 | 64)  (unsigned | signed)
 // ST is the signed version of T

openmp/libomptarget/deviceRTLs/common/src/reduction.cu

 //===---- reduction.cu - GPU OpenMP reduction implementation ----- CUDA -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the implementation of reduction with KMPC interface.
 //
 //===----------------------------------------------------------------------===//
 #pragma omp declare target
 
 #include "common/omptarget.h"
+#include "target/shuffle.h"
 #include "target_impl.h"
 
 EXTERN
 void __kmpc_nvptx_end_reduce(int32_t global_tid) {}
 
 EXTERN
 void __kmpc_nvptx_end_reduce_nowait(int32_t global_tid) {}
 
-EXTERN int32_t __kmpc_shuffle_int32(int32_t val, int16_t delta, int16_t size) {
-  return __kmpc_impl_shfl_down_sync(__kmpc_impl_all_lanes, val, delta, size);
-}
-
-EXTERN int64_t __kmpc_shuffle_int64(int64_t val, int16_t delta, int16_t size) {
-  uint32_t lo, hi;
-  __kmpc_impl_unpack(val, lo, hi);
-  hi = __kmpc_impl_shfl_down_sync(__kmpc_impl_all_lanes, hi, delta, size);
-  lo = __kmpc_impl_shfl_down_sync(__kmpc_impl_all_lanes, lo, delta, size);
-  return __kmpc_impl_pack(lo, hi);
-}
-
 INLINE static void gpu_regular_warp_reduce(void *reduce_data,
                                            kmp_ShuffleReductFctPtr shflFct) {
   for (uint32_t mask = WARPSIZE / 2; mask > 0; mask /= 2) {
     shflFct(reduce_data, /*LaneId - not used= */ 0,
             /*Offset = */ mask, /*AlgoVersion=*/0);
   }
 }
 

openmp/libomptarget/deviceRTLs/common/src/shuffle.cpp

+//===--- shuffle.cpp - Implementation of the external shuffle idiom API -*-===//
+//
+// 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 "target/shuffle.h"
+
+#pragma omp declare target
+
+static constexpr uint64_t AllLanes = -1;
+
+int32_t __kmpc_shuffle_int32(int32_t val, int16_t delta, int16_t size) {
+  return __kmpc_impl_shfl_down_sync(AllLanes, val, delta, size);
+}
+
+int64_t __kmpc_shuffle_int64(int64_t val, int16_t delta, int16_t size) {
+  uint32_t lo, hi;
+  __kmpc_impl_unpack(val, lo, hi);
+  hi = __kmpc_impl_shfl_down_sync(AllLanes, hi, delta, size);
+  lo = __kmpc_impl_shfl_down_sync(AllLanes, lo, delta, size);
+  return __kmpc_impl_pack(lo, hi);
+}
+
+#pragma omp end declare target

openmp/libomptarget/deviceRTLs/nvptx/CMakeLists.txt

set(cuda_src_files
   ${devicertl_common_directory}/src/loop.cu
   ${devicertl_common_directory}/src/omp_data.cu
   ${devicertl_common_directory}/src/omptarget.cu
   ${devicertl_common_directory}/src/parallel.cu
   ${devicertl_common_directory}/src/reduction.cu
   ${devicertl_common_directory}/src/support.cu
   ${devicertl_common_directory}/src/sync.cu
   ${devicertl_common_directory}/src/task.cu
+  ${devicertl_common_directory}/src/shuffle.cpp
   src/target_impl.cu
 )
 
 # Set flags for LLVM Bitcode compilation.
 set(bc_flags -S -x c++ -O1 -std=c++14
              -target nvptx64
              -Xclang -emit-llvm-bc
              -Xclang -aux-triple -Xclang ${aux_triple}
              -fopenmp -fopenmp-cuda-mode -Xclang -fopenmp-is-device
              -Xclang -target-feature -Xclang +ptx61
              -D__CUDACC__
              -I${devicertl_base_directory}
+             -I${devicertl_common_directory}/include
              -I${devicertl_nvptx_directory}/src)
 
 if(${LIBOMPTARGET_NVPTX_DEBUG})
   list(APPEND bc_flags -DOMPTARGET_NVPTX_DEBUG=-1)
 else()
   list(APPEND bc_flags -DOMPTARGET_NVPTX_DEBUG=0)
 endif()
 

openmp/libomptarget/deviceRTLs/nvptx/src/target_impl.cu

 }
 
 DEVICE __kmpc_impl_lanemask_t __kmpc_impl_activemask() {
   unsigned int Mask;
   asm volatile("activemask.b32 %0;" : "=r"(Mask));
   return Mask;
 }
 
-DEVICE int32_t __kmpc_impl_shfl_sync(__kmpc_impl_lanemask_t Mask, int32_t Var,
-                                     int32_t SrcLane) {
-  return __nvvm_shfl_sync_idx_i32(Mask, Var, SrcLane, 0x1f);
-}
-
-DEVICE int32_t __kmpc_impl_shfl_down_sync(__kmpc_impl_lanemask_t Mask,
-                                          int32_t Var, uint32_t Delta,
-                                          int32_t Width) {
-  int32_t T = ((WARPSIZE - Width) << 8) | 0x1f;
-  return __nvvm_shfl_sync_down_i32(Mask, Var, Delta, T);
-}
-
 DEVICE void __kmpc_impl_syncthreads() { __syncthreads(); }
 
 DEVICE void __kmpc_impl_syncwarp(__kmpc_impl_lanemask_t Mask) {
   __nvvm_bar_warp_sync(Mask);
 }
 
 // NVPTX specific kernel initialization
 DEVICE void __kmpc_impl_target_init() { /* nvptx needs no extra setup */

openmp/libomptarget/deviceRTLs/target_interface.h

 EXTERN void __kmpc_impl_unpack(uint64_t val, uint32_t &lo, uint32_t &hi);
 EXTERN uint64_t __kmpc_impl_pack(uint32_t lo, uint32_t hi);
 EXTERN __kmpc_impl_lanemask_t __kmpc_impl_lanemask_lt();
 EXTERN __kmpc_impl_lanemask_t __kmpc_impl_lanemask_gt();
 EXTERN uint32_t __kmpc_impl_smid();
 
 EXTERN __kmpc_impl_lanemask_t __kmpc_impl_activemask();
 
-EXTERN int32_t __kmpc_impl_shfl_sync(__kmpc_impl_lanemask_t Mask, int32_t Var,
-                                     int32_t SrcLane);
-EXTERN int32_t __kmpc_impl_shfl_down_sync(__kmpc_impl_lanemask_t Mask,
-                                          int32_t Var, uint32_t Delta,
-                                          int32_t Width);
-
 EXTERN void __kmpc_impl_syncthreads();
 EXTERN void __kmpc_impl_syncwarp(__kmpc_impl_lanemask_t Mask);
 
 // Kernel initialization
 EXTERN void __kmpc_impl_target_init();
 
 // Memory
 EXTERN void *__kmpc_impl_malloc(size_t);
 EXTERN void __kmpc_impl_free(void *);
 
 // Barrier until num_threads arrive.
 EXTERN void __kmpc_impl_named_sync(uint32_t num_threads);
 
 #endif // _OMPTARGET_TARGET_INTERFACE_H_