[OpenMP] Use common interface to access GPU Grid Values

[saiislam]

Summary

Use common interface for accessing target specific GPU grid values in NVPTX
OpenMP codegen as proposed in https://reviews.llvm.org/D80917

Originally authored by Greg Rodgers (@gregrodgers).

[JonChesterfield]

Changing to getGridValue would be useful for sharing parts of this with amdgcn.

The aomp toolchain handles codegen for amdgcn by adding if (isAMDGCN) to this file. Until such time as tregions obsoletes this code, I think we should go with layers instead of scattered conditionals.

I.e. rename CGOpenMPRuntimeNVPTX to CGOpenMPRuntimeGPU which contains code that is common to nvptx and amdgcn. That probably uses getGridValue() as a way to abstract over minor differences. Derive CGOpenMPRuntimeAMDGCN and CGOpenMPRuntimeNVPTX from CGOpenMPRuntimeGPU to implement (virtual) functions which are different between the two.

clang/lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp

628 ↗

(On Diff #276767)

This looks unrelated to using the constants. Amdgcn doesn't have an nvvm_read_ptx_sreg_warpsize so does need a different means of accessing the wave size, but that's not directly related to using OMPGridValues

[jdoerfert]

clang/lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp

654 ↗

(On Diff #276767)

Why do we keep the enum value with this name then?

[saiislam]

In D83492#2142086, @JonChesterfield wrote:

Changing to getGridValue would be useful for sharing parts of this with amdgcn.

The aomp toolchain handles codegen for amdgcn by adding if (isAMDGCN) to this file. Until such time as tregions obsoletes this code, I think we should go with layers instead of scattered conditionals.

I.e. rename CGOpenMPRuntimeNVPTX to CGOpenMPRuntimeGPU which contains code that is common to nvptx and amdgcn. That probably uses getGridValue() as a way to abstract over minor differences. Derive CGOpenMPRuntimeAMDGCN and CGOpenMPRuntimeNVPTX from CGOpenMPRuntimeGPU to implement (virtual) functions which are different between the two.

Here is an implementation: D83723 . It also provides target specific implementation of getNVPTXWarpSize() as proof of concept.

[saiislam]

clang/lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp

654 ↗

(On Diff #276767)

Enums in OMPGridValues.h has been named like this because many of the values are dependent on a different value in the same enum. For example, hardcoded value of number of bits required to represent the max number of threads in a warp (LaneIDBits here), is computed as log of WarpSize. Similarly, mask of this value (LaneIDMask below) is based on WarpSize. A future change in WarpSize thus will require update in all the enums based on (and prefixed as) WarpSize, so less chances of a manual error.

See this: https://github.com/llvm/llvm-project/blob/8b6821a5843bb321b3738e2519beae7142e62928/llvm/include/llvm/Frontend/OpenMP/OMPGridValues.h#L120

[saiislam]

Removed getNVPTXWarpSize() changes for a separate patch and rebased.

[saiislam]

Corrected.

[jdoerfert]

LGTM

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

 #include "CGOpenMPRuntimeNVPTX.h"
 #include "CodeGenFunction.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/DeclOpenMP.h"
 #include "clang/AST/StmtOpenMP.h"
 #include "clang/AST/StmtVisitor.h"
 #include "clang/Basic/Cuda.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Frontend/OpenMP/OMPGridValues.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 
 using namespace clang;
 using namespace CodeGen;
 using namespace llvm::omp;
 
 namespace {
 enum OpenMPRTLFunctionNVPTX {
public:
   }
 };
 
 /// GPU Configuration:  This information can be derived from cuda registers,
 /// however, providing compile time constants helps generate more efficient
 /// code.  For all practical purposes this is fine because the configuration
 /// is the same for all known NVPTX architectures.
 enum MachineConfiguration : unsigned {
-  WarpSize = 32,
-  /// Number of bits required to represent a lane identifier, which is
-  /// computed as log_2(WarpSize).
-  LaneIDBits = 5,
-  LaneIDMask = WarpSize - 1,
+  /// See "llvm/Frontend/OpenMP/OMPGridValues.h" for various related target
+  /// specific Grid Values like GV_Warp_Size, GV_Warp_Size_Log2,
+  /// and GV_Warp_Size_Log2_Mask.
 
   /// Global memory alignment for performance.
   GlobalMemoryAlignment = 128,
 
   /// Maximal size of the shared memory buffer.
   SharedMemorySize = 128,
 };
 
for (const CapturedStmt::Capture &C : S->captures()) {
       }
     }
   }
 
   void buildRecordForGlobalizedVars(bool IsInTTDRegion) {
     assert(!GlobalizedRD &&
            "Record for globalized variables is built already.");
     ArrayRef<const ValueDecl *> EscapedDeclsForParallel, EscapedDeclsForTeams;
+    unsigned WarpSize = CGF.getTarget().getGridValue(llvm::omp::GV_Warp_Size);
     if (IsInTTDRegion)
       EscapedDeclsForTeams = EscapedDecls.getArrayRef();
     else
       EscapedDeclsForParallel = EscapedDecls.getArrayRef();
     GlobalizedRD = ::buildRecordForGlobalizedVars(
         CGF.getContext(), EscapedDeclsForParallel, EscapedDeclsForTeams,
         MappedDeclsFields, WarpSize);
   }
return CGF.EmitRuntimeCall(
       "nvptx_tid");
 }
 
 /// Get the id of the warp in the block.
 /// We assume that the warp size is 32, which is always the case
 /// on the NVPTX device, to generate more efficient code.
 static llvm::Value *getNVPTXWarpID(CodeGenFunction &CGF) {
   CGBuilderTy &Bld = CGF.Builder;
+  unsigned LaneIDBits =
+      CGF.getTarget().getGridValue(llvm::omp::GV_Warp_Size_Log2);
   return Bld.CreateAShr(getNVPTXThreadID(CGF), LaneIDBits, "nvptx_warp_id");
 }
 
 /// Get the id of the current lane in the Warp.
 /// We assume that the warp size is 32, which is always the case
 /// on the NVPTX device, to generate more efficient code.
 static llvm::Value *getNVPTXLaneID(CodeGenFunction &CGF) {
   CGBuilderTy &Bld = CGF.Builder;
+  unsigned LaneIDMask = CGF.getContext().getTargetInfo().getGridValue(
+      llvm::omp::GV_Warp_Size_Log2_Mask);
   return Bld.CreateAnd(getNVPTXThreadID(CGF), Bld.getInt32(LaneIDMask),
                        "nvptx_lane_id");
 }
 
 /// Get the maximum number of threads in a block of the GPU.
 static llvm::Value *getNVPTXNumThreads(CodeGenFunction &CGF) {
   return CGF.EmitRuntimeCall(
       llvm::Intrinsic::getDeclaration(

 llvm::Function *CGOpenMPRuntimeGPU::emitTeamsOutlinedFunction(
     const OMPExecutableDirective &D, const VarDecl *ThreadIDVar,
     OpenMPDirectiveKind InnermostKind, const RegionCodeGenTy &CodeGen) {
   SourceLocation Loc = D.getBeginLoc();
 
   const RecordDecl *GlobalizedRD = nullptr;
   llvm::SmallVector<const ValueDecl *, 4> LastPrivatesReductions;
   llvm::SmallDenseMap<const ValueDecl *, const FieldDecl *> MappedDeclsFields;
+  unsigned WarpSize = CGM.getTarget().getGridValue(llvm::omp::GV_Warp_Size);
   // Globalize team reductions variable unconditionally in all modes.
   if (getExecutionMode() != CGOpenMPRuntimeGPU::EM_SPMD)
     getTeamsReductionVars(CGM.getContext(), D, LastPrivatesReductions);
   if (getExecutionMode() == CGOpenMPRuntimeGPU::EM_SPMD) {
     getDistributeLastprivateVars(CGM.getContext(), D, LastPrivatesReductions);
     if (!LastPrivatesReductions.empty()) {
       GlobalizedRD = ::buildRecordForGlobalizedVars(
           CGM.getContext(), llvm::None, LastPrivatesReductions,
static llvm::Value *emitInterWarpCopyFunction(CodeGenModule &CGM,
   // (reduction across warps).  The array is placed in NVPTX __shared__ memory
   // for reduced latency, as well as to have a distinct copy for concurrently
   // executing target regions.  The array is declared with common linkage so
   // as to be shared across compilation units.
   StringRef TransferMediumName =
       "__openmp_nvptx_data_transfer_temporary_storage";
   llvm::GlobalVariable *TransferMedium =
       M.getGlobalVariable(TransferMediumName);
+  unsigned WarpSize = CGF.getTarget().getGridValue(llvm::omp::GV_Warp_Size);
   if (!TransferMedium) {
     auto *Ty = llvm::ArrayType::get(CGM.Int32Ty, WarpSize);
     unsigned SharedAddressSpace = C.getTargetAddressSpace(LangAS::cuda_shared);
     TransferMedium = new llvm::GlobalVariable(
         M, Ty, /*isConstant=*/false, llvm::GlobalVariable::CommonLinkage,
         llvm::Constant::getNullValue(Ty), TransferMediumName,
         /*InsertBefore=*/nullptr, llvm::GlobalVariable::NotThreadLocal,
         SharedAddressSpace);