[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.
GitOrigin-RevId: 1580877555f91c5ab7eed7926bb984b2f6b70bb4
diff --git a/libomptarget/plugins-nextgen/common/src/JIT.cpp b/libomptarget/plugins-nextgen/common/src/JIT.cpp
index 188fb6d..08080c9 100644
--- a/libomptarget/plugins-nextgen/common/src/JIT.cpp
+++ b/libomptarget/plugins-nextgen/common/src/JIT.cpp
@@ -47,18 +47,19 @@
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 @@
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;
- }
+ LLVMContext Context;
+ SMDiagnostic Diagnostic;
+ std::unique_ptr<Module> M =
+ llvm::getLazyIRModule(std::move(MB), Diagnostic, Context,
+ /*ShouldLazyLoadMetadata=*/true);
- auto ActualTriple = FOrErr->TheReader.getTargetTriple();
- auto BitcodeTA = Triple(ActualTriple).getArch();
- BitcodeImageMap[Image.ImageStart] = BitcodeTA;
-
- DP("Is%s IR Image\n", BitcodeTA == TT.getArch() ? " " : " NOT");
-
- return BitcodeTA == TT.getArch();
+ return M && Triple(M->getTargetTriple()).getArch() == TT.getArch();
}
