[llvm] [offload][OpenMP] Set StaticBlockMemSize in kernel launch info (PR #228756)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-offload

Author: Robert Imschweiler (ro-i)

<details>
<summary>Changes</summary>

#<!-- -->223761 moved the dynamic block memory preparation to libomptarget, which reads the kernel's static block memory size from KernelLaunchInfoTy, but nothing set that field. -> It's 0 and libomptarget does not apply the dyn_groupprivate fallback when the dynamic size alone fits but static + dynamic memory exceeds the maximum, and the plugin then rejects the launch.

Claude assisted with this patch.

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


2 Files Affected:

- (modified) offload/libompaccsupport/PluginManager.cpp (+1) 
- (modified) offload/test/offloading/dyn_groupprivate.cpp (+18) 


``````````diff
diff --git a/offload/libompaccsupport/PluginManager.cpp b/offload/libompaccsupport/PluginManager.cpp
index 19207685385d73..921d999997175b 100644
--- a/offload/libompaccsupport/PluginManager.cpp
+++ b/offload/libompaccsupport/PluginManager.cpp
@@ -511,6 +511,7 @@ static int loadImagesOntoDevice(DeviceTy &Device) {
             LaunchInfo.Mode =
                 static_cast<llvm::omp::OMPTgtExecModeFlags>(Cfg.ExecMode);
             LaunchInfo.ReductionDataSize = Cfg.ReductionDataSize;
+            LaunchInfo.StaticBlockMemSize = Kernel->getStaticBlockMemSize();
             // Max = Config.Max > 0 ? min(Config.Max, Device.Max) : Device.Max,
             // further clamped to the kernel function's own driver-reported
             // maximum.
diff --git a/offload/test/offloading/dyn_groupprivate.cpp b/offload/test/offloading/dyn_groupprivate.cpp
index 5a478a0c5a9171..ee2392740f8639 100644
--- a/offload/test/offloading/dyn_groupprivate.cpp
+++ b/offload/test/offloading/dyn_groupprivate.cpp
@@ -109,6 +109,24 @@ int main() {
       ++Failed;
   }
 
+  // The dynamic size alone fits, but not together with the kernel's static
+  // block memory, e.g., the device runtime's team state.
+  size_t AlmostMaxSize = MaxSize - 1;
+
+// Verify that the fallback(default_mem) also triggers in that case.
+#pragma omp target dyn_groupprivate(fallback(default_mem) : AlmostMaxSize)     \
+    map(tofrom : Failed)
+  {
+    if (!omp_get_dyn_gprivate_ptr(0))
+      ++Failed;
+    if (omp_get_dyn_gprivate_nofb_ptr(0))
+      ++Failed;
+    if (omp_get_dyn_gprivate_size() != AlmostMaxSize)
+      ++Failed;
+    if (omp_get_dyn_gprivate_memspace() != omp_default_mem_space)
+      ++Failed;
+  }
+
 // Verify that the fallback(abort) modifier works.
 #pragma omp target dyn_groupprivate(fallback(abort) : N) map(tofrom : Failed)
   {

``````````

</details>


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