[llvm] [Offload] Cache symbols in program (PR #148209)

[Ross Brunton via llvm-commits]

https://github.com/RossBrunton created https://github.com/llvm/llvm-project/pull/148209

When creating a new symbol, check that it already exists. If it does,
return that pointer rather than building a new symbol structure.


>From 19060bebaa0c4ea3d34a2d99b2c99448b043a765 Mon Sep 17 00:00:00 2001
From: Ross Brunton <ross at codeplay.com>
Date: Fri, 11 Jul 2025 11:52:22 +0100
Subject: [PATCH] [Offload] Cache symbols in program

When creating a new symbol, check that it already exists. If it does,
return that pointer rather than building a new symbol structure.
---
 offload/liboffload/src/OffloadImpl.cpp        | 21 ++++++++++++++++---
 .../OffloadAPI/kernel/olGetKernel.cpp         |  8 +++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index fa5d18c044048..4382aad08ede0 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -85,14 +85,16 @@ struct ol_program_impl_t {
   plugin::DeviceImageTy *Image;
   std::unique_ptr<llvm::MemoryBuffer> ImageData;
   std::vector<std::unique_ptr<ol_symbol_impl_t>> Symbols;
+  std::mutex SymbolListMutex;
   __tgt_device_image DeviceImage;
 };
 
 struct ol_symbol_impl_t {
-  ol_symbol_impl_t(GenericKernelTy *Kernel)
-      : PluginImpl(Kernel), Kind(OL_SYMBOL_KIND_KERNEL) {}
+  ol_symbol_impl_t(const char *Name, GenericKernelTy *Kernel)
+      : PluginImpl(Kernel), Kind(OL_SYMBOL_KIND_KERNEL), Name(Name) {}
   std::variant<GenericKernelTy *> PluginImpl;
   ol_symbol_kind_t Kind;
+  const char *Name;
 };
 
 namespace llvm {
@@ -662,6 +664,18 @@ Error olDestroyProgram_impl(ol_program_handle_t Program) {
 
 Error olGetKernel_impl(ol_program_handle_t Program, const char *KernelName,
                        ol_symbol_handle_t *Kernel) {
+  std::lock_guard<std::mutex> Lock{Program->SymbolListMutex};
+
+  // If it already exists, return an existing handle
+  auto Check = std::find_if(Program->Symbols.begin(), Program->Symbols.end(),
+                            [&](auto &Sym) {
+                              return Sym->Kind == OL_SYMBOL_KIND_KERNEL &&
+                                     !std::strcmp(Sym->Name, KernelName);
+                            });
+  if (Check != Program->Symbols.end()) {
+    *Kernel = Check->get();
+    return Error::success();
+  }
 
   auto &Device = Program->Image->getDevice();
   auto KernelImpl = Device.constructKernel(KernelName);
@@ -672,7 +686,8 @@ Error olGetKernel_impl(ol_program_handle_t Program, const char *KernelName,
     return Err;
 
   *Kernel = Program->Symbols
-                .emplace_back(std::make_unique<ol_symbol_impl_t>(&*KernelImpl))
+                .emplace_back(std::make_unique<ol_symbol_impl_t>(
+                    KernelImpl->getName(), &*KernelImpl))
                 .get();
 
   return Error::success();
diff --git a/offload/unittests/OffloadAPI/kernel/olGetKernel.cpp b/offload/unittests/OffloadAPI/kernel/olGetKernel.cpp
index 34870f1fbf0a3..8c222aa851dda 100644
--- a/offload/unittests/OffloadAPI/kernel/olGetKernel.cpp
+++ b/offload/unittests/OffloadAPI/kernel/olGetKernel.cpp
@@ -19,6 +19,14 @@ TEST_P(olGetKernelTest, Success) {
   ASSERT_NE(Kernel, nullptr);
 }
 
+TEST_P(olGetKernelTest, SuccessSamePtr) {
+  ol_symbol_handle_t KernelA = nullptr;
+  ol_symbol_handle_t KernelB = nullptr;
+  ASSERT_SUCCESS(olGetKernel(Program, "foo", &KernelA));
+  ASSERT_SUCCESS(olGetKernel(Program, "foo", &KernelB));
+  ASSERT_EQ(KernelA, KernelB);
+}
+
 TEST_P(olGetKernelTest, InvalidNullProgram) {
   ol_symbol_handle_t Kernel = nullptr;
   ASSERT_ERROR(OL_ERRC_INVALID_NULL_HANDLE,