[Openmp-commits] [openmp] [Libomptarget] Remove bitcode image map used for JIT processing (PR #75672)

[Joseph Huber via Openmp-commits]

https://github.com/jhuber6 created https://github.com/llvm/llvm-project/pull/75672

Summary:
Libomptarget supports JIT by treating an LLVM-IR file as a regular input
image. The handling here used a global map to keep track of triples once
it was parsed. This was done to same time, however this created a global
constructor as well as an extra mutex to handle it. This patch removes
the use of this map.

Instead, we simply use the file magic to perform a quick check if the
input image is valid bitcode. If not, we then create a lazy module. This
should roughly equivalent to the old handling that create an IR symbol
table. Here we can prevent the module from materializing everything but
the single triple metadata we read in later.


>From 17d022557cdaad4b0fe5de71e32cd383c53f3663 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Fri, 15 Dec 2023 18:07:31 -0600
Subject: [PATCH] [Libomptarget] Remove bitcode image map used for JIT
 processing

Summary:
Libomptarget supports JIT by treating an LLVM-IR file as a regular input
image. The handling here used a global map to keep track of triples once
it was parsed. This was done to same time, however this created a global
constructor as well as an extra mutex to handle it. This patch removes
the use of this map.

Instead, we simply use the file magic to perform a quick check if the
input image is valid bitcode. If not, we then create a lazy module. This
should roughly equivalent to the old handling that create an IR symbol
table. Here we can prevent the module from materializing everything but
the single triple metadata we read in later.
---
 .../plugins-nextgen/common/src/JIT.cpp        | 51 +++++++------------
 1 file changed, 19 insertions(+), 32 deletions(-)

diff --git a/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp b/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp
index 188fb6dac84b9c..759a4f94523c73 100644
--- a/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp
@@ -47,18 +47,19 @@ using namespace llvm::object;
 using namespace omp;
 using namespace omp::target;
 
-static codegen::RegisterCodeGenFlags RCGF;
-
 namespace {
 
-/// A map from a bitcode image start address to its corresponding triple. If the
-/// image is not in the map, it is not a bitcode image.
-DenseMap<void *, Triple::ArchType> BitcodeImageMap;
-std::shared_mutex BitcodeImageMapMutex;
+bool isImageBitcode(const __tgt_device_image &Image) {
+  StringRef Binary(reinterpret_cast<const char *>(Image.ImageStart),
+                   target::getPtrDiff(Image.ImageEnd, Image.ImageStart));
+
+  return identify_magic(Binary) == file_magic::bitcode;
+}
 
 std::once_flag InitFlag;
 
 void init(Triple TT) {
+  codegen::RegisterCodeGenFlags();
 #ifdef LIBOMPTARGET_JIT_NVPTX
   if (TT.isNVPTX()) {
     LLVMInitializeNVPTXTargetInfo();
@@ -323,44 +324,30 @@ JITEngine::process(const __tgt_device_image &Image,
     return Device.doJITPostProcessing(std::move(MB));
   };
 
-  {
-    std::shared_lock<std::shared_mutex> SharedLock(BitcodeImageMapMutex);
-    auto Itr = BitcodeImageMap.find(Image.ImageStart);
-    if (Itr != BitcodeImageMap.end() && Itr->second == TT.getArch())
-      return compile(Image, ComputeUnitKind, PostProcessing);
-  }
+  if (isImageBitcode(Image))
+    return compile(Image, ComputeUnitKind, PostProcessing);
 
   return &Image;
 }
 
 bool JITEngine::checkBitcodeImage(const __tgt_device_image &Image) {
   TimeTraceScope TimeScope("Check bitcode image");
-  std::lock_guard<std::shared_mutex> Lock(BitcodeImageMapMutex);
 
-  {
-    auto Itr = BitcodeImageMap.find(Image.ImageStart);
-    if (Itr != BitcodeImageMap.end() && Itr->second == TT.getArch())
-      return true;
-  }
+  if (!isImageBitcode(Image))
+    return false;
 
   StringRef Data(reinterpret_cast<const char *>(Image.ImageStart),
                  target::getPtrDiff(Image.ImageEnd, Image.ImageStart));
-  std::unique_ptr<MemoryBuffer> MB = MemoryBuffer::getMemBuffer(
-      Data, /* BufferName */ "", /* RequiresNullTerminator */ false);
+  auto MB = MemoryBuffer::getMemBuffer(Data, /*BufferName=*/"",
+                                       /*RequiresNullTerminator=*/false);
   if (!MB)
     return false;
 
-  Expected<object::IRSymtabFile> FOrErr = object::readIRSymtab(*MB);
-  if (!FOrErr) {
-    consumeError(FOrErr.takeError());
-    return false;
-  }
-
-  auto ActualTriple = FOrErr->TheReader.getTargetTriple();
-  auto BitcodeTA = Triple(ActualTriple).getArch();
-  BitcodeImageMap[Image.ImageStart] = BitcodeTA;
-
-  DP("Is%s IR Image\n", BitcodeTA == TT.getArch() ? " " : " NOT");
+  LLVMContext Context;
+  SMDiagnostic Diagnostic;
+  std::unique_ptr<Module> M =
+      llvm::getLazyIRModule(std::move(MB), Diagnostic, Context,
+                            /*ShouldLazyLoadMetadata=*/true);
 
-  return BitcodeTA == TT.getArch();
+  return Triple(M->getTargetTriple()).getArch() == TT.getArch();
 }