[llvm] [OFFLOAD][L0] Handle group sizes correctly for multidimensional bare kernels (PR #191770)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-offload

Author: Alex Duran (adurang)

<details>
<summary>Changes</summary>

Don't use the L0 heuristics if all the dimensions are specified by the user code.

---
Full diff: https://github.com/llvm/llvm-project/pull/191770.diff


2 Files Affected:

- (modified) offload/plugins-nextgen/level_zero/src/L0Kernel.cpp (+29-24) 
- (modified) offload/test/offloading/ompx_bare_multi_dim.cpp (-1) 


``````````diff
diff --git a/offload/plugins-nextgen/level_zero/src/L0Kernel.cpp b/offload/plugins-nextgen/level_zero/src/L0Kernel.cpp
index e30e55bba8f1b..b87af310e192d 100644
--- a/offload/plugins-nextgen/level_zero/src/L0Kernel.cpp
+++ b/offload/plugins-nextgen/level_zero/src/L0Kernel.cpp
@@ -368,35 +368,40 @@ Error L0KernelTy::setKernelGroups(L0DeviceTy &l0Device, L0LaunchEnvTy &KEnv,
                                   uint32_t NumThreads[3],
                                   uint32_t NumBlocks[3]) const {
 
-  if (KernelEnvironment.Configuration.ExecMode != OMP_TGT_EXEC_MODE_BARE) {
-    // For non-bare mode, the groups are already set in the launch.
-    KEnv.GroupCounts = {NumBlocks[0], NumBlocks[1], NumBlocks[2]};
-    CALL_ZE_RET_ERROR(zeKernelSetGroupSize, getZeKernel(), NumThreads[0],
-                      NumThreads[1], NumThreads[2]);
-    return Plugin::success();
-  }
-
-  int32_t NumTeams = NumBlocks[0];
-  int32_t ThreadLimit = NumThreads[0];
-  if (NumTeams < 0)
-    NumTeams = 0;
-  if (ThreadLimit < 0)
-    ThreadLimit = 0;
+  bool userDefinedGroups = NumThreads[0] != 0 && NumThreads[1] != 0 &&
+                           NumThreads[2] != 0 && NumBlocks[0] != 0 &&
+                           NumBlocks[1] != 0 && NumBlocks[2] != 0;
 
   uint32_t GroupSizes[3];
-  auto DeviceId = l0Device.getDeviceId();
-  auto &KernelPR = KEnv.KernelPR;
-  // Check if we can reuse previous group parameters.
-  bool GroupParamsReused =
-      KernelPR.reuseGroupParams(NumTeams, ThreadLimit, GroupSizes, KEnv);
+  bool GroupParamsReused = false;
 
-  if (!GroupParamsReused) {
-    if (auto Err =
-            getGroupsShape(l0Device, NumTeams, ThreadLimit, GroupSizes, KEnv))
-      return Err;
-    KernelPR.cacheGroupParams(NumTeams, ThreadLimit, GroupSizes, KEnv);
+  if (userDefinedGroups) {
+    KEnv.GroupCounts = {NumBlocks[0], NumBlocks[1], NumBlocks[2]};
+    GroupSizes[0] = NumThreads[0];
+    GroupSizes[1] = NumThreads[1];
+    GroupSizes[2] = NumThreads[2];
+  } else {
+    int32_t NumTeams = NumBlocks[0];
+    int32_t ThreadLimit = NumThreads[0];
+    if (NumTeams < 0)
+      NumTeams = 0;
+    if (ThreadLimit < 0)
+      ThreadLimit = 0;
+
+    auto &KernelPR = KEnv.KernelPR;
+    // Check if we can reuse previous group parameters.
+    GroupParamsReused =
+        KernelPR.reuseGroupParams(NumTeams, ThreadLimit, GroupSizes, KEnv);
+
+    if (!GroupParamsReused) {
+      if (auto Err =
+              getGroupsShape(l0Device, NumTeams, ThreadLimit, GroupSizes, KEnv))
+        return Err;
+      KernelPR.cacheGroupParams(NumTeams, ThreadLimit, GroupSizes, KEnv);
+    }
   }
 
+  auto DeviceId = l0Device.getDeviceId();
   INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
        "Team sizes = {%" PRIu32 ", %" PRIu32 ", %" PRIu32 "}\n", GroupSizes[0],
        GroupSizes[1], GroupSizes[2]);
diff --git a/offload/test/offloading/ompx_bare_multi_dim.cpp b/offload/test/offloading/ompx_bare_multi_dim.cpp
index 46d5f03b44692..aa854114f9914 100644
--- a/offload/test/offloading/ompx_bare_multi_dim.cpp
+++ b/offload/test/offloading/ompx_bare_multi_dim.cpp
@@ -2,7 +2,6 @@
 // RUN: %libomptarget-compilexx-generic
 // RUN: env LIBOMPTARGET_INFO=63 %libomptarget-run-generic 2>&1 | %fcheck-generic
 // REQUIRES: gpu
-// XFAIL: intelgpu
 // clang-format on
 
 #include <ompx.h>

``````````

</details>


https://github.com/llvm/llvm-project/pull/191770