[Offload] Cache symbols in program

offload/liboffload/src/OffloadImpl.cpp

@@ -84,17 +84,20 @@ struct ol_program_impl_t {
         DeviceImage(DeviceImage) {}
   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;
+  llvm::StringMap<std::unique_ptr<ol_symbol_impl_t>> KernelSymbols;
+  llvm::StringMap<std::unique_ptr<ol_symbol_impl_t>> GlobalSymbols;
 };
 
 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;
+  llvm::StringRef Name;
 };
 
 namespace llvm {
@@ -714,30 +717,53 @@ 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 CheckCache = [&](StringMap<std::unique_ptr<ol_symbol_impl_t>> &Map)
+      -> std::optional<ol_symbol_handle_t> {
+    if (Map.contains(Name))
+      return Map[Name].get();
+    return std::nullopt;
+  };
+
   switch (Kind) {
   case OL_SYMBOL_KIND_KERNEL: {
+    if (auto Cache = CheckCache(Program->KernelSymbols)) {
+      *Symbol = *Cache;
+      return Plugin::success();
+    }
+
     auto KernelImpl = Device.constructKernel(Name);
     if (!KernelImpl)
       return KernelImpl.takeError();
 
     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->KernelSymbols
+                  .insert({Name, std::make_unique<ol_symbol_impl_t>(
+                                     KernelImpl->getName(), &*KernelImpl)})
+                  .first->getValue()
+                  .get();
     return Error::success();
   }
   case OL_SYMBOL_KIND_GLOBAL_VARIABLE: {
+    if (auto Cache = CheckCache(Program->GlobalSymbols)) {
+      *Symbol = *Cache;
+      return Plugin::success();
+    }
+
     GlobalTy GlobalObj{Name};
     if (auto Res = Device.Plugin.getGlobalHandler().getGlobalMetadataFromDevice(
             Device, *Program->Image, GlobalObj))
       return Res;
 
-    *Symbol = Program->Symbols
-                  .emplace_back(
-                      std::make_unique<ol_symbol_impl_t>(std::move(GlobalObj)))
+    *Symbol = Program->GlobalSymbols
+                  .insert({Name, std::make_unique<ol_symbol_impl_t>(
+                                     GlobalObj.getName().c_str(),
+                                     std::move(GlobalObj))})
+                  .first->getValue()
                   .get();
 
     return Error::success();

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(