[Offload] Cleanups for bare kernels after recent upstream change. (llvm#2085)

Author: ronlieb

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

@@ -1048,10 +1048,9 @@ void CGOpenMPRuntimeGPU::emitTargetOutlinedFunction(
                     CGM.emitNxResult("[No-Loop/Big-Jump-Loop/Xteam]", D,
                                      CodeGenModule::NxNonSPMD));
   }
-  // Note that bare kernels always run in SPMD mode.
   setPropertyExecutionMode(
       CGM, OutlinedFn->getName(),
-      IsBareKernel ? OMP_TGT_EXEC_MODE_SPMD
+      IsBareKernel ? OMP_TGT_EXEC_MODE_BARE
                    : computeExecutionMode(Mode, DirectiveStmt, CGM));
 
   if (Mode && DirectiveStmt)

offload/plugins-nextgen/amdgpu/src/rtl.cpp

@@ -961,11 +961,10 @@ struct AMDGPUKernelTy : public GenericKernelTy {
   /// user-defined threads and block clauses.
   uint32_t getNumThreads(GenericDeviceTy &GenericDevice,
                          uint32_t ThreadLimitClause[3]) const override {
-    // On amd-staging, bare kernels go through this codepath. Legacy flang
-    // kernels show up as bare kernels since kernel-env is not generated.
-    // In order to accomodate bare ekrnels, disable this assert.
-    // assert(ThreadLimitClause[1] == 1 && ThreadLimitClause[2] == 1 &&
-    //       "Multi dimensional launch not supported yet.");
+    assert(!isBareMode() && "bare kernel should not call this function");
+
+    assert(ThreadLimitClause[1] == 1 && ThreadLimitClause[2] == 1 &&
+           "Multi dimensional launch not supported yet.");
 
     // Honor OMP_TEAMS_THREAD_LIMIT environment variable and
     // num_threads/thread_limit clause for BigJumpLoop and NoLoop kernel types.
@@ -1014,11 +1013,10 @@ struct AMDGPUKernelTy : public GenericKernelTy {
                         uint32_t NumTeamsClause[3], uint64_t LoopTripCount,
                         uint32_t &NumThreads,
                         bool IsNumThreadsFromUser) const override {
-    // On amd-staging, bare kernels go through this codepath. Legacy flang
-    // kernels show up as bare kernels since kernel-env is not generated.
-    // In order to accomodate bare ekrnels, disable this assert.
-    // assert(NumTeamsClause[1] == 1 && NumTeamsClause[2] == 1 &&
-    //      "Multi dimensional launch not supported yet.");
+    assert(!isBareMode() && "bare kernel should not call this function");
+
+    assert(NumTeamsClause[1] == 1 && NumTeamsClause[2] == 1 &&
+           "Multi dimensional launch not supported yet.");
 
     const auto getNumGroupsFromThreadsAndTripCount =
         [](const uint64_t TripCount, const uint32_t NumThreads) {

offload/plugins-nextgen/common/include/PluginInterface.h

@@ -356,6 +356,7 @@ struct GenericKernelTy {
     return ExecutionMode == OMP_TGT_EXEC_MODE_GENERIC;
   }
   bool isSPMDMode() const { return ExecutionMode == OMP_TGT_EXEC_MODE_SPMD; }
+  bool isBareMode() const { return ExecutionMode == OMP_TGT_EXEC_MODE_BARE; }
 
   /// AMD-only execution modes
   bool isBigJumpLoopMode() const {
@@ -374,7 +375,7 @@ struct GenericKernelTy {
 protected:
   /// Get the execution mode name of the kernel.
   const char *getExecutionModeName() const {
-    switch (KernelEnvironment.Configuration.ExecMode) {
+    switch (ExecutionMode) {
     case OMP_TGT_EXEC_MODE_BARE:
       return "BARE";
     case OMP_TGT_EXEC_MODE_SPMD:
@@ -440,10 +441,6 @@ struct GenericKernelTy {
                                 uint32_t BlockLimitClause[3],
                                 uint64_t LoopTripCount, uint32_t &NumThreads,
                                 bool IsNumThreadsFromUser) const;
-  /// Indicate if the kernel works in Generic SPMD, Generic or SPMD mode.
-  bool isBareMode() const {
-    return KernelEnvironment.Configuration.ExecMode == OMP_TGT_EXEC_MODE_BARE;
-  }
 
   /// The kernel name.
   const char *Name;

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -514,10 +514,10 @@ Error GenericKernelTy::init(GenericDeviceTy &GenericDevice,
     // Consume the error since it is acceptable to fail.
     [[maybe_unused]] std::string ErrStr = toString(std::move(Err));
     DP("Failed to read execution mode for '%s': %s\n"
-       "Using default SPMD (2) execution mode\n",
+       "Using default Bare (0) execution mode\n",
        Name, ErrStr.data());
 
-    ExecutionMode = OMP_TGT_EXEC_MODE_SPMD;
+    ExecutionMode = OMP_TGT_EXEC_MODE_BARE;
   } else {
     // Check that the retrieved execution mode is valid.
     if (!GenericKernelTy::isValidExecutionMode(ExecModeGlobal.getValue()))
@@ -744,12 +744,7 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
         KernelRecord->getLaunchParamsForKernel(*this, GenericDevice);
     NumBlocks[0] = Teams;
     NumThreads[0] = Threads;
-  } else {
-
-    // TODO fix workaround since IsBareKernel is not properly set for legacy
-    // flang and specialized kernels since they don't use kernel-env. While
-    // we can check for specialized kernels, we can't for legacy flang. So,
-    // on amd-staging, all kernels including bare ones use this codepath.
+  } else if (!isBareMode()) {
     NumThreads[0] = getNumThreads(GenericDevice, NumThreads);
 
     std::pair<bool, uint32_t> AdjustInfo = adjustNumThreadsForLowTripCount(
@@ -823,8 +818,7 @@ KernelLaunchParamsTy GenericKernelTy::prepareArgs(
 
 uint32_t GenericKernelTy::getNumThreads(GenericDeviceTy &GenericDevice,
                                         uint32_t ThreadLimitClause[3]) const {
-  // TODO fix workaround since IsBareKernel is not properly set for all kernels.
-  // assert(!IsBareKernel && "bare kernel should not call this function");
+  assert(!isBareMode() && "bare kernel should not call this function");
 
   assert(ThreadLimitClause[1] == 1 && ThreadLimitClause[2] == 1 &&
          "Multi dimensional launch not supported yet.");
@@ -846,8 +840,7 @@ uint32_t GenericKernelTy::getNumBlocks(GenericDeviceTy &GenericDevice,
                                        uint64_t LoopTripCount,
                                        uint32_t &NumThreads,
                                        bool IsNumThreadsFromUser) const {
-  // TODO fix workaround since IsBareKernel is not properly set for all kernels.
-  // assert(!IsBareKernel && "bare kernel should not call this function");
+  assert(!isBareMode() && "bare kernel should not call this function");
 
   assert(NumTeamsClause[1] == 1 && NumTeamsClause[2] == 1 &&
          "Multi dimensional launch not supported yet.");