[OpenMP][JIT] Introduce environment variables to deal with JIT IR

[jdoerfert]

Summary

We can now dump the IR before and after JIT optimizations into the
files passed via LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE and
LIBOMPTARGET_JIT_POST_OPT_IR_MODULE, respectively.

Similarly, users can set LIBOMPTARGET_JIT_REPLACEMENT_MODULE to
replace the IR in the image with a custom IR module in a file.
All options take file paths, documentation was added.

[tianshilei1992]

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

298

We don’t check the error code?

[jdoerfert]

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

298

I'll add it.

openmp/docs/design/Runtimes.rst

variables is defined below.
     * ``LIBOMPTARGET_DEBUG=<Num>``
     * ``LIBOMPTARGET_PROFILE=<Filename>``
     * ``LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD=<Num>``
     * ``LIBOMPTARGET_INFO=<Num>``
     * ``LIBOMPTARGET_HEAP_SIZE=<Num>``
     * ``LIBOMPTARGET_STACK_SIZE=<Num>``
     * ``LIBOMPTARGET_SHARED_MEMORY_SIZE=<Num>``
     * ``LIBOMPTARGET_MAP_FORCE_ATOMIC=[TRUE/FALSE] (default TRUE)``
+    * ``LIBOMPTARGET_JIT_OPT_LEVEL={0,1,2,3} (default 3)``
+    * ``LIBOMPTARGET_JIT_REPLACEMENT_MODULE=<in:Filename> (LLVM-IR file)``
+    * ``LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE=<out:Filename> (LLVM-IR file)``
+    * ``LIBOMPTARGET_JIT_POST_OPT_IR_MODULE=<out:Filename> (LLVM-IR file)``
 
 LIBOMPTARGET_DEBUG
 """"""""""""""""""
 
 ``LIBOMPTARGET_DEBUG`` controls whether or not debugging information will be
 displayed. This feature is only availible if ``libomptarget`` was built with
 ``-DOMPTARGET_DEBUG``. The debugging output provided is intended for use by
 ``libomptarget`` developers. More user-friendly output is presented when using

 concurrently map the same memory. If the memory is already mapped and the
 map clauses will only modify the reference counter from a non-zero count to
 another non-zero count, concurrent map clauses are supported regardless of
 this option. To disable forced atomic map clauses use "false"/"FALSE" as the
 value of the ``LIBOMPTARGET_MAP_FORCE_ATOMIC`` environment variable.
 The default behavior of LLVM 14 is to force atomic maps clauses, prior versions
 of LLVM did not.
 
+
+LIBOMPTARGET_JIT_OPT_LEVEL
+""""""""""""""""""""""""""
+
+This environment variable can be used to change the optimization pipeleine used
+to optimize the embedded device code as part of the device JIT. The value is
+corresponds to the ``-O{0,1,2,3}`` command line argument passed to ``clang``.
+
+
+LIBOMPTARGET_JIT_REPLACEMENT_MODULE
+"""""""""""""""""""""""""""""""""""
+
+This environment variable can be used to replace the embedded device code
+before the device JIT finishes compilation for the target. The value is
+expected to be a filename to an LLVM-IR file, thus containing an LLVM-IR module
+for the respective target. To obtain a device code image compatible with the
+embedded one it is recommended to extract the embedded one either before or
+after IR optimization. This can be done at compile time, after compile time via
+llvm tools (llvm-objdump), or, simply, by setting the
+:ref:`LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE` or
+:ref:`LIBOMPTARGET_JIT_POST_OPT_IR_MODULE` environment variables.
+
+
+LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE
+""""""""""""""""""""""""""""""""""
+
+This environment variable can be used to extract the embedded device code
+before the device JIT runs additional IR optimizations on it (see
+:ref:`LIBOMPTARGET_JIT_OPT_LEVEL`). The value is expected to be a filename into
+which the LLVM-IR module is written. The module can be the analyzed, and
+transformed and loaded back into the JIT pipeline via
+:ref:`LIBOMPTARGET_JIT_REPLACEMENT_MODULE`.
+
+
+LIBOMPTARGET_JIT_POST_OPT_IR_MODULE
+""""""""""""""""""""""""""""""""""
+
+This environment variable can be used to extract the embedded device code after
+the device JIT runs additional IR optimizations on it (see
+:ref:`LIBOMPTARGET_JIT_OPT_LEVEL`). The value is expected to be a filename into
+which the LLVM-IR module is written. The module can be the analyzed, and
+transformed and loaded back into the JIT pipeline via
+:ref:`LIBOMPTARGET_JIT_REPLACEMENT_MODULE`.
+
+
+
 .. _libomptarget_plugin:
 
 LLVM/OpenMP Target Host Runtime Plugins (``libomptarget.rtl.XXXX``)
 -------------------------------------------------------------------
 
 .. _device_runtime:
 
 

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 "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/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 <system_error>
 
 using namespace llvm;
 using namespace llvm::object;
 using namespace omp;
 
 static codegen::RegisterCodeGenFlags RCGF;
 
 namespace {
#endif
   initializeWriteBitcodePassPass(Registry);
   initializeHardwareLoopsPass(Registry);
   initializeTypePromotionLegacyPass(Registry);
   initializeReplaceWithVeclibLegacyPass(Registry);
   initializeJMCInstrumenterPass(Registry);
 }
 
 Expected<std::unique_ptr<Module>>
-createModuleFromImage(__tgt_device_image *Image, LLVMContext &Context) {
-  StringRef Data((const char *)Image->ImageStart,
-                 (char *)Image->ImageEnd - (char *)Image->ImageStart);
-  std::unique_ptr<MemoryBuffer> MB = MemoryBuffer::getMemBuffer(
-      Data, /* BufferName */ "", /* RequiresNullTerminator */ false);
+createModuleFromMemoryBuffer(std::unique_ptr<MemoryBuffer> &MB,
+                             LLVMContext &Context) {
   SMDiagnostic Err;
   auto Mod = parseIR(*MB, Err, Context);
   if (!Mod)
     return make_error<StringError>("Failed to create module",
                                    inconvertibleErrorCode());
   return Mod;
 }
+Expected<std::unique_ptr<Module>>
+createModuleFromImage(__tgt_device_image *Image, LLVMContext &Context) {
+  StringRef Data((const char *)Image->ImageStart,
+                 (char *)Image->ImageEnd - (char *)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:
     return CodeGenOpt::None;
   case 1:
     return CodeGenOpt::Less;
   case 2:
return make_error<StringError>("Failed to create target machine",
                                    inconvertibleErrorCode());
   return TM;
 }
 
 ///
 class JITEngine {
 public:
   JITEngine(Triple::ArchType TA, std::string MCpu)
-      : TT(Triple::getArchTypeName(TA)), CPU(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);
   }
 
   /// 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;
Expected<std::unique_ptr<MemoryBuffer>> JITEngine::backend(Module &M,
 
   auto TMOrErr = createTargetMachine(M, CPU, OptLevel);
   if (!TMOrErr)
     return TMOrErr.takeError();
 
   std::unique_ptr<TargetMachine> TM = std::move(*TMOrErr);
   TargetLibraryInfoImpl TLII(TT);
 
+  if (PreOptIRModuleFileName.isPresent()) {
+    std::error_code EC;
+    raw_fd_stream FD(PreOptIRModuleFileName.get(), EC);
+    M.print(FD, nullptr);
+  }
+
   opt(TM.get(), &TLII, M, OptLevel);
 
+  if (PostOptIRModuleFileName.isPresent()) {
+    std::error_code EC;
+    raw_fd_stream FD(PostOptIRModuleFileName.get(), EC);
+    M.print(FD, nullptr);
+  }
+
   // Prepare the output buffer and stream for codegen.
   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) {
+  Module *Mod = nullptr;
+  // Check if the user replaces the module at runtime or we read it from the
+  // image.
+  if (!ReplacementModuleFileName.isPresent()) {
     auto ModOrErr = createModuleFromImage(Image, Context);
     if (!ModOrErr)
       return ModOrErr.takeError();
-
-  auto Mod = std::move(*ModOrErr);
+    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);
+    if (!ModOrErr)
+      return ModOrErr.takeError();
+    Mod = ModOrErr->release();
+  }
 
   auto MBOrError = backend(*Mod, OptLevel);
   if (!MBOrError)
     return MBOrError.takeError();
 
   return PostProcessing(std::move(*MBOrError));
 }
 

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

 Expected<__tgt_target_table *>
 GenericDeviceTy::loadBinary(GenericPluginTy &Plugin,
                             const __tgt_device_image *InputTgtImage) {
   __tgt_device_image *TgtImage = nullptr;
   bool isBitcodeImage = false;
   // 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 = getArch();
 
     jit::PostProcessingFn PostProcessing =
         [this](std::unique_ptr<MemoryBuffer> MB)
         -> Expected<std::unique_ptr<MemoryBuffer>> {
       return doJITPostProcessing(std::move(MB));