[OpenMP][JIT] Cleanup JIT interface, caching, and races

[jdoerfert]

Summary

The JIT interface was somewhat irregular as it used multiple global
functions. It also did not cache the results of the JIT, hence multiple
GPU systems would perform the work multiple times. Finally, there might
have been races on the state if we have multi-threaded initialization of
different embedded images, or one image initialized on multiple devices.

This patch tries to rectify all of the above. The JITEngine is now a
part of the GenericPluginTy and tied to one target triple. To support
multiple "ComputeUnitKind"s (previously confusingly called Arch or
[M]CPU) and to avoid re-jitting for the same ComputeUnitKind, we keep a
map of JIT results per ComputeUnitKind. All interaction with the JIT
happens through the JITEngine directly, two functions are exposed. Both
use (shared) locks to avoid races and cache the result. All JIT-related
environment variables are now defined together.

[jdoerfert]

Forgot the context, will update it later if necessary.

[jdoerfert]

openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.h

98

I need to rename CPU to ComputeUnit here too.

[tianshilei1992]

When do we do multi-threading initialization?

[jdoerfert]

Include context, fix some oversights.

[jdoerfert]

In D141081#4029796, @tianshilei1992 wrote:

When do we do multi-threading initialization?

TBH, I don't know. Even if not, it seems appropriate (given past errors) to design everything with multi-threading in mind, except if it is explicitly excluded. The cost for the 2 mutexes should be minimal.
I'm also not super certain we can't have two threads calling __tgt_rtl_load_binary on two devices concurrently.

[tianshilei1992]

LG, the name is really weird though.

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   Error initImpl(GenericPluginTy &Plugin) override {
     // First setup all the memory pools.
     if (auto Err = initMemoryPools())
       return Err;
 
     char GPUName[64];
     if (auto Err = getDeviceAttr(HSA_AGENT_INFO_NAME, GPUName))
       return Err;
-    Arch = GPUName;
+    ComputeUnitKind = GPUName;
 
     // Get the wavefront size.
     uint32_t WavefrontSize = 0;
     if (auto Err = getDeviceAttr(HSA_AGENT_INFO_WAVEFRONT_SIZE, WavefrontSize))
       return Err;
     GridValues.GV_Warp_Size = WavefrontSize;
 
     // Load the grid values dependending on the wavefront.
doJITPostProcessing(std::unique_ptr<MemoryBuffer> MB) const override {
     if (!ErrorOrPath)
       return createStringError(inconvertibleErrorCode(),
                                "Failed to find `lld` on the PATH.");
 
     std::string LLDPath = ErrorOrPath.get();
     INFO(OMP_INFOTYPE_PLUGIN_KERNEL, getDeviceId(),
          "Using `%s` to link JITed amdgcn ouput.", LLDPath.c_str());
 
-    std::string MCPU = "-plugin-opt=mcpu=" + getArch();
+    std::string MCPU = "-plugin-opt=mcpu=" + getComputeUnitKind();
 
     StringRef Args[] = {LLDPath,
                         "-flavor",
                         "gnu",
                         "--no-undefined",
                         "-shared",
                         MCPU,
                         "-o",
                         LinkerOutputFilePath.data(),
                         LinkerInputFilePath.data()};
 
     std::string Error;
     int RC = sys::ExecuteAndWait(LLDPath, Args, std::nullopt, {}, 0, 0, &Error);
     if (RC)
       return createStringError(inconvertibleErrorCode(),
                                "Linking optimized bitcode failed: %s",
                                Error.c_str());
 
     return std::move(
         MemoryBuffer::getFileOrSTDIN(LinkerOutputFilePath.data()).get());
   }
 
-  std::string getArch() const override { return Arch; }
+  /// See GenericDeviceTy::getComputeUnitKind().
+  std::string getComputeUnitKind() const override { return ComputeUnitKind; }
 
   /// Allocate and construct an AMDGPU kernel.
   Expected<GenericKernelTy *>
   constructKernelEntry(const __tgt_offload_entry &KernelEntry,
                        DeviceImageTy &Image) override {
 
     Expected<OMPTgtExecModeFlags> ExecModeOrErr =
         getExecutionModeForKernel(KernelEntry.name, Image);
private:
 
   /// Signal manager for AMDGPU signals.
   AMDGPUSignalManagerTy AMDGPUSignalManager;
 
   /// The agent handler corresponding to the device.
   hsa_agent_t Agent;
 
   /// The GPU architecture.
-  std::string Arch;
+  std::string ComputeUnitKind;
 
   /// Reference to the host device.
   AMDHostDeviceTy &HostDevice;
 
   /// List of device packet queues.
   std::vector<AMDGPUQueueTy> Queues;
 };
 
Error getGlobalMetadataFromELF(const DeviceImageTy &Image,
 
     return Plugin::success();
   }
 };
 
 /// Class implementing the AMDGPU-specific functionalities of the plugin.
 struct AMDGPUPluginTy final : public GenericPluginTy {
   /// Create an AMDGPU plugin and initialize the AMDGPU driver.
-  AMDGPUPluginTy() : GenericPluginTy(), HostDevice(nullptr) {}
+  AMDGPUPluginTy() : GenericPluginTy(getTripleArch()), HostDevice(nullptr) {}
 
   /// This class should not be copied.
   AMDGPUPluginTy(const AMDGPUPluginTy &) = delete;
   AMDGPUPluginTy(AMDGPUPluginTy &&) = delete;
 
   /// Initialize the plugin and return the number of devices.
   Expected<int32_t> initImpl() override {
     hsa_status_t Status = hsa_init();

openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.h

 //===- JIT.h - Target independent JIT infrastructure ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H
 #define OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H
 
+#include "Utilities.h"
+
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/IR/Module.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Target/TargetMachine.h"
 
 #include <functional>
 #include <memory>
+#include <shared_mutex>
 #include <string>
 
 struct __tgt_device_image;
 
 namespace llvm {
 class MemoryBuffer;
 
 namespace omp {
-namespace jit {
+namespace target {
+namespace plugin {
+struct GenericDeviceTy;
+} // namespace plugin
 
+/// The JIT infrastructure and caching mechanism.
+struct JITEngine {
   /// Function type for a callback that will be called after the backend is
   /// called.
-using PostProcessingFn = std::function<Expected<std::unique_ptr<MemoryBuffer>>(
+  using PostProcessingFn =
+      std::function<Expected<std::unique_ptr<MemoryBuffer>>(
           std::unique_ptr<MemoryBuffer>)>;
 
-/// Check if \p Image contains bitcode with triple \p Triple.
-bool checkBitcodeImage(__tgt_device_image *Image, Triple::ArchType TA);
+  JITEngine(Triple::ArchType TA);
 
-/// Compile the bitcode image \p Image and generate the binary image that can be
-/// loaded to the target device of the triple \p Triple architecture \p MCpu. \p
-/// PostProcessing will be called after codegen to handle cases such as assember
-/// as an external tool.
-Expected<__tgt_device_image *> compile(__tgt_device_image *Image,
-                                       Triple::ArchType TA, std::string MCpu,
-                                       unsigned OptLevel,
-                                       PostProcessingFn PostProcessing);
-} // namespace jit
+  /// Run jit compilation if \p Image is a bitcode image, otherwise simply
+  /// return \p Image. It is expected to return a memory buffer containing the
+  /// generated device image that could be loaded to the device directly.
+  Expected<const __tgt_device_image *>
+  process(const __tgt_device_image &Image,
+          target::plugin::GenericDeviceTy &Device);
+
+  /// Return true if \p Image is a bitcode image that can be JITed for the given
+  /// architecture.
+  bool checkBitcodeImage(const __tgt_device_image &Image);
+
+private:
+  /// Compile the bitcode image \p Image and generate the binary image that can
+  /// be loaded to the target device of the triple \p Triple architecture \p
+  /// MCpu. \p PostProcessing will be called after codegen to handle cases such
+  /// as assember as an external tool.
+  Expected<const __tgt_device_image *>
+  compile(const __tgt_device_image &Image, const std::string &ComputeUnitKind,
+
+  /// Run backend, which contains optimization and code generation.
+  Expected<std::unique_ptr<MemoryBuffer>>
+  backend(Module &M, const std::string &ComputeUnitKind, unsigned OptLevel);
+
+  /// Run optimization pipeline.
+  void opt(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
+           unsigned OptLevel);
+
+  /// Run code generation.
+  void codegen(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
+               raw_pwrite_stream &OS);
+
+  /// The target triple used by the JIT.
+  const Triple TT;
+
+  struct ComputeUnitInfo {
+    /// LLVM Context in which the modules will be constructed.
+    LLVMContext Context;
+
+    /// Output images generated from LLVM backend.
+    SmallVector<std::unique_ptr<MemoryBuffer>, 4> JITImages;
+
+    /// A map of embedded IR images to JITed images.
+    DenseMap<const __tgt_device_image *, __tgt_device_image *> TgtImageMap;
+  };
+
+  /// Map from (march) "CPUs" (e.g., sm_80, or gfx90a), which we call compute
+  /// units as they are not CPUs, to the image information we cached for them.
+  StringMap<ComputeUnitInfo> ComputeUnitMap;
+  std::mutex ComputeUnitMapMutex;
+
+  /// Control environment variables.
+  target::StringEnvar ReplacementModuleFileName =
+      target::StringEnvar("LIBOMPTARGET_JIT_REPLACEMENT_MODULE");
+  target::StringEnvar PreOptIRModuleFileName =
+      target::StringEnvar("LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE");
+  target::StringEnvar PostOptIRModuleFileName =
+      target::StringEnvar("LIBOMPTARGET_JIT_POST_OPT_IR_MODULE");
+  target::UInt32Envar JITOptLevel =
+      target::UInt32Envar("LIBOMPTARGET_JIT_OPT_LEVEL", 3);
+};
+
+} // namespace target
 } // namespace omp
 } // namespace llvm
 
 #endif // OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H

openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.cpp

 //===- JIT.cpp - Target independent JIT infrastructure --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //===----------------------------------------------------------------------===//
 
 #include "JIT.h"
 #include "Debug.h"
 
+#include "PluginInterface.h"
 #include "Utilities.h"
 #include "omptarget.h"
 
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/CommandFlags.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LLVMRemarkStreamer.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/MC/SubtargetFeature.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Object/IRObjectFile.h"
 #include "llvm/Passes/OptimizationLevel.h"
 #include "llvm/Passes/PassBuilder.h"
-#include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/TimeProfiler.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 
 #include <mutex>
+#include <shared_mutex>
 #include <system_error>
 
 using namespace llvm;
 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;
+
 std::once_flag InitFlag;
 
 void init(Triple TT) {
   bool JITTargetInitialized = false;
 #ifdef LIBOMPTARGET_JIT_NVPTX
   if (TT.isNVPTX()) {
     LLVMInitializeNVPTXTargetInfo();
     LLVMInitializeNVPTXTarget();
     LLVMInitializeNVPTXTargetMC();
     LLVMInitializeNVPTXAsmPrinter();
     JITTargetInitialized = true;
   }
 #endif
 #ifdef LIBOMPTARGET_JIT_AMDGPU
   if (TT.isAMDGPU()) {
     LLVMInitializeAMDGPUTargetInfo();
     LLVMInitializeAMDGPUTarget();
     LLVMInitializeAMDGPUTargetMC();
     LLVMInitializeAMDGPUAsmPrinter();
     JITTargetInitialized = true;
   }
 #endif
-  if (!JITTargetInitialized) {
-    FAILURE_MESSAGE("unsupported JIT target: %s\n", TT.str().c_str());
-    abort();
-  }
+  if (!JITTargetInitialized)
+    return;
 
   // Initialize passes
   PassRegistry &Registry = *PassRegistry::getPassRegistry();
   initializeCore(Registry);
   initializeScalarOpts(Registry);
   initializeVectorization(Registry);
   initializeIPO(Registry);
   initializeAnalysis(Registry);
createModuleFromMemoryBuffer(std::unique_ptr<MemoryBuffer> &MB,
   SMDiagnostic Err;
   auto Mod = parseIR(*MB, Err, Context);
   if (!Mod)
     return make_error<StringError>("Failed to create module",
                                    inconvertibleErrorCode());
   return std::move(Mod);
 }
 Expected<std::unique_ptr<Module>>
-createModuleFromImage(__tgt_device_image *Image, LLVMContext &Context) {
-  StringRef Data((const char *)Image->ImageStart,
-                 target::getPtrDiff(Image->ImageEnd, Image->ImageStart));
+createModuleFromImage(const __tgt_device_image &Image, LLVMContext &Context) {
+  StringRef Data((const char *)Image.ImageStart,
+                 target::getPtrDiff(Image.ImageEnd, Image.ImageStart));
   std::unique_ptr<MemoryBuffer> MB = MemoryBuffer::getMemBuffer(
       Data, /* BufferName */ "", /* RequiresNullTerminator */ false);
   return createModuleFromMemoryBuffer(MB, Context);
 }
 
 CodeGenOpt::Level getCGOptLevel(unsigned OptLevel) {
   switch (OptLevel) {
   case 0:
std::unique_ptr<TargetMachine> TM(
       T->createTargetMachine(M.getTargetTriple(), CPU, Features.getString(),
                              Options, RelocModel, CodeModel, CGOptLevel));
   if (!TM)
     return make_error<StringError>("Failed to create target machine",
                                    inconvertibleErrorCode());
   return std::move(TM);
 }
 
-///
-class JITEngine {
-public:
-  JITEngine(Triple::ArchType TA, std::string MCpu)
-      : TT(Triple::getArchTypeName(TA)), CPU(MCpu),
-        ReplacementModuleFileName("LIBOMPTARGET_JIT_REPLACEMENT_MODULE"),
-        PreOptIRModuleFileName("LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE"),
-        PostOptIRModuleFileName("LIBOMPTARGET_JIT_POST_OPT_IR_MODULE") {
-    std::call_once(InitFlag, init, TT);
-  }
+} // namespace
+
+JITEngine::JITEngine(Triple::ArchType TA) : TT(Triple::getArchTypeName(TA)) {
 
-  /// Run jit compilation. It is expected to get a memory buffer containing the
-  /// generated device image that could be loaded to the device directly.
-  Expected<std::unique_ptr<MemoryBuffer>>
-  run(__tgt_device_image *Image, unsigned OptLevel,
-      jit::PostProcessingFn PostProcessing);
-
-private:
-  /// Run backend, which contains optimization and code generation.
-  Expected<std::unique_ptr<MemoryBuffer>> backend(Module &M, unsigned OptLevel);
-
-  /// Run optimization pipeline.
-  void opt(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
-           unsigned OptLevel);
-
-  /// Run code generation.
-  void codegen(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
-               raw_pwrite_stream &OS);
-
-  LLVMContext Context;
-  const Triple TT;
-  const std::string CPU;
-
-  /// Control environment variables.
-  target::StringEnvar ReplacementModuleFileName;
-  target::StringEnvar PreOptIRModuleFileName;
-  target::StringEnvar PostOptIRModuleFileName;
-};
-
 void JITEngine::opt(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
                     unsigned OptLevel) {
   PipelineTuningOptions PTO;
   std::optional<PGOOptions> PGOOpt;
 
   LoopAnalysisManager LAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
MachineModuleInfoWrapperPass *MMIWP = new MachineModuleInfoWrapperPass(
       reinterpret_cast<LLVMTargetMachine *>(TM));
   TM->addPassesToEmitFile(PM, OS, nullptr,
                           TT.isNVPTX() ? CGFT_AssemblyFile : CGFT_ObjectFile,
                           /* DisableVerify */ false, MMIWP);
 
   PM.run(M);
 }
 
-Expected<std::unique_ptr<MemoryBuffer>> JITEngine::backend(Module &M,
-                                                           unsigned OptLevel) {
+Expected<std::unique_ptr<MemoryBuffer>>
+JITEngine::backend(Module &M, const std::string &ComputeUnitKind,
 
   auto RemarksFileOrErr = setupLLVMOptimizationRemarks(
-      Context, /* RemarksFilename */ "", /* RemarksPasses */ "",
+      M.getContext(), /* RemarksFilename */ "", /* RemarksPasses */ "",
       /* RemarksFormat */ "", /* RemarksWithHotness */ false);
   if (Error E = RemarksFileOrErr.takeError())
     return std::move(E);
   if (*RemarksFileOrErr)
     (*RemarksFileOrErr)->keep();
 
-  auto TMOrErr = createTargetMachine(M, CPU, OptLevel);
+  auto TMOrErr = createTargetMachine(M, ComputeUnitKind, OptLevel);
   if (!TMOrErr)
     return TMOrErr.takeError();
 
   std::unique_ptr<TargetMachine> TM = std::move(*TMOrErr);
   TargetLibraryInfoImpl TLII(TT);
 
   if (PreOptIRModuleFileName.isPresent()) {
     std::error_code EC;
JITEngine::backend(Module &M, const std::string &ComputeUnitKind,
   SmallVector<char> CGOutputBuffer;
   raw_svector_ostream OS(CGOutputBuffer);
 
   codegen(TM.get(), &TLII, M, OS);
 
   return MemoryBuffer::getMemBufferCopy(OS.str());
 }
 
-Expected<std::unique_ptr<MemoryBuffer>>
-JITEngine::run(__tgt_device_image *Image, unsigned OptLevel,
-               jit::PostProcessingFn PostProcessing) {
+Expected<const __tgt_device_image *>
+JITEngine::compile(const __tgt_device_image &Image,
+                   const std::string &ComputeUnitKind,
+                   PostProcessingFn PostProcessing) {
+  std::lock_guard<std::mutex> Lock(ComputeUnitMapMutex);
+
+  // Check if we JITed this image for the given compute unit kind before.
+  ComputeUnitInfo &CUI = ComputeUnitMap[ComputeUnitKind];
+  if (__tgt_device_image *JITedImage = CUI.TgtImageMap.lookup(&Image))
+    return JITedImage;
+
   Module *Mod = nullptr;
   // Check if the user replaces the module at runtime or we read it from the
   // image.
+  // TODO: Allow the user to specify images per device (Arch + ComputeUnitKind).
   if (!ReplacementModuleFileName.isPresent()) {
-    auto ModOrErr = createModuleFromImage(Image, Context);
+    auto ModOrErr = createModuleFromImage(Image, CUI.Context);
     if (!ModOrErr)
       return ModOrErr.takeError();
     Mod = ModOrErr->release();
   } else {
     auto MBOrErr =
         MemoryBuffer::getFileOrSTDIN(ReplacementModuleFileName.get());
     if (!MBOrErr)
       return createStringError(MBOrErr.getError(),
                                "Could not read replacement module from %s\n",
                                ReplacementModuleFileName.get().c_str());
-    auto ModOrErr = createModuleFromMemoryBuffer(MBOrErr.get(), Context);
+    auto ModOrErr = createModuleFromMemoryBuffer(MBOrErr.get(), CUI.Context);
     if (!ModOrErr)
       return ModOrErr.takeError();
     Mod = ModOrErr->release();
   }
 
-  auto MBOrError = backend(*Mod, OptLevel);
+  auto MBOrError = backend(*Mod, ComputeUnitKind, JITOptLevel);
   if (!MBOrError)
     return MBOrError.takeError();
 
-  return PostProcessing(std::move(*MBOrError));
+  auto ImageMBOrErr = PostProcessing(std::move(*MBOrError));
+  if (!ImageMBOrErr)
+    return ImageMBOrErr.takeError();
+
+  CUI.JITImages.push_back(std::move(*ImageMBOrErr));
+  __tgt_device_image *&JITedImage = CUI.TgtImageMap[&Image];
+  JITedImage = new __tgt_device_image();
+  *JITedImage = Image;
+
+  auto &ImageMB = CUI.JITImages.back();
+
+  JITedImage->ImageStart = (void *)ImageMB->getBufferStart();
+  JITedImage->ImageEnd = (void *)ImageMB->getBufferEnd();
+
+  return JITedImage;
 }
 
-/// 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;
+Expected<const __tgt_device_image *>
+JITEngine::process(const __tgt_device_image &Image,
+                   target::plugin::GenericDeviceTy &Device) {
+  const std::string &ComputeUnitKind = Device.getComputeUnitKind();
 
-/// Output images generated from LLVM backend.
-SmallVector<std::unique_ptr<MemoryBuffer>, 4> JITImages;
+  PostProcessingFn PostProcessing = [&Device](std::unique_ptr<MemoryBuffer> MB)
+      -> Expected<std::unique_ptr<MemoryBuffer>> {
+    return Device.doJITPostProcessing(std::move(MB));
+  };
 
-/// A list of __tgt_device_image images.
-std::list<__tgt_device_image> TgtImages;
-} // namespace
+  {
+    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);
+  }
 
-namespace llvm {
-namespace omp {
-namespace jit {
-bool checkBitcodeImage(__tgt_device_image *Image, Triple::ArchType TA) {
+  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 == TA)
+    auto Itr = BitcodeImageMap.find(Image.ImageStart);
+    if (Itr != BitcodeImageMap.end() && Itr->second == TT.getArch())
       return true;
   }
 
-  StringRef Data(reinterpret_cast<const char *>(Image->ImageStart),
-                 target::getPtrDiff(Image->ImageEnd, Image->ImageStart));
+  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);
   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;
 
-  if (Triple(ActualTriple).getArch() == TA) {
-    BitcodeImageMap[Image->ImageStart] = TA;
-    return true;
-  }
-
-  return false;
+  return BitcodeTA == TT.getArch();
 }
-
-Expected<__tgt_device_image *> compile(__tgt_device_image *Image,
-                                       Triple::ArchType TA, std::string MCPU,
-                                       unsigned OptLevel,
-                                       PostProcessingFn PostProcessing) {
-  JITEngine J(TA, MCPU);
-
-  auto ImageMBOrErr = J.run(Image, OptLevel, PostProcessing);
-  if (!ImageMBOrErr)
-    return ImageMBOrErr.takeError();
-
-  JITImages.push_back(std::move(*ImageMBOrErr));
-  TgtImages.push_back(*Image);
-
-  auto &ImageMB = JITImages.back();
-  auto *NewImage = &TgtImages.back();
-
-  NewImage->ImageStart = (void *)ImageMB->getBufferStart();
-  NewImage->ImageEnd = (void *)ImageMB->getBufferEnd();
-
-  return NewImage;
-}
-
-} // namespace jit
-} // namespace omp
-} // namespace llvm

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

 #include <list>
 #include <map>
 #include <shared_mutex>
 #include <vector>
 
 #include "Debug.h"
 #include "DeviceEnvironment.h"
 #include "GlobalHandler.h"
+#include "JIT.h"
 #include "MemoryManager.h"
 #include "Utilities.h"
 #include "omptarget.h"
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MemoryBufferRef.h"
 
 namespace llvm {
 namespace omp {
 namespace target {
+
 namespace plugin {
 
 struct GenericPluginTy;
 struct GenericKernelTy;
 struct GenericDeviceTy;
 
 /// Class that wraps the __tgt_async_info to simply its usage. In case the
 /// object is constructed without a valid __tgt_async_info, the object will use
struct GenericDeviceTy : public DeviceAllocatorTy {
   uint32_t getDefaultNumThreads() const {
     return GridValues.GV_Default_WG_Size;
   }
   uint64_t getDefaultNumBlocks() const {
     return GridValues.GV_Default_Num_Teams;
   }
   uint32_t getDynamicMemorySize() const { return OMPX_SharedMemorySize; }
 
-  /// Get target architecture.
-  virtual std::string getArch() const {
-    return "unknown";
-  }
+  /// Get target compute unit kind (e.g., sm_80, or gfx908).
+  virtual std::string getComputeUnitKind() const { return "unknown"; }
 
   /// Post processing after jit backend. The ownership of \p MB will be taken.
   virtual Expected<std::unique_ptr<MemoryBuffer>>
   doJITPostProcessing(std::unique_ptr<MemoryBuffer> MB) const {
     return std::move(MB);
   }
 
 private:

 };
 
 /// Class implementing common functionalities of offload plugins. Each plugin
 /// should define the specific plugin class, derive from this generic one, and
 /// implement the necessary virtual function members.
 struct GenericPluginTy {
 
   /// Construct a plugin instance.
-  GenericPluginTy()
-      : RequiresFlags(OMP_REQ_UNDEFINED), GlobalHandler(nullptr) {}
+  GenericPluginTy(Triple::ArchType TA)
+      : RequiresFlags(OMP_REQ_UNDEFINED), GlobalHandler(nullptr), JIT(TA) {}
 
   virtual ~GenericPluginTy() {}
 
   /// Initialize the plugin.
   Error init();
 
   /// Initialize the plugin and return the number of available devices.
   virtual Expected<int32_t> initImpl() = 0;
struct GenericPluginTy {
 
   /// Get the number of active devices.
   int32_t getNumDevices() const { return NumDevices; }
 
   /// Get the ELF code to recognize the binary image of this plugin.
   virtual uint16_t getMagicElfBits() const = 0;
 
   /// Get the target triple of this plugin.
-  virtual Triple::ArchType getTripleArch() const {
-    return Triple::ArchType::UnknownArch;
-  }
+  virtual Triple::ArchType getTripleArch() const = 0;
 
   /// Allocate a structure using the internal allocator.
   template <typename Ty> Ty *allocate() {
     return reinterpret_cast<Ty *>(Allocator.Allocate(sizeof(Ty), alignof(Ty)));
   }
 
   /// Get the reference to the global handler of this plugin.
   GenericGlobalHandlerTy &getGlobalHandler() {
     assert(GlobalHandler && "Global handler not initialized");
     return *GlobalHandler;
   }
 
+  /// Get the reference to the JIT used for all devices connected to this
+  /// plugin.
+  JITEngine &getJIT() { return JIT; }
+
   /// Get the OpenMP requires flags set for this plugin.
   int64_t getRequiresFlags() const { return RequiresFlags; }
 
   /// Set the OpenMP requires flags for this plugin.
   void setRequiresFlag(int64_t Flags) { RequiresFlags = Flags; }
 
   /// Initialize a device within the plugin.
   Error initDevice(int32_t DeviceId);
private:
   /// OpenMP requires flags.
   int64_t RequiresFlags;
 
   /// Pointer to the global handler for this plugin.
   GenericGlobalHandlerTy *GlobalHandler;
 
   /// Internal allocator for different structures.
   BumpPtrAllocator Allocator;
+
+  /// The JIT engine shared by all devices connected to this plugin.
+  JITEngine JIT;
 };
 
 /// Class for simplifying the getter operation of the plugin. Anywhere on the
 /// code, the current plugin can be retrieved by Plugin::get(). The class also
 /// declares functions to create plugin-specific object instances. The check(),
 /// createPlugin(), createDevice() and createGlobalHandler() functions should be
 /// defined by each plugin implementation.
 class Plugin {

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

if (MemoryManager)
     delete MemoryManager;
   MemoryManager = nullptr;
 
   return deinitImpl();
 }
 
 Expected<__tgt_target_table *>
 GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
-                            const __tgt_device_image *TgtImage) {
-  DP("Load data from image " DPxMOD "\n", DPxPTR(TgtImage->ImageStart));
+                            const __tgt_device_image *InputTgtImage) {
+  assert(InputTgtImage && "Expected non-null target image");
+  DP("Load data from image " DPxMOD "\n", DPxPTR(InputTgtImage->ImageStart));
+
+  auto PostJITImageOrErr = Plugin.getJIT().process(*InputTgtImage, *this);
+  if (!PostJITImageOrErr) {
+    auto Err = PostJITImageOrErr.takeError();
+    REPORT("Failure to jit IR image %p on device %d: %s\n", InputTgtImage,
+           DeviceId, toString(std::move(Err)).data());
+    return nullptr;
+  }
 
   // Load the binary and allocate the image object. Use the next available id
   // for the image id, which is the number of previously loaded images.
-  auto ImageOrErr = loadBinaryImpl(TgtImage, LoadedImages.size());
+  auto ImageOrErr =
+      loadBinaryImpl(PostJITImageOrErr.get(), LoadedImages.size());
   if (!ImageOrErr)
     return ImageOrErr.takeError();
 
   DeviceImageTy *Image = *ImageOrErr;
   assert(Image != nullptr && "Invalid image");
 
   // Add the image to list.
   LoadedImages.push_back(Image);

 
 int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *TgtImage) {
   if (!Plugin::isActive())
     return false;
 
   if (elf_check_machine(TgtImage, Plugin::get().getMagicElfBits()))
     return true;
 
-  return jit::checkBitcodeImage(TgtImage, Plugin::get().getTripleArch());
+  return Plugin::get().getJIT().checkBitcodeImage(*TgtImage);
 }
 
 int32_t __tgt_rtl_is_valid_binary_info(__tgt_device_image *TgtImage,
                                        __tgt_image_info *Info) {
   if (!Plugin::isActive())
     return false;
 
   if (!__tgt_rtl_is_valid_binary(TgtImage))
int32_t __tgt_rtl_is_data_exchangable(int32_t SrcDeviceId,
   return Plugin::get().isDataExchangable(SrcDeviceId, DstDeviceId);
 }
 
 __tgt_target_table *__tgt_rtl_load_binary(int32_t DeviceId,
                                           __tgt_device_image *TgtImage) {
   GenericPluginTy &Plugin = Plugin::get();
   GenericDeviceTy &Device = Plugin.getDevice(DeviceId);
 
-  // If it is a bitcode image, we have to jit the binary image before loading to
-  // the device.
-  {
-    // TODO: Move this (at least the environment variable) into the JIT.h.
-    UInt32Envar JITOptLevel("LIBOMPTARGET_JIT_OPT_LEVEL", 3);
-    Triple::ArchType TA = Plugin.getTripleArch();
-    std::string Arch = Device.getArch();
-
-    jit::PostProcessingFn PostProcessing =
-        [&Device](std::unique_ptr<MemoryBuffer> MB)
-        -> Expected<std::unique_ptr<MemoryBuffer>> {
-      return Device.doJITPostProcessing(std::move(MB));
-    };
-
-    if (jit::checkBitcodeImage(TgtImage, TA)) {
-      auto TgtImageOrErr =
-          jit::compile(TgtImage, TA, Arch, JITOptLevel, PostProcessing);
-      if (!TgtImageOrErr) {
-        auto Err = TgtImageOrErr.takeError();
-        REPORT("Failure to jit binary image from bitcode image %p on device "
-               "%d: %s\n",
-               TgtImage, DeviceId, toString(std::move(Err)).data());
-        return nullptr;
-      }
-
-      TgtImage = *TgtImageOrErr;
-    }
-  }
 
   auto TableOrErr = Device.loadBinary(Plugin, TgtImage);
   if (!TableOrErr) {
     auto Err = TableOrErr.takeError();
     REPORT("Failure to load binary image %p on device %d: %s\n", TgtImage,
            DeviceId, toString(std::move(Err)).data());
     return nullptr;
   }

openmp/libomptarget/plugins-nextgen/cuda/src/rtl.cpp

struct CUDADeviceTy : public GenericDeviceTy {
   /// CUDA-specific function to get device attributes.
   Error getDeviceAttr(uint32_t Kind, uint32_t &Value) {
     // TODO: Warn if the new value is larger than the old.
     CUresult Res =
         cuDeviceGetAttribute((int *)&Value, (CUdevice_attribute)Kind, Device);
     return Plugin::check(Res, "Error in cuDeviceGetAttribute: %s");
   }
 
-  /// See GenericDeviceTy::getArch().
-  std::string getArch() const override { return ComputeCapability.str(); }
+  /// See GenericDeviceTy::getComputeUnitKind().
+  std::string getComputeUnitKind() const override {
+    return ComputeCapability.str();
+  }
 
 private:
   using CUDAStreamManagerTy = GenericDeviceResourceManagerTy<CUDAStreamRef>;
   using CUDAEventManagerTy = GenericDeviceResourceManagerTy<CUDAEventRef>;
 
   /// Stream manager for CUDA streams.
   CUDAStreamManagerTy CUDAStreamManager;
 
Error getGlobalMetadataFromDevice(GenericDeviceTy &Device,
     DeviceGlobal.setPtr(reinterpret_cast<void *>(CUPtr));
     return Plugin::success();
   }
 };
 
 /// Class implementing the CUDA-specific functionalities of the plugin.
 struct CUDAPluginTy final : public GenericPluginTy {
   /// Create a CUDA plugin.
-  CUDAPluginTy() : GenericPluginTy() {}
+  CUDAPluginTy() : GenericPluginTy(getTripleArch()) {}
 
   /// This class should not be copied.
   CUDAPluginTy(const CUDAPluginTy &) = delete;
   CUDAPluginTy(CUDAPluginTy &&) = delete;
 
   /// Initialize the plugin and return the number of devices.
   Expected<int32_t> initImpl() override {
     CUresult Res = cuInit(0);

openmp/libomptarget/plugins-nextgen/generic-elf-64bit/src/rtl.cpp

Error getGlobalMetadataFromDevice(GenericDeviceTy &GenericDevice,
 
     return Plugin::success();
   }
 };
 
 /// Class implementing the plugin functionalities for GenELF64.
 struct GenELF64PluginTy final : public GenericPluginTy {
   /// Create the GenELF64 plugin.
-  GenELF64PluginTy() : GenericPluginTy() {}
+  GenELF64PluginTy() : GenericPluginTy(getTripleArch()) {}
 
   /// This class should not be copied.
   GenELF64PluginTy(const GenELF64PluginTy &) = delete;
   GenELF64PluginTy(GenELF64PluginTy &&) = delete;
 
   /// Initialize the plugin and return the number of devices.
   Expected<int32_t> initImpl() override { return NUM_DEVICES; }