[OpenMP] Honor `thread_limit` value when choosing grid size

[tianshilei1992]

Summary

D152014 introduced an optimization that favors more smaller blocks over
fewer larger blocks, even if user sets thread_limit explicitly. This patch changes
the behavior to honor user value.

[jdoerfert]

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

379

Just check if the user set anything (ThreadLimit[0] > 0), we might still reduce it before but we should not reduce it further here.

448–449

You don't need the changes above if you add || IsNumThreadsFromUser above. That makes it clear we keep the value and just compute num teams based on it. Rather than having all these intermediate computations we do not use.

471

[tianshilei1992]

rebase and fix comments

[jdoerfert]

See comment. LG

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

471

With the change in the condition above, the old comment was good.

[tianshilei1992]

rebase and fix comment

[jplehr]

This broke the AMD GPU buildbot (https://lab.llvm.org/buildbot/#/builders/193/builds/37371)

[tianshilei1992]

In D158802#4620652, @jplehr wrote:

This broke the AMD GPU buildbot (https://lab.llvm.org/buildbot/#/builders/193/builds/37371)

In D158802#4620652, @jplehr wrote:

This broke the AMD GPU buildbot (https://lab.llvm.org/buildbot/#/builders/193/builds/37371)

Thanks. I’ll fix it today. Just need one line of change.

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

void *prepareArgs(GenericDeviceTy &GenericDevice, void **ArgPtrs,
                     llvm::SmallVectorImpl<void *> &Ptrs) const;
 
   /// Get the number of threads and blocks for the kernel based on the
   /// user-defined threads and block clauses.
   uint32_t getNumThreads(GenericDeviceTy &GenericDevice,
                          uint32_t ThreadLimitClause[3]) const;
 
   /// The number of threads \p NumThreads can be adjusted by this method.
+  /// \p IsNumThreadsFromUser is true is \p NumThreads is defined by user via
+  /// thread_limit clause.
   uint64_t getNumBlocks(GenericDeviceTy &GenericDevice,
                         uint32_t BlockLimitClause[3], uint64_t LoopTripCount,
-                        uint32_t &NumThreads) const;
+                        uint32_t &NumThreads, bool IsNumThreadsFromUser) const;
 
   /// Indicate if the kernel works in Generic SPMD, Generic or SPMD mode.
   bool isGenericSPMDMode() const {
     return ExecutionMode == OMP_TGT_EXEC_MODE_GENERIC_SPMD;
   }
   bool isGenericMode() const {
     return ExecutionMode == OMP_TGT_EXEC_MODE_GENERIC;
   }

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
                               AsyncInfoWrapperTy &AsyncInfoWrapper) const {
   llvm::SmallVector<void *, 16> Args;
   llvm::SmallVector<void *, 16> Ptrs;
 
   void *KernelArgsPtr = prepareArgs(GenericDevice, ArgPtrs, ArgOffsets,
                                     KernelArgs.NumArgs, Args, Ptrs);
 
   uint32_t NumThreads = getNumThreads(GenericDevice, KernelArgs.ThreadLimit);
-  uint64_t NumBlocks = getNumBlocks(GenericDevice, KernelArgs.NumTeams,
-                                    KernelArgs.Tripcount, NumThreads);
+  uint64_t NumBlocks =
+      getNumBlocks(GenericDevice, KernelArgs.NumTeams, KernelArgs.Tripcount,
+                   NumThreads, KernelArgs.ThreadLimit[0] > 0);
 
   if (auto Err =
           printLaunchInfo(GenericDevice, KernelArgs, NumThreads, NumBlocks))
     return Err;
 
   return launchImpl(GenericDevice, NumThreads, NumBlocks, KernelArgs,
                     KernelArgsPtr, AsyncInfoWrapper);
 }
uint32_t GenericKernelTy::getNumThreads(GenericDeviceTy &GenericDevice,
   return std::min(MaxNumThreads, (ThreadLimitClause[0] > 0)
                                      ? ThreadLimitClause[0]
                                      : PreferredNumThreads);
 }
 
 uint64_t GenericKernelTy::getNumBlocks(GenericDeviceTy &GenericDevice,
                                        uint32_t NumTeamsClause[3],
                                        uint64_t LoopTripCount,
-                                       uint32_t &NumThreads) const {
+                                       uint32_t &NumThreads,
+                                       bool IsNumThreadsFromUser) const {
   assert(NumTeamsClause[1] == 0 && NumTeamsClause[2] == 0 &&
          "Multi dimensional launch not supported yet.");
 
   if (NumTeamsClause[0] > 0) {
     // TODO: We need to honor any value and consequently allow more than the
     // block limit. For this we might need to start multiple kernels or let the
     // blocks start again until the requested number has been started.
     return std::min(NumTeamsClause[0], GenericDevice.getBlockLimit());
   }
 
   uint64_t DefaultNumBlocks = GenericDevice.getDefaultNumBlocks();
   uint64_t TripCountNumBlocks = std::numeric_limits<uint64_t>::max();
   if (LoopTripCount > 0) {
     if (isSPMDMode()) {
       // We have a combined construct, i.e. `target teams distribute
       // parallel for [simd]`. We launch so many teams so that each thread
       // will execute one iteration of the loop; rounded up to the nearest
       // integer. However, if that results in too few teams, we artificially
       // reduce the thread count per team to increase the outer parallelism.
       auto MinThreads = GenericDevice.getMinThreadsForLowTripCountLoop();
       MinThreads = std::min(MinThreads, NumThreads);
 
       // Honor the thread_limit clause; only lower the number of threads.
       [[maybe_unused]] auto OldNumThreads = NumThreads;
-      if (LoopTripCount >= DefaultNumBlocks * NumThreads) {
+      if (LoopTripCount >= DefaultNumBlocks * NumThreads ||
+          IsNumThreadsFromUser) {
         // Enough parallelism for teams and threads.
         TripCountNumBlocks = ((LoopTripCount - 1) / NumThreads) + 1;
         assert(TripCountNumBlocks >= DefaultNumBlocks &&
                "Expected sufficient outer parallelism.");
       } else if (LoopTripCount >= DefaultNumBlocks * MinThreads) {
         // Enough parallelism for teams, limit threads.
 
         // This case is hard; for now, we force "full warps":
         // First, compute a thread count assuming DefaultNumBlocks.
         auto NumThreadsDefaultBlocks =
             (LoopTripCount + DefaultNumBlocks - 1) / DefaultNumBlocks;
         // Now get a power of two that is larger or equal.
         auto NumThreadsDefaultBlocksP2 =
             llvm::PowerOf2Ceil(NumThreadsDefaultBlocks);
         // Do not increase a thread limit given be the user.
         NumThreads = std::min(NumThreads, uint32_t(NumThreadsDefaultBlocksP2));
         assert(NumThreads >= MinThreads &&
                "Expected sufficient inner parallelism.");
         TripCountNumBlocks = ((LoopTripCount - 1) / NumThreads) + 1;
       } else {
         // Not enough parallelism for teams and threads, limit both.
         NumThreads = std::min(NumThreads, MinThreads);
         TripCountNumBlocks = ((LoopTripCount - 1) / NumThreads) + 1;
       }
 
       assert(NumThreads * TripCountNumBlocks >= LoopTripCount &&
              "Expected sufficient parallelism");
       assert(OldNumThreads >= NumThreads &&
              "Number of threads cannot be increased!");
     } else {