[Offload] Cache symbols in program

[RossBrunton]

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

[llvmbot]

@llvm/pr-subscribers-offload

Author: Ross Brunton (RossBrunton)

Changes

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

---

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

2 Files Affected:

• (modified) offload/liboffload/src/OffloadImpl.cpp (+24-10)
• (modified) offload/unittests/OffloadAPI/symbol/olGetSymbol.cpp (+18)

diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index af07a6786cfea..40ad1dd7ff617 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -85,16 +85,18 @@ 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(GlobalTy &&Global)
-      : PluginImpl(Global), Kind(OL_SYMBOL_KIND_GLOBAL_VARIABLE) {}
+  ol_symbol_impl_t(const char *Name, GenericKernelTy *Kernel)
+      : PluginImpl(Kernel), Kind(OL_SYMBOL_KIND_KERNEL), Name(Name) {}
+  ol_symbol_impl_t(const char *Name, GlobalTy &&Global)
+      : PluginImpl(Global), Kind(OL_SYMBOL_KIND_GLOBAL_VARIABLE), Name(Name) {}
   std::variant<GenericKernelTy *, GlobalTy> PluginImpl;
   ol_symbol_kind_t Kind;
+  const char *Name;
 };
 
 namespace llvm {
@@ -714,6 +716,18 @@ Error olGetSymbol_impl(ol_program_handle_t Program, const char *Name,
                        ol_symbol_kind_t Kind, ol_symbol_handle_t *Symbol) {
   auto &Device = Program->Image->getDevice();
 
+  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 == Kind && !std::strcmp(Sym->Name, Name);
+      });
+  if (Check != Program->Symbols.end()) {
+    *Symbol = Check->get();
+    return Error::success();
+  }
+
   switch (Kind) {
   case OL_SYMBOL_KIND_KERNEL: {
     auto KernelImpl = Device.constructKernel(Name);
@@ -723,10 +737,10 @@ Error olGetSymbol_impl(ol_program_handle_t Program, const char *Name,
     if (auto Err = KernelImpl->init(Device, *Program->Image))
       return Err;
 
-    *Symbol =
-        Program->Symbols
-            .emplace_back(std::make_unique<ol_symbol_impl_t>(&*KernelImpl))
-            .get();
+    *Symbol = Program->Symbols
+                  .emplace_back(std::make_unique<ol_symbol_impl_t>(
+                      KernelImpl->getName(), &*KernelImpl))
+                  .get();
     return Error::success();
   }
   case OL_SYMBOL_KIND_GLOBAL_VARIABLE: {
@@ -736,8 +750,8 @@ Error olGetSymbol_impl(ol_program_handle_t Program, const char *Name,
       return Res;
 
     *Symbol = Program->Symbols
-                  .emplace_back(
-                      std::make_unique<ol_symbol_impl_t>(std::move(GlobalObj)))
+                  .emplace_back(std::make_unique<ol_symbol_impl_t>(
+                      GlobalObj.getName().c_str(), std::move(GlobalObj)))
                   .get();
 
     return Error::success();
diff --git a/offload/unittests/OffloadAPI/symbol/olGetSymbol.cpp b/offload/unittests/OffloadAPI/symbol/olGetSymbol.cpp
index 5e87ab5b29621..1f496b9c6e1ae 100644
--- a/offload/unittests/OffloadAPI/symbol/olGetSymbol.cpp
+++ b/offload/unittests/OffloadAPI/symbol/olGetSymbol.cpp
@@ -41,6 +41,14 @@ TEST_P(olGetSymbolKernelTest, Success) {
   ASSERT_NE(Kernel, nullptr);
 }
 
+TEST_P(olGetSymbolKernelTest, SuccessSamePtr) {
+  ol_symbol_handle_t KernelA = nullptr;
+  ol_symbol_handle_t KernelB = nullptr;
+  ASSERT_SUCCESS(olGetSymbol(Program, "foo", OL_SYMBOL_KIND_KERNEL, &KernelA));
+  ASSERT_SUCCESS(olGetSymbol(Program, "foo", OL_SYMBOL_KIND_KERNEL, &KernelB));
+  ASSERT_EQ(KernelA, KernelB);
+}
+
 TEST_P(olGetSymbolKernelTest, InvalidNullProgram) {
   ol_symbol_handle_t Kernel = nullptr;
   ASSERT_ERROR(OL_ERRC_INVALID_NULL_HANDLE,
@@ -72,6 +80,16 @@ TEST_P(olGetSymbolGlobalTest, Success) {
   ASSERT_NE(Global, nullptr);
 }
 
+TEST_P(olGetSymbolGlobalTest, SuccessSamePtr) {
+  ol_symbol_handle_t GlobalA = nullptr;
+  ol_symbol_handle_t GlobalB = nullptr;
+  ASSERT_SUCCESS(
+      olGetSymbol(Program, "global", OL_SYMBOL_KIND_GLOBAL_VARIABLE, &GlobalA));
+  ASSERT_SUCCESS(
+      olGetSymbol(Program, "global", OL_SYMBOL_KIND_GLOBAL_VARIABLE, &GlobalB));
+  ASSERT_EQ(GlobalA, GlobalB);
+}
+
 TEST_P(olGetSymbolGlobalTest, InvalidNullProgram) {
   ol_symbol_handle_t Global = nullptr;
   ASSERT_ERROR(

[jhuber6]

Shouldn't we just have contains()? If not, use the llvm::find_if at least.

[RossBrunton]

It's a vector, so sadly no "contains". Updated to use llvm::find_if though (which should be in the C++ standard library, IMO...).

I have updated it to use SmallVector rather than std::vector though.

[jhuber6]

Every symbol lookup is going to search an O(N) list? It would be faster to just use the ELF's symbol table at that point because it's hashed. Shouldn't we be able to use a map or something?