Fine-grained control of kernel execution mode

Author: skatrak

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

@@ -20,6 +20,7 @@
 #include "clang/AST/StmtVisitor.h"
 #include "clang/Basic/Cuda.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Frontend/OpenMP/OMPDeviceConstants.h"
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
 
 using namespace clang;
@@ -748,7 +749,11 @@ void CGOpenMPRuntimeGPU::emitKernelInit(const OMPExecutableDirective &D,
   computeMinAndMaxThreadsAndTeams(D, CGF, Attrs);
 
   CGBuilderTy &Bld = CGF.Builder;
-  Bld.restoreIP(OMPBuilder.createTargetInit(Bld, IsSPMD, Attrs));
+  Bld.restoreIP(OMPBuilder.createTargetInit(
+      Bld,
+      IsSPMD ? llvm::omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_SPMD
+             : llvm::omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_GENERIC,
+      Attrs));
   if (!IsSPMD)
     emitGenericVarsProlog(CGF, EST.Loc);
 }

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -2243,21 +2243,21 @@ class OpenMPIRBuilder {
 
   /// Container to pass LLVM IR runtime values or constants related to the
   /// number of teams and threads with which the kernel must be launched, as
-  /// well as the trip count of the SPMD loop, if it is an SPMD kernel. These
-  /// must be defined in the host prior to the call to the kernel launch OpenMP
-  /// RTL function.
+  /// well as the trip count of the loop, if it is an SPMD or Generic-SPMD
+  /// kernel. These must be defined in the host prior to the call to the kernel
+  /// launch OpenMP RTL function.
   struct TargetKernelRuntimeAttrs {
     SmallVector<Value *, 3> MaxTeams = {nullptr};
     Value *MinTeams = nullptr;
     SmallVector<Value *, 3> TargetThreadLimit = {nullptr};
     SmallVector<Value *, 3> TeamsThreadLimit = {nullptr};
 
-    /// 'parallel' construct 'num_threads' clause value, if present and it is a
-    /// target SPMD kernel.
+    /// 'parallel' construct 'num_threads' clause value, if present and it is an
+    /// SPMD kernel.
     Value *MaxThreads = nullptr;
 
-    /// Total number of iterations of the target SPMD kernel or null if it is a
-    /// generic kernel.
+    /// Total number of iterations of the SPMD or Generic-SPMD kernel or null if
+    /// it is a generic kernel.
     Value *LoopTripCount = nullptr;
   };
 
@@ -2763,11 +2763,12 @@ class OpenMPIRBuilder {
   /// Create a runtime call for kmpc_target_init
   ///
   /// \param Loc The insert and source location description.
+  /// \param ExecFlags Kernel execution mode flags.
   /// \param IsSPMD Flag to indicate if the kernel is an SPMD kernel or not.
   /// \param Attrs Structure containing the default numbers of threads and teams
   ///        to launch the kernel with.
   InsertPointTy createTargetInit(
-      const LocationDescription &Loc, bool IsSPMD,
+      const LocationDescription &Loc, omp::OMPTgtExecModeFlags ExecFlags,
       const llvm::OpenMPIRBuilder::TargetKernelDefaultAttrs &Attrs);
 
   /// Create a runtime call for kmpc_target_deinit
@@ -2929,7 +2930,7 @@ class OpenMPIRBuilder {
   ///
   /// \param Loc where the target data construct was encountered.
   /// \param IsOffloadEntry whether it is an offload entry.
-  /// \param IsSPMD whether it is a target SPMD kernel.
+  /// \param ExecFlags kernel execution mode flags.
   /// \param CodeGenIP The insertion point where the call to the outlined
   /// function should be emitted.
   /// \param EntryInfo The entry information about the function.
@@ -2946,7 +2947,8 @@ class OpenMPIRBuilder {
   // dependency information as passed in the depend clause
   // \param HasNowait Whether the target construct has a `nowait` clause or not.
   InsertPointOrErrorTy createTarget(
-      const LocationDescription &Loc, bool IsOffloadEntry, bool IsSPMD,
+      const LocationDescription &Loc, bool IsOffloadEntry,
+      omp::OMPTgtExecModeFlags ExecFlags,
       OpenMPIRBuilder::InsertPointTy AllocaIP,
       OpenMPIRBuilder::InsertPointTy CodeGenIP,
       TargetRegionEntryInfo &EntryInfo,

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -6124,7 +6124,7 @@ CallInst *OpenMPIRBuilder::createCachedThreadPrivate(
 }
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTargetInit(
-    const LocationDescription &Loc, bool IsSPMD,
+    const LocationDescription &Loc, omp::OMPTgtExecModeFlags ExecFlags,
     const llvm::OpenMPIRBuilder::TargetKernelDefaultAttrs &Attrs) {
   assert(!Attrs.MaxThreads.empty() && !Attrs.MaxTeams.empty() &&
          "expected num_threads and num_teams to be specified");
@@ -6135,9 +6135,9 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTargetInit(
   uint32_t SrcLocStrSize;
   Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
   Constant *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
-  Constant *IsSPMDVal = ConstantInt::getSigned(
-      Int8, IsSPMD ? OMP_TGT_EXEC_MODE_SPMD : OMP_TGT_EXEC_MODE_GENERIC);
-  Constant *UseGenericStateMachineVal = ConstantInt::getSigned(Int8, !IsSPMD);
+  Constant *IsSPMDVal = ConstantInt::getSigned(Int8, ExecFlags);
+  Constant *UseGenericStateMachineVal =
+      ConstantInt::getSigned(Int8, ExecFlags != omp::OMP_TGT_EXEC_MODE_SPMD);
   Constant *MayUseNestedParallelismVal = ConstantInt::getSigned(Int8, true);
   Constant *DebugIndentionLevelVal = ConstantInt::getSigned(Int16, 0);
 
@@ -6742,7 +6742,8 @@ FunctionCallee OpenMPIRBuilder::createDispatchDeinitFunction() {
 }
 
 static Expected<Function *> createOutlinedFunction(
-    OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder, bool IsSPMD,
+    OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
+    omp::OMPTgtExecModeFlags ExecFlags,
     const OpenMPIRBuilder::TargetKernelDefaultAttrs &DefaultAttrs,
     StringRef FuncName, SmallVectorImpl<Value *> &Inputs,
     OpenMPIRBuilder::TargetBodyGenCallbackTy &CBFunc,
@@ -6773,8 +6774,7 @@ static Expected<Function *> createOutlinedFunction(
       Function::Create(FuncType, GlobalValue::InternalLinkage, FuncName, M);
 
   if (OMPBuilder.Config.isTargetDevice()) {
-    Value *ExecMode = OMPBuilder.emitKernelExecutionMode(
-        FuncName, IsSPMD ? OMP_TGT_EXEC_MODE_SPMD : OMP_TGT_EXEC_MODE_GENERIC);
+    Value *ExecMode = OMPBuilder.emitKernelExecutionMode(FuncName, ExecFlags);
     OMPBuilder.emitUsed("llvm.compiler.used", {ExecMode});
   }
 
@@ -6818,7 +6818,7 @@ static Expected<Function *> createOutlinedFunction(
   // Insert target init call in the device compilation pass.
   if (OMPBuilder.Config.isTargetDevice())
     Builder.restoreIP(
-        OMPBuilder.createTargetInit(Builder, IsSPMD, DefaultAttrs));
+        OMPBuilder.createTargetInit(Builder, ExecFlags, DefaultAttrs));
 
   BasicBlock *UserCodeEntryBB = Builder.GetInsertBlock();
 
@@ -7014,7 +7014,7 @@ static Function *emitTargetTaskProxyFunction(OpenMPIRBuilder &OMPBuilder,
 
 static Error emitTargetOutlinedFunction(
     OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder, bool IsOffloadEntry,
-    bool IsSPMD, TargetRegionEntryInfo &EntryInfo,
+    omp::OMPTgtExecModeFlags ExecFlags, TargetRegionEntryInfo &EntryInfo,
     const OpenMPIRBuilder::TargetKernelDefaultAttrs &DefaultAttrs,
     Function *&OutlinedFn, Constant *&OutlinedFnID,
     SmallVectorImpl<Value *> &Inputs,
@@ -7023,8 +7023,8 @@ static Error emitTargetOutlinedFunction(
 
   OpenMPIRBuilder::FunctionGenCallback &&GenerateOutlinedFunction =
       [&](StringRef EntryFnName) {
-        return createOutlinedFunction(OMPBuilder, Builder, IsSPMD, DefaultAttrs,
-                                      EntryFnName, Inputs, CBFunc,
+        return createOutlinedFunction(OMPBuilder, Builder, ExecFlags,
+                                      DefaultAttrs, EntryFnName, Inputs, CBFunc,
                                       ArgAccessorFuncCB);
       };
 
@@ -7484,9 +7484,9 @@ emitTargetCall(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
 }
 
 OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTarget(
-    const LocationDescription &Loc, bool IsOffloadEntry, bool IsSPMD,
-    InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
-    TargetRegionEntryInfo &EntryInfo,
+    const LocationDescription &Loc, bool IsOffloadEntry,
+    omp::OMPTgtExecModeFlags ExecFlags, InsertPointTy AllocaIP,
+    InsertPointTy CodeGenIP, TargetRegionEntryInfo &EntryInfo,
     const TargetKernelDefaultAttrs &DefaultAttrs,
     const TargetKernelRuntimeAttrs &RuntimeAttrs,
     SmallVectorImpl<Value *> &Args, GenMapInfoCallbackTy GenMapInfoCB,
@@ -7505,7 +7505,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTarget(
   // the target region itself is generated using the callbacks CBFunc
   // and ArgAccessorFuncCB
   if (Error Err = emitTargetOutlinedFunction(
-          *this, Builder, IsOffloadEntry, IsSPMD, EntryInfo, DefaultAttrs,
+          *this, Builder, IsOffloadEntry, ExecFlags, EntryInfo, DefaultAttrs,
           OutlinedFn, OutlinedFnID, Args, CBFunc, ArgAccessorFuncCB))
     return Err;
 

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

@@ -6189,7 +6189,8 @@ TEST_F(OpenMPIRBuilderTest, TargetRegion) {
   RuntimeAttrs.TeamsThreadLimit[0] = Builder.getInt32(30);
   RuntimeAttrs.MaxThreads = Builder.getInt32(40);
   OpenMPIRBuilder::InsertPointOrErrorTy AfterIP = OMPBuilder.createTarget(
-      OmpLoc, /*IsOffloadEntry=*/true, /*IsSPMD=*/false, Builder.saveIP(),
+      OmpLoc, /*IsOffloadEntry=*/true,
+      omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_GENERIC, Builder.saveIP(),
       Builder.saveIP(), EntryInfo, DefaultAttrs, RuntimeAttrs, Inputs,
       GenMapInfoCB, BodyGenCB, SimpleArgAccessorCB);
   assert(AfterIP && "unexpected error");
@@ -6340,7 +6341,8 @@ TEST_F(OpenMPIRBuilderTest, TargetRegionDevice) {
       /*MaxTeams=*/{-1}, /*MinTeams=*/0, /*MaxThreads=*/{0}, /*MinThreads=*/0};
   OpenMPIRBuilder::TargetKernelRuntimeAttrs RuntimeAttrs;
   OpenMPIRBuilder::InsertPointOrErrorTy AfterIP = OMPBuilder.createTarget(
-      Loc, /*IsOffloadEntry=*/true, /*IsSPMD=*/false, EntryIP, EntryIP,
+      Loc, /*IsOffloadEntry=*/true,
+      omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_GENERIC, EntryIP, EntryIP,
       EntryInfo, DefaultAttrs, RuntimeAttrs, CapturedArgs, GenMapInfoCB,
       BodyGenCB, SimpleArgAccessorCB);
   assert(AfterIP && "unexpected error");
@@ -6480,7 +6482,8 @@ TEST_F(OpenMPIRBuilderTest, TargetRegionSPMD) {
   OpenMPIRBuilder::TargetKernelRuntimeAttrs RuntimeAttrs;
   RuntimeAttrs.LoopTripCount = Builder.getInt64(1000);
   OpenMPIRBuilder::InsertPointOrErrorTy AfterIP = OMPBuilder.createTarget(
-      OmpLoc, /*IsOffloadEntry=*/true, /*IsSPMD=*/true, Builder.saveIP(),
+      OmpLoc, /*IsOffloadEntry=*/true,
+      omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_SPMD, Builder.saveIP(),
       Builder.saveIP(), EntryInfo, DefaultAttrs, RuntimeAttrs, Inputs,
       GenMapInfoCB, BodyGenCB, SimpleArgAccessorCB);
   assert(AfterIP && "unexpected error");
@@ -6580,7 +6583,8 @@ TEST_F(OpenMPIRBuilderTest, TargetRegionDeviceSPMD) {
       /*MaxTeams=*/{-1}, /*MinTeams=*/0, /*MaxThreads=*/{0}, /*MinThreads=*/0};
   OpenMPIRBuilder::TargetKernelRuntimeAttrs RuntimeAttrs;
   OpenMPIRBuilder::InsertPointOrErrorTy AfterIP = OMPBuilder.createTarget(
-      Loc, /*IsOffloadEntry=*/true, /*IsSPMD=*/true, EntryIP, EntryIP,
+      Loc, /*IsOffloadEntry=*/true,
+      omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_SPMD, EntryIP, EntryIP,
       EntryInfo, DefaultAttrs, RuntimeAttrs, CapturedArgs, GenMapInfoCB,
       BodyGenCB, SimpleArgAccessorCB);
   assert(AfterIP && "unexpected error");
@@ -6686,7 +6690,8 @@ TEST_F(OpenMPIRBuilderTest, ConstantAllocaRaise) {
       /*MaxTeams=*/{-1}, /*MinTeams=*/0, /*MaxThreads=*/{0}, /*MinThreads=*/0};
   OpenMPIRBuilder::TargetKernelRuntimeAttrs RuntimeAttrs;
   OpenMPIRBuilder::InsertPointOrErrorTy AfterIP = OMPBuilder.createTarget(
-      Loc, /*IsOffloadEntry=*/true, /*IsSPMD=*/false, EntryIP, EntryIP,
+      Loc, /*IsOffloadEntry=*/true,
+      omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_GENERIC, EntryIP, EntryIP,
       EntryInfo, DefaultAttrs, RuntimeAttrs, CapturedArgs, GenMapInfoCB,
       BodyGenCB, SimpleArgAccessorCB);
   assert(AfterIP && "unexpected error");

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -3975,9 +3975,10 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
 
   llvm::OpenMPIRBuilder::InsertPointOrErrorTy afterIP =
       moduleTranslation.getOpenMPBuilder()->createTarget(
-          ompLoc, isOffloadEntry, /*IsSPMD=*/false, allocaIP, builder.saveIP(),
-          entryInfo, defaultAttrs, runtimeAttrs, kernelInput, genMapInfoCB,
-          bodyCB, argAccessorCB, dds, targetOp.getNowait());
+          ompLoc, isOffloadEntry, llvm::omp::OMP_TGT_EXEC_MODE_GENERIC,
+          allocaIP, builder.saveIP(), entryInfo, defaultAttrs, runtimeAttrs,
+          kernelInput, genMapInfoCB, bodyCB, argAccessorCB, dds,
+          targetOp.getNowait());
 
   if (failed(handleError(afterIP, opInst)))
     return failure();