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

[via Openmp-commits]

Author: Joseph Huber
Date: 2023-12-18T09:28:06-06:00
New Revision: 1580877555f91c5ab7eed7926bb984b2f6b70bb4

URL: https://github.com/llvm/llvm-project/commit/1580877555f91c5ab7eed7926bb984b2f6b70bb4
DIFF: https://github.com/llvm/llvm-project/commit/1580877555f91c5ab7eed7926bb984b2f6b70bb4.diff

LOG: [Libomptarget] Remove bitcode image map used for JIT processing (#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.

Added: 
    

Modified: 
    openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp

Removed: 
    


################################################################################
diff  --git a/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp b/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp
index 188fb6dac84b9c..08080c9d6091bb 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 M && Triple(M->getTargetTriple()).getArch() == TT.getArch();
 }