[OpenMP][Offload][AMDGPU] Add envar for setting CU multiplier

This PR added a new envar for configuration of the CU multiplier for cross-team-reduction kernel (SGN: 8).

Now, we can use the following envars for setting CU multipliers on different kernel types:

LIBOMPTARGET_AMDGPU_TEAMS_PER_CU - SPMD
LIBOMPTARGET_AMDGPU_GENERIC_SPMD_TEAMS_PER_CU - Generic SPMD
LIBOMPTARGET_AMDGPU_BIG_JUMP_LOOP_TEAMS_PER_CU - Big-jump-loop
LIBOMPTARGET_AMDGPU_XTEAM_RED_TEAMS_PER_CU - Cross-team-reduction

Author: Kewen12

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

@@ -1095,11 +1095,19 @@ struct AMDGPUKernelTy : public GenericKernelTy {
       // Honor OMP_NUM_TEAMS environment variable for XteamReduction kernel
       // type, if possible.
       int32_t NumTeamsEnvVar = GenericDevice.getOMPNumTeams();
+      // CU mulitiplier from envar.
+      uint32_t EnvarCUMultiplier = GenericDevice.getXTeamRedTeamsPerCU();
+      // Disabled if the value is 0.
+      if (EnvarCUMultiplier == 0) {
+        EnvarCUMultiplier = UINT_MAX;
+      }
 
       if (GenericDevice.isFastReductionEnabled()) {
         // When fast reduction is enabled, the number of teams is capped by
         // the MaxCUMultiplier constant.
-        MaxNumGroups = DeviceNumCUs * llvm::omp::xteam_red::MaxCUMultiplier;
+        MaxNumGroups =
+            DeviceNumCUs * std::min(llvm::omp::xteam_red::MaxCUMultiplier,
+                                    static_cast<int16_t>(EnvarCUMultiplier));
       } else {
         // When fast reduction is not enabled, the number of teams is capped
         // by the metadata that clang CodeGen created. The number of teams
@@ -1110,7 +1118,7 @@ struct AMDGPUKernelTy : public GenericKernelTy {
         // ConstWGSize is the block size that CodeGen used.
         uint32_t CUMultiplier =
             llvm::omp::xteam_red::getXteamRedCUMultiplier(ConstWGSize);
-        MaxNumGroups = DeviceNumCUs * CUMultiplier;
+        MaxNumGroups = DeviceNumCUs * std::min(CUMultiplier, EnvarCUMultiplier);
       }
 
       // If envar OMPX_XTEAMREDUCTION_OCCUPANCY_BASED_OPT is set and no
@@ -2915,6 +2923,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
             "LIBOMPTARGET_AMDGPU_GENERIC_SPMD_TEAMS_PER_CU", 6),
         OMPX_BigJumpLoopTeamsPerCU(
             "LIBOMPTARGET_AMDGPU_BIG_JUMP_LOOP_TEAMS_PER_CU", 0),
+        OMPX_XTeamRedTeamsPerCU("LIBOMPTARGET_AMDGPU_XTEAM_RED_TEAMS_PER_CU",
+                                0),
         OMPX_BigJumpLoopMaxTotalTeams(
             "LIBOMPTARGET_AMDGPU_BIG_JUMP_LOOP_MAX_TOTAL_TEAMS", 1024 * 1024),
         OMPX_LowTripCount("LIBOMPTARGET_AMDGPU_LOW_TRIPCOUNT", 9000),
@@ -2980,6 +2990,9 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   virtual uint32_t getOMPXBigJumpLoopTeamsPerCU() const override {
     return OMPX_BigJumpLoopTeamsPerCU;
   }
+  virtual uint32_t getXTeamRedTeamsPerCU() const override {
+    return OMPX_XTeamRedTeamsPerCU;
+  }
   virtual uint32_t getOMPXBigJumpLoopMaxTotalTeams() const override {
     return OMPX_BigJumpLoopMaxTotalTeams;
   }
@@ -4427,6 +4440,12 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   /// OMPX_BigJumpLoopTeamsPerCU * #CUs.
   UInt32Envar OMPX_BigJumpLoopTeamsPerCU;
 
+  /// Envar for controlling the number of teams relative to the number of
+  /// compute units (CUs) for cross-team-reduction kernels. 0 indicates that
+  /// this value is not specified. If non-zero, the number of teams =
+  /// OMPX_XTeamRedTeamsPerCU * #CUs.
+  UInt32Envar OMPX_XTeamRedTeamsPerCU;
+
   /// Envar controlling the maximum number of teams per device for
   /// Big-Jump-Loop kernels.
   UInt32Envar OMPX_BigJumpLoopMaxTotalTeams;

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

@@ -1017,6 +1017,9 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
   virtual uint32_t getOMPXBigJumpLoopTeamsPerCU() const {
     llvm_unreachable("Unimplemented");
   }
+  virtual uint32_t getXTeamRedTeamsPerCU() const {
+    llvm_unreachable("Unimplemented");
+  }
   virtual uint32_t getOMPXBigJumpLoopMaxTotalTeams() const {
     llvm_unreachable("Unimplemented");
   }