[llvm] [Offload] Add olGetKernelMaxGroupSize (PR #142950)

Ross Brunton via llvm-commits llvm-commits at lists.llvm.org
Tue Jul 15 07:43:53 PDT 2025


https://github.com/RossBrunton updated https://github.com/llvm/llvm-project/pull/142950

>From 37d0791b50f45523fb54fc6959bdeb213e3554f2 Mon Sep 17 00:00:00 2001
From: Ross Brunton <ross at codeplay.com>
Date: Thu, 5 Jun 2025 12:21:52 +0100
Subject: [PATCH 1/2] [Offload] Add olKernelMaxGroupSize

This is equivalent to `cuOccupancyMaxPotentialBlockSize`. It is currently
only implented on Cuda; AMDGPU and Host return the legal-but-suboptimal
value of `1`.

Co-Authored-By: Callum Fare <callum at codeplay.com>
---
 offload/liboffload/API/Kernel.td              | 17 +++++++-
 offload/liboffload/src/OffloadImpl.cpp        | 17 ++++++++
 offload/plugins-nextgen/amdgpu/src/rtl.cpp    |  8 ++++
 .../common/include/PluginInterface.h          |  3 ++
 .../cuda/dynamic_cuda/cuda.cpp                |  1 +
 .../plugins-nextgen/cuda/dynamic_cuda/cuda.h  |  3 ++
 offload/plugins-nextgen/cuda/src/rtl.cpp      | 14 ++++++
 offload/plugins-nextgen/host/src/rtl.cpp      |  7 +++
 offload/unittests/OffloadAPI/CMakeLists.txt   |  1 +
 .../kernel/olGetKernelMaxGroupSize.cpp        | 43 +++++++++++++++++++
 10 files changed, 113 insertions(+), 1 deletion(-)
 create mode 100644 offload/unittests/OffloadAPI/kernel/olGetKernelMaxGroupSize.cpp

diff --git a/offload/liboffload/API/Kernel.td b/offload/liboffload/API/Kernel.td
index 1e9537452820d..a229685b2dde3 100644
--- a/offload/liboffload/API/Kernel.td
+++ b/offload/liboffload/API/Kernel.td
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file contains Offload API definitions related to launching kernels
+// This file contains Offload API definitions related to kernels
 //
 //===----------------------------------------------------------------------===//
 
@@ -44,3 +44,18 @@ def : Function {
         Return<"OL_ERRC_SYMBOL_KIND", ["The provided symbol is not a kernel"]>,
     ];
 }
+
+def : Function {
+    let name = "olGetKernelMaxGroupSize";
+    let desc = "Get the maximum block size needed to achieve maximum occupancy.";
+    let details = [];
+    let params = [
+        Param<"ol_device_handle_t", "Device", "device intended to run the kernel", PARAM_IN>,
+        Param<"ol_symbol_handle_t", "Kernel", "handle of the kernel", PARAM_IN>,
+        Param<"size_t", "SharedMemory", "dynamic shared memory required", PARAM_IN>,
+        Param<"size_t*", "GroupSize", "maximum block size", PARAM_OUT>
+    ];
+    let returns = [
+        Return<"OL_ERRC_SYMBOL_KIND", ["The provided symbol is not a kernel"]>,
+    ];
+}
diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index c4e7f9689a900..9cd36a5f8ffad 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -662,6 +662,23 @@ Error olDestroyProgram_impl(ol_program_handle_t Program) {
   return olDestroy(Program);
 }
 
+Error olGetKernelMaxGroupSize_impl(ol_device_handle_t Device,
+                                   ol_symbol_handle_t Kernel,
+                                   size_t DynamicMemSize, size_t *GroupSize) {
+  if (Kernel->Kind != OL_SYMBOL_KIND_KERNEL)
+    return createOffloadError(ErrorCode::SYMBOL_KIND,
+                              "provided symbol is not a kernel");
+  auto *KernelImpl = std::get<GenericKernelTy *>(Kernel->PluginImpl);
+
+  auto Res = KernelImpl->maxGroupSize(*Device->Device, DynamicMemSize);
+  if (auto Err = Res.takeError())
+    return Err;
+
+  *GroupSize = *Res;
+
+  return Error::success();
+}
+
 Error olLaunchKernel_impl(ol_queue_handle_t Queue, ol_device_handle_t Device,
                           ol_symbol_handle_t Kernel, const void *ArgumentsData,
                           size_t ArgumentsSize,
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 12c7cc62905c9..dd4106f61c2c6 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -570,6 +570,14 @@ struct AMDGPUKernelTy : public GenericKernelTy {
                    KernelLaunchParamsTy LaunchParams,
                    AsyncInfoWrapperTy &AsyncInfoWrapper) const override;
 
+  /// Return maximum block size for maximum occupancy
+  ///
+  /// TODO: This needs to be implemented for amdgpu
+  Expected<size_t> maxGroupSize(GenericDeviceTy &GenericDevice,
+                                size_t DynamicMemSize) const override {
+    return 1;
+  }
+
   /// Print more elaborate kernel launch info for AMDGPU
   Error printLaunchInfoDetails(GenericDeviceTy &GenericDevice,
                                KernelArgsTy &KernelArgs, uint32_t NumThreads[3],
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index 162b149ab483e..6fe18677b5eae 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -316,6 +316,9 @@ struct GenericKernelTy {
                            KernelLaunchParamsTy LaunchParams,
                            AsyncInfoWrapperTy &AsyncInfoWrapper) const = 0;
 
+  virtual Expected<size_t> maxGroupSize(GenericDeviceTy &GenericDevice,
+                                        size_t DynamicMemSize) const = 0;
+
   /// Get the kernel name.
   const char *getName() const { return Name.c_str(); }
 
diff --git a/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.cpp b/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.cpp
index 361a781e8f9b6..c003d0b2f9451 100644
--- a/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.cpp
+++ b/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.cpp
@@ -72,6 +72,7 @@ DLWRAP(cuDevicePrimaryCtxGetState, 3)
 DLWRAP(cuDevicePrimaryCtxSetFlags, 2)
 DLWRAP(cuDevicePrimaryCtxRetain, 2)
 DLWRAP(cuModuleLoadDataEx, 5)
+DLWRAP(cuOccupancyMaxPotentialBlockSize, 6)
 
 DLWRAP(cuDeviceCanAccessPeer, 3)
 DLWRAP(cuCtxEnablePeerAccess, 2)
diff --git a/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.h b/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.h
index b6c022c8e7e8b..5f1c44364c143 100644
--- a/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.h
+++ b/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.h
@@ -290,6 +290,7 @@ static inline void *CU_LAUNCH_PARAM_BUFFER_POINTER = (void *)0x01;
 static inline void *CU_LAUNCH_PARAM_BUFFER_SIZE = (void *)0x02;
 
 typedef void (*CUstreamCallback)(CUstream, CUresult, void *);
+typedef size_t (*CUoccupancyB2DSize)(int);
 
 CUresult cuCtxGetDevice(CUdevice *);
 CUresult cuDeviceGet(CUdevice *, int);
@@ -372,5 +373,7 @@ CUresult cuMemSetAccess(CUdeviceptr ptr, size_t size,
 CUresult cuMemGetAllocationGranularity(size_t *granularity,
                                        const CUmemAllocationProp *prop,
                                        CUmemAllocationGranularity_flags option);
+CUresult cuOccupancyMaxPotentialBlockSize(int *, int *, CUfunction,
+                                          CUoccupancyB2DSize, size_t, int);
 
 #endif
diff --git a/offload/plugins-nextgen/cuda/src/rtl.cpp b/offload/plugins-nextgen/cuda/src/rtl.cpp
index b787376eb1770..4e178824e5530 100644
--- a/offload/plugins-nextgen/cuda/src/rtl.cpp
+++ b/offload/plugins-nextgen/cuda/src/rtl.cpp
@@ -157,6 +157,20 @@ struct CUDAKernelTy : public GenericKernelTy {
                    KernelLaunchParamsTy LaunchParams,
                    AsyncInfoWrapperTy &AsyncInfoWrapper) const override;
 
+  /// Return maximum block size for maximum occupancy
+  Expected<size_t> maxGroupSize(GenericDeviceTy &,
+                                size_t DynamicMemSize) const override {
+    int minGridSize;
+    int maxBlockSize;
+    auto Res = cuOccupancyMaxPotentialBlockSize(
+        &minGridSize, &maxBlockSize, Func, NULL, DynamicMemSize, INT_MAX);
+    if (auto Err = Plugin::check(
+            Res, "error in cuOccupancyMaxPotentialBlockSize: %s")) {
+      return Err;
+    }
+    return maxBlockSize;
+  }
+
 private:
   /// The CUDA kernel function to execute.
   CUfunction Func;
diff --git a/offload/plugins-nextgen/host/src/rtl.cpp b/offload/plugins-nextgen/host/src/rtl.cpp
index d950572265b4c..d5e7e991d7e07 100644
--- a/offload/plugins-nextgen/host/src/rtl.cpp
+++ b/offload/plugins-nextgen/host/src/rtl.cpp
@@ -114,6 +114,13 @@ struct GenELF64KernelTy : public GenericKernelTy {
     return Plugin::success();
   }
 
+  /// Return maximum block size for maximum occupancy
+  Expected<size_t> maxGroupSize(GenericDeviceTy &Device,
+                                size_t DynamicMemSize) const override {
+    // TODO
+    return 1;
+  }
+
 private:
   /// The kernel function to execute.
   void (*Func)(void);
diff --git a/offload/unittests/OffloadAPI/CMakeLists.txt b/offload/unittests/OffloadAPI/CMakeLists.txt
index d76338612210d..dee3b196f85be 100644
--- a/offload/unittests/OffloadAPI/CMakeLists.txt
+++ b/offload/unittests/OffloadAPI/CMakeLists.txt
@@ -19,6 +19,7 @@ add_offload_unittest("init"
 target_compile_definitions("init.unittests" PRIVATE DISABLE_WRAPPER)
 
 add_offload_unittest("kernel"
+    kernel/olGetKernelMaxGroupSize.cpp
     kernel/olLaunchKernel.cpp)
 
 add_offload_unittest("memory"
diff --git a/offload/unittests/OffloadAPI/kernel/olGetKernelMaxGroupSize.cpp b/offload/unittests/OffloadAPI/kernel/olGetKernelMaxGroupSize.cpp
new file mode 100644
index 0000000000000..7923b10ea3030
--- /dev/null
+++ b/offload/unittests/OffloadAPI/kernel/olGetKernelMaxGroupSize.cpp
@@ -0,0 +1,43 @@
+//===------- Offload API tests - olGetKernelMaxGroupSize ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "../common/Fixtures.hpp"
+#include <OffloadAPI.h>
+#include <gtest/gtest.h>
+
+using olKernelGetMaxGroupSizeTest = OffloadKernelTest;
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olKernelGetMaxGroupSizeTest);
+
+TEST_P(olKernelGetMaxGroupSizeTest, Success) {
+  size_t Size{0};
+  ASSERT_SUCCESS(olGetKernelMaxGroupSize(Device, Kernel, 0, &Size));
+  ASSERT_GT(Size, 0u);
+}
+
+TEST_P(olKernelGetMaxGroupSizeTest, SuccessMem) {
+  size_t Size{0};
+  ASSERT_SUCCESS(olGetKernelMaxGroupSize(Device, Kernel, 1024, &Size));
+  ASSERT_GT(Size, 0u);
+}
+
+TEST_P(olKernelGetMaxGroupSizeTest, NullKernel) {
+  size_t Size;
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_HANDLE,
+               olGetKernelMaxGroupSize(Device, nullptr, 0, &Size));
+}
+
+TEST_P(olKernelGetMaxGroupSizeTest, NullDevice) {
+  size_t Size;
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_HANDLE,
+               olGetKernelMaxGroupSize(nullptr, Kernel, 0, &Size));
+}
+
+TEST_P(olKernelGetMaxGroupSizeTest, NullOutput) {
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER,
+               olGetKernelMaxGroupSize(Device, Kernel, 0, nullptr));
+}

>From 4bdf74896f72bbf386ea5489ae3d8400ea965079 Mon Sep 17 00:00:00 2001
From: Ross Brunton <ross at codeplay.com>
Date: Mon, 30 Jun 2025 16:22:33 +0100
Subject: [PATCH 2/2] Use uint64_t rather than size_t

---
 offload/plugins-nextgen/amdgpu/src/rtl.cpp               | 4 ++--
 offload/plugins-nextgen/common/include/PluginInterface.h | 4 ++--
 offload/plugins-nextgen/cuda/src/rtl.cpp                 | 4 ++--
 offload/plugins-nextgen/host/src/rtl.cpp                 | 4 ++--
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index dd4106f61c2c6..6cf9d260ee2ff 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -573,8 +573,8 @@ struct AMDGPUKernelTy : public GenericKernelTy {
   /// Return maximum block size for maximum occupancy
   ///
   /// TODO: This needs to be implemented for amdgpu
-  Expected<size_t> maxGroupSize(GenericDeviceTy &GenericDevice,
-                                size_t DynamicMemSize) const override {
+  Expected<uint64_t> maxGroupSize(GenericDeviceTy &GenericDevice,
+                                  uint64_t DynamicMemSize) const override {
     return 1;
   }
 
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index 6fe18677b5eae..d3f2c117838b9 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -316,8 +316,8 @@ struct GenericKernelTy {
                            KernelLaunchParamsTy LaunchParams,
                            AsyncInfoWrapperTy &AsyncInfoWrapper) const = 0;
 
-  virtual Expected<size_t> maxGroupSize(GenericDeviceTy &GenericDevice,
-                                        size_t DynamicMemSize) const = 0;
+  virtual Expected<uint64_t> maxGroupSize(GenericDeviceTy &GenericDevice,
+                                          uint64_t DynamicMemSize) const = 0;
 
   /// Get the kernel name.
   const char *getName() const { return Name.c_str(); }
diff --git a/offload/plugins-nextgen/cuda/src/rtl.cpp b/offload/plugins-nextgen/cuda/src/rtl.cpp
index 4e178824e5530..f7fc33e7a6656 100644
--- a/offload/plugins-nextgen/cuda/src/rtl.cpp
+++ b/offload/plugins-nextgen/cuda/src/rtl.cpp
@@ -158,8 +158,8 @@ struct CUDAKernelTy : public GenericKernelTy {
                    AsyncInfoWrapperTy &AsyncInfoWrapper) const override;
 
   /// Return maximum block size for maximum occupancy
-  Expected<size_t> maxGroupSize(GenericDeviceTy &,
-                                size_t DynamicMemSize) const override {
+  Expected<uint64_t> maxGroupSize(GenericDeviceTy &,
+                                  uint64_t DynamicMemSize) const override {
     int minGridSize;
     int maxBlockSize;
     auto Res = cuOccupancyMaxPotentialBlockSize(
diff --git a/offload/plugins-nextgen/host/src/rtl.cpp b/offload/plugins-nextgen/host/src/rtl.cpp
index d5e7e991d7e07..92b67d524430c 100644
--- a/offload/plugins-nextgen/host/src/rtl.cpp
+++ b/offload/plugins-nextgen/host/src/rtl.cpp
@@ -115,8 +115,8 @@ struct GenELF64KernelTy : public GenericKernelTy {
   }
 
   /// Return maximum block size for maximum occupancy
-  Expected<size_t> maxGroupSize(GenericDeviceTy &Device,
-                                size_t DynamicMemSize) const override {
+  Expected<uint64_t> maxGroupSize(GenericDeviceTy &Device,
+                                  uint64_t DynamicMemSize) const override {
     // TODO
     return 1;
   }



More information about the llvm-commits mailing list