[OpenMP][NFCI] Cleanup new device RT mapping interface

[jdoerfert]

Summary

Minimize the impl interface and clean up some uses of mapping
functions.

[jhuber6]

openmp/libomptarget/DeviceRTL/src/Mapping.cpp

242–251

Why do we want to get rid of these? I thought the idea was to keep these functions so we can generate generic IR in Clang that gets dispatched by the runtime once it's linked in.

[jdoerfert]

openmp/libomptarget/DeviceRTL/src/Mapping.cpp

242–251

I think we should replace runtime call folding with store->load forwarding. The new runtime does generic store->load forwarding for IsSPMD, as an example, and it works better as we do not need to keep calls around. That said, there is an argument made for both to exist. However, even if, we can simply look for the mangled names of the new runtime once we switched over rather than some __kmpc names.

[jhuber6]

openmp/libomptarget/DeviceRTL/src/Mapping.cpp

242–251

We use these in OpenMPOpt and CGOpenMPRuntimeGPU, could we replace the definitions in OMPKinds.def with the mangled names?

[JonChesterfield]

Cleanups good, be a shame to lose the functions clang calls into just as we're on the edge of deleting the amdgcn and nvptx subclasses from openmp codegen.

Is getNumberOfProcessorElements the number of SMs / CUs? That may only be known accurately at application runtime

openmp/libomptarget/DeviceRTL/src/Mapping.cpp

242–251

The _kmpc prefix is used for the rest of the functions clang emits calls to, not that keen to replace all those with mangled symbols written in clang

[JonChesterfield]

openmp/libomptarget/DeviceRTL/include/Interface.h

204 ↗

(On Diff #380997)

These, plus one more nyi, let clang emit 'gpu' IR with calls to these instead of platform specific intrinsics

[jdoerfert]

I can put them back if clang needs them.

[tianshilei1992]

openmp/libomptarget/DeviceRTL/src/Mapping.cpp

242–251

We also (plan to) use these function calls to replace the emission of intrinsics directly from front end, aka those functions in CGOpenMPRuntimeGPUAMDGCN and CGOpenMPRuntimeGPUNVPTX.

[jdoerfert]

Add the generic external interface back (but somewhat hidden)

[jhuber6]

Did you test this with SPMDzation? I'm pretty sure we generate a call to __kmpc_get_hardware_thread_id_in_block for that, and needed to keep it alive or else there would be no definition in the module.

[jdoerfert]

Add assertions, keep external interface, fix amdgpu

[jdoerfert]

Fix oversight

[jhuber6]

LGTM

[jdoerfert]

Fix more rebasing oversights

openmp/libomptarget/DeviceRTL/src/Mapping.cpp

 //===------- Mapping.cpp - OpenMP device runtime mapping helpers -- C++ -*-===//
 //
 // 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 "Mapping.h"
+#include "Interface.h"
 #include "State.h"
 #include "Types.h"
 #include "Utils.h"
 
 #pragma omp declare target
 
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
 

 }
 
 uint32_t getWorkgroupDim(uint32_t group_id, uint32_t grid_size,
                          uint16_t group_size) {
   uint32_t r = grid_size - group_id * group_size;
   return (r < group_size) ? r : group_size;
 }
 
+uint32_t getNumHardwareThreadsInBlock() {
+  return getWorkgroupDim(__builtin_amdgcn_workgroup_id_x(),
+                         __builtin_amdgcn_grid_size_x(),
+                         __builtin_amdgcn_workgroup_size_x());
+}
+
 LaneMaskTy activemask() { return __builtin_amdgcn_read_exec(); }
 
 LaneMaskTy lanemaskLT() {
   uint32_t Lane = mapping::getThreadIdInWarp();
   int64_t Ballot = mapping::activemask();
   uint64_t Mask = ((uint64_t)1 << Lane) - (uint64_t)1;
   return Mask & Ballot;
 }
 
 LaneMaskTy lanemaskGT() {
   uint32_t Lane = mapping::getThreadIdInWarp();
   if (Lane == (mapping::getWarpSize() - 1))
     return 0;
   int64_t Ballot = mapping::activemask();
   uint64_t Mask = (~((uint64_t)0)) << (Lane + 1);
   return Mask & Ballot;
 }
 
 uint32_t getThreadIdInWarp() {
   return __builtin_amdgcn_mbcnt_hi(~0u, __builtin_amdgcn_mbcnt_lo(~0u, 0u));
 }
 
 uint32_t getThreadIdInBlock() { return __builtin_amdgcn_workitem_id_x(); }
 
-uint32_t getBlockSize() {
-  // TODO: verify this logic for generic mode.
-  return getWorkgroupDim(__builtin_amdgcn_workgroup_id_x(),
-                         __builtin_amdgcn_grid_size_x(),
-                         __builtin_amdgcn_workgroup_size_x());
-}
-
 uint32_t getKernelSize() { return __builtin_amdgcn_grid_size_x(); }
 
 uint32_t getBlockId() { return __builtin_amdgcn_workgroup_id_x(); }
 
 uint32_t getNumberOfBlocks() {
   return getGridDim(__builtin_amdgcn_grid_size_x(),
                     __builtin_amdgcn_workgroup_size_x());
 }
 
-uint32_t getNumberOfProcessorElements() {
-  return getBlockSize();
-}
-
 uint32_t getWarpId() {
-  return mapping::getThreadIdInBlock() / mapping::getWarpSize();
+  return impl::getThreadIdInBlock() / mapping::getWarpSize();
 }
 
 uint32_t getNumberOfWarpsInBlock() {
   return mapping::getBlockSize() / mapping::getWarpSize();
 }
 
 #pragma omp end declare variant
 ///}
 
 /// NVPTX Implementation
 ///
 ///{
 #pragma omp begin declare variant match(                                       \
     device = {arch(nvptx, nvptx64)}, implementation = {extension(match_any)})
 
+uint32_t getNumHardwareThreadsInBlock() {
+  return __nvvm_read_ptx_sreg_ntid_x();
+}
+
 constexpr const llvm::omp::GV &getGridValue() {
   return llvm::omp::NVPTXGridValues;
 }
 
 LaneMaskTy activemask() {
   unsigned int Mask;
   asm("activemask.b32 %0;" : "=r"(Mask));
   return Mask;
 }
 
 LaneMaskTy lanemaskLT() {
   __kmpc_impl_lanemask_t Res;
   asm("mov.u32 %0, %%lanemask_lt;" : "=r"(Res));
   return Res;
 }
 
 LaneMaskTy lanemaskGT() {
   __kmpc_impl_lanemask_t Res;
   asm("mov.u32 %0, %%lanemask_gt;" : "=r"(Res));
   return Res;
 }
 
-uint32_t getThreadIdInWarp() {
-  return mapping::getThreadIdInBlock() & (mapping::getWarpSize() - 1);
-}
-
 uint32_t getThreadIdInBlock() { return __nvvm_read_ptx_sreg_tid_x(); }
 
-uint32_t getBlockSize() {
-  return __nvvm_read_ptx_sreg_ntid_x() -
-         (!mapping::isSPMDMode() * mapping::getWarpSize());
+uint32_t getThreadIdInWarp() {
+  return impl::getThreadIdInBlock() & (mapping::getWarpSize() - 1);
 }
 
-uint32_t getKernelSize() { return __nvvm_read_ptx_sreg_nctaid_x(); }
+uint32_t getKernelSize() {
+  return __nvvm_read_ptx_sreg_nctaid_x() *
+         mapping::getNumberOfProcessorElements();
+}
 
 uint32_t getBlockId() { return __nvvm_read_ptx_sreg_ctaid_x(); }
 
 uint32_t getNumberOfBlocks() { return __nvvm_read_ptx_sreg_nctaid_x(); }
 
-uint32_t getNumberOfProcessorElements() {
-  return __nvvm_read_ptx_sreg_ntid_x();
-}
-
 uint32_t getWarpId() {
-  return mapping::getThreadIdInBlock() / mapping::getWarpSize();
+  return impl::getThreadIdInBlock() / mapping::getWarpSize();
 }
 
 uint32_t getNumberOfWarpsInBlock() {
   return (mapping::getBlockSize() + mapping::getWarpSize() - 1) /
          mapping::getWarpSize();
 }
 
 #pragma omp end declare variant
 ///}
 
 uint32_t getWarpSize() { return getGridValue().GV_Warp_Size; }
 
 } // namespace impl
 } // namespace _OMP
 
+/// We have to be deliberate about the distinction of `mapping::` and `impl::`
+/// below to avoid repeating assumptions or including irrelevant ones.
+///{
+
 static bool isInLastWarp() {
   uint32_t MainTId = (mapping::getNumberOfProcessorElements() - 1) &
                      ~(mapping::getWarpSize() - 1);
   return mapping::getThreadIdInBlock() == MainTId;
 }
 
 bool mapping::isMainThreadInGenericMode(bool IsSPMD) {
   if (IsSPMD || icv::Level)

 }
 
 LaneMaskTy mapping::activemask() { return impl::activemask(); }
 
 LaneMaskTy mapping::lanemaskLT() { return impl::lanemaskLT(); }
 
 LaneMaskTy mapping::lanemaskGT() { return impl::lanemaskGT(); }
 
-uint32_t mapping::getThreadIdInWarp() { return impl::getThreadIdInWarp(); }
+uint32_t mapping::getThreadIdInWarp() {
+  uint32_t ThreadIdInWarp = impl::getThreadIdInWarp();
+  ASSERT(ThreadIdInWarp < impl::getWarpSize());
+  return ThreadIdInWarp;
+}
+
+uint32_t mapping::getThreadIdInBlock() {
+  uint32_t ThreadIdInBlock = impl::getThreadIdInBlock();
+  ASSERT(ThreadIdInBlock < impl::getNumHardwareThreadsInBlock());
+  return ThreadIdInBlock;
+}
 
-uint32_t mapping::getThreadIdInBlock() { return impl::getThreadIdInBlock(); }
+uint32_t mapping::getWarpSize() { return impl::getWarpSize(); }
 
-uint32_t mapping::getBlockSize() { return impl::getBlockSize(); }
+uint32_t mapping::getBlockSize() {
+  uint32_t BlockSize = mapping::getNumberOfProcessorElements() -
+                       (!mapping::isSPMDMode() * impl::getWarpSize());
+  return BlockSize;
+}
 
 uint32_t mapping::getKernelSize() { return impl::getKernelSize(); }
 
-uint32_t mapping::getBlockId() { return impl::getBlockId(); }
-
-uint32_t mapping::getNumberOfBlocks() { return impl::getNumberOfBlocks(); }
+uint32_t mapping::getWarpId() {
+  uint32_t WarpID = impl::getWarpId();
+  ASSERT(WarpID < impl::getNumberOfWarpsInBlock());
+  return WarpID;
+}
 
-uint32_t mapping::getNumberOfProcessorElements() {
-  return impl::getNumberOfProcessorElements();
+uint32_t mapping::getBlockId() {
+  uint32_t BlockId = impl::getBlockId();
+  ASSERT(BlockId < impl::getNumberOfBlocks());
+  return BlockId;
 }
 
-uint32_t mapping::getWarpId() { return impl::getWarpId(); }
+uint32_t mapping::getNumberOfWarpsInBlock() {
+  uint32_t NumberOfWarpsInBlocks = impl::getNumberOfWarpsInBlock();
+  ASSERT(impl::getWarpId() < NumberOfWarpsInBlocks);
+  return NumberOfWarpsInBlocks;
+}
 
-uint32_t mapping::getWarpSize() { return impl::getWarpSize(); }
+uint32_t mapping::getNumberOfBlocks() {
+  uint32_t NumberOfBlocks = impl::getNumberOfBlocks();
+  ASSERT(impl::getBlockId() < NumberOfBlocks);
+  return NumberOfBlocks;
+}
 
-uint32_t mapping::getNumberOfWarpsInBlock() {
-  return impl::getNumberOfWarpsInBlock();
+uint32_t mapping::getNumberOfProcessorElements() {
+  uint32_t NumberOfProcessorElements = impl::getNumHardwareThreadsInBlock();
+  ASSERT(impl::getThreadIdInBlock() < NumberOfProcessorElements);
+  return NumberOfProcessorElements;
 }
 
+///}
+
 /// Execution mode
 ///
 ///{
 static int SHARED(IsSPMDMode);
 
 void mapping::init(bool IsSPMD) {
   if (mapping::isInitialThreadInLevel0(IsSPMD))
     IsSPMDMode = IsSPMD;
 }
 
 bool mapping::isSPMDMode() { return IsSPMDMode; }
 
 bool mapping::isGenericMode() { return !isSPMDMode(); }
 ///}
 
 extern "C" {
 __attribute__((noinline)) uint32_t __kmpc_get_hardware_thread_id_in_block() {
   FunctionTracingRAII();
   return mapping::getThreadIdInBlock();
 }
 
 __attribute__((noinline)) uint32_t __kmpc_get_hardware_num_threads_in_block() {
   FunctionTracingRAII();
-  return mapping::getNumberOfProcessorElements();
+  return impl::getNumHardwareThreadsInBlock();
 }
 }
 #pragma omp end declare target