[Offload] Add olGetGlobalVariable

[RossBrunton]

Looks up a global variable as a symbol

[llvmbot]

@llvm/pr-subscribers-offload

Author: Ross Brunton (RossBrunton)

Changes

Looks up a global variable as a symbol

---

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

6 Files Affected:

• (modified) offload/liboffload/API/Symbol.td (+15)
• (modified) offload/liboffload/src/OffloadImpl.cpp (+21-1)
• (modified) offload/unittests/OffloadAPI/CMakeLists.txt (+3)
• (modified) offload/unittests/OffloadAPI/device_code/global.c (+1)
• (modified) offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp (+8)
• (added) offload/unittests/OffloadAPI/symbol/olGetGlobalVariable.cpp (+56)

diff --git a/offload/liboffload/API/Symbol.td b/offload/liboffload/API/Symbol.td
index cf4d45b09f035..68ccbc4e68c4b 100644
--- a/offload/liboffload/API/Symbol.td
+++ b/offload/liboffload/API/Symbol.td
@@ -15,5 +15,20 @@ def : Enum {
   let desc = "The kind of a symbol";
   let etors =[
     Etor<"KERNEL", "a kernel object">,
+    Etor<"GLOBAL_VARIABLE", "a global variable">,
   ];
 }
+
+def : Function {
+    let name = "olGetGlobalVariable";
+    let desc = "Get a global variable named `GlobalName` in the given program.";
+    let details = [
+        "Symbol handles are owned by the program, so do not need to be destroyed."
+    ];
+    let params = [
+        Param<"ol_program_handle_t", "Program", "handle of the program", PARAM_IN>,
+        Param<"const char*", "GlobalName", "name of the global variable in the program", PARAM_IN>,
+        Param<"ol_symbol_handle_t*", "Global", "output pointer for the fetched global", PARAM_OUT>
+    ];
+    let returns = [];
+}
diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index fa5d18c044048..b4a545e013d6f 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -91,7 +91,9 @@ struct ol_program_impl_t {
 struct ol_symbol_impl_t {
   ol_symbol_impl_t(GenericKernelTy *Kernel)
       : PluginImpl(Kernel), Kind(OL_SYMBOL_KIND_KERNEL) {}
-  std::variant<GenericKernelTy *> PluginImpl;
+  ol_symbol_impl_t(GlobalTy &&Global)
+      : PluginImpl(Global), Kind(OL_SYMBOL_KIND_GLOBAL_VARIABLE) {}
+  std::variant<GenericKernelTy *, GlobalTy> PluginImpl;
   ol_symbol_kind_t Kind;
 };
 
@@ -726,5 +728,23 @@ Error olLaunchKernel_impl(ol_queue_handle_t Queue, ol_device_handle_t Device,
   return Error::success();
 }
 
+Error olGetGlobalVariable_impl(ol_program_handle_t Program,
+                               const char *GlobalName,
+                               ol_symbol_handle_t *Global) {
+  auto &Device = Program->Image->getDevice();
+
+  GlobalTy GlobalObj{GlobalName};
+  if (auto Res = Device.Plugin.getGlobalHandler().getGlobalMetadataFromDevice(
+          Device, *Program->Image, GlobalObj))
+    return Res;
+
+  *Global = Program->Symbols
+                .emplace_back(
+                    std::make_unique<ol_symbol_impl_t>(std::move(GlobalObj)))
+                .get();
+
+  return Error::success();
+}
+
 } // namespace offload
 } // namespace llvm
diff --git a/offload/unittests/OffloadAPI/CMakeLists.txt b/offload/unittests/OffloadAPI/CMakeLists.txt
index 93e5fd2f6cd26..5d734df635cb4 100644
--- a/offload/unittests/OffloadAPI/CMakeLists.txt
+++ b/offload/unittests/OffloadAPI/CMakeLists.txt
@@ -41,3 +41,6 @@ add_offload_unittest("queue"
     queue/olDestroyQueue.cpp
     queue/olGetQueueInfo.cpp
     queue/olGetQueueInfoSize.cpp)
+
+add_offload_unittest("symbol"
+    symbol/olGetGlobalVariable.cpp)
diff --git a/offload/unittests/OffloadAPI/device_code/global.c b/offload/unittests/OffloadAPI/device_code/global.c
index b30e406fb98c7..9f27f9424324f 100644
--- a/offload/unittests/OffloadAPI/device_code/global.c
+++ b/offload/unittests/OffloadAPI/device_code/global.c
@@ -1,6 +1,7 @@
 #include <gpuintrin.h>
 #include <stdint.h>
 
+[[gnu::visibility("default")]]
 uint32_t global[64];
 
 __gpu_kernel void write() {
diff --git a/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp b/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp
index acda4795edec2..cb77acf1bd21a 100644
--- a/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp
+++ b/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp
@@ -223,6 +223,14 @@ TEST_P(olLaunchKernelGlobalTest, Success) {
   ASSERT_SUCCESS(olMemFree(Mem));
 }
 
+TEST_P(olLaunchKernelGlobalTest, InvalidNotAKernel) {
+  ol_symbol_handle_t Global = nullptr;
+  ASSERT_SUCCESS(olGetGlobalVariable(Program, "global", &Global));
+  ASSERT_ERROR(
+      OL_ERRC_SYMBOL_KIND,
+      olLaunchKernel(Queue, Device, Global, nullptr, 0, &LaunchArgs, nullptr));
+}
+
 TEST_P(olLaunchKernelGlobalCtorTest, Success) {
   void *Mem;
   ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED,
diff --git a/offload/unittests/OffloadAPI/symbol/olGetGlobalVariable.cpp b/offload/unittests/OffloadAPI/symbol/olGetGlobalVariable.cpp
new file mode 100644
index 0000000000000..e5600f399890b
--- /dev/null
+++ b/offload/unittests/OffloadAPI/symbol/olGetGlobalVariable.cpp
@@ -0,0 +1,56 @@
+//===------- Offload API tests - olGetGlobalVariable ----------------------===//
+//
+// 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>
+
+struct olGetGlobalVariableTest : OffloadQueueTest {
+  void SetUp() override {
+    RETURN_ON_FATAL_FAILURE(OffloadQueueTest::SetUp());
+    ASSERT_TRUE(TestEnvironment::loadDeviceBinary("global", Device, DeviceBin));
+    ASSERT_GE(DeviceBin->getBufferSize(), 0lu);
+    ASSERT_SUCCESS(olCreateProgram(Device, DeviceBin->getBufferStart(),
+                                   DeviceBin->getBufferSize(), &Program));
+  }
+
+  void TearDown() override {
+    if (Program) {
+      olDestroyProgram(Program);
+    }
+    RETURN_ON_FATAL_FAILURE(OffloadQueueTest::TearDown());
+  }
+
+  std::unique_ptr<llvm::MemoryBuffer> DeviceBin;
+  ol_program_handle_t Program = nullptr;
+  ol_kernel_launch_size_args_t LaunchArgs{};
+};
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olGetGlobalVariableTest);
+
+TEST_P(olGetGlobalVariableTest, Success) {
+  ol_symbol_handle_t Global = nullptr;
+  ASSERT_SUCCESS(olGetGlobalVariable(Program, "global", &Global));
+  ASSERT_NE(Global, nullptr);
+}
+
+TEST_P(olGetGlobalVariableTest, InvalidNullProgram) {
+  ol_symbol_handle_t Global = nullptr;
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_HANDLE,
+               olGetGlobalVariable(nullptr, "global", &Global));
+}
+
+TEST_P(olGetGlobalVariableTest, InvalidNullGlobalPointer) {
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER,
+               olGetGlobalVariable(Program, "global", nullptr));
+}
+
+TEST_P(olGetGlobalVariableTest, InvalidGlobalName) {
+  ol_symbol_handle_t Global = nullptr;
+  ASSERT_ERROR(OL_ERRC_NOT_FOUND,
+               olGetGlobalVariable(Program, "invalid_global", &Global));
+}

[jhuber6]

Debating whether or not we should just have olGetSymbol and let the user assume whether or not it's a kernel and let them use the info to verify if they're unsure.

[RossBrunton]

@jhuber6 I'm not sure that's possible generically. The AMD backend at least needs to mangle kernel names to add ".kd" on the end, so we'd need to know ahead of time whether it's a kernel or not.

However, olGetSymbol could work if we pass in the symbol kind as a parameter, then it dispatches to either kernel or global lookup depending on that.

[jhuber6]

@jhuber6 I'm not sure that's possible generically. The AMD backend at least needs to mangle kernel names to add ".kd" on the end, so we'd need to know ahead of time whether it's a kernel or not.

However, olGetSymbol could work if we pass in the symbol kind as a parameter, then it dispatches to either kernel or global lookup depending on that.

That's true, it's a little confusing because the AMD version implicity requires the .kd for its kernels and we can't just slice it, hm.

[RossBrunton]

@jhuber6 To follow up on the discussion at #147972 (review) , I don't think it's worth doing the type lookup. The caller presumably knows whether it's requesting a global variable or a kernel, so poking around and querying the binary just feels like wasted cycles.

There's also the chance that a binary could have both a global variable and a kernel with the same name, although I'm not sure whether in practice that's a thing to be concerned about.

Shall I replace olGetGlobalVariable(Program, Name, Out) and olGetKernel(Program, Name, Out) with a single olGetSymbol(Program, Name, Kind, Out) API, or does having to pass Kind make it not worth it?

Edit: Decided to just do it. #148221