[OpenMP][JIT] Introduce more debugging configuration options

The JIT is a great debugging tool since we can modify the IR manually
before launching it in an existing test case. The new flasks allow to
skip optimizations, to use the exact given IR, as well as to provide a
finished object file. The latter is useful to try out different backend
options and to have complete freedom with pass pipelines.

Documentation is included. Minimal refactoring was performed to make the
second object fit in nicely.

openmp/docs/design/Runtimes.rst

@@ -714,6 +714,8 @@ variables is defined below.
     * ``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_SKIP_OPT=[TRUE/FALSE] (default FALSE)``
+    * ``LIBOMPTARGET_JIT_REPLACEMENT_OBJECT=<in:Filename> (object file)``
     * ``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)``
@@ -1062,6 +1064,22 @@ 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_SKIP_OPT
+""""""""""""""""""""""""""
+
+This environment variable can be used to skip the optimization pipeline during
+JIT compilation. If set, the image will only be passed through the backend. The
+backend is invoked with the ``LIBOMPTARGET_JIT_OPT_LEVEL`` flag.
+
+LIBOMPTARGET_JIT_REPLACEMENT_OBJECT
+"""""""""""""""""""""""""""""""""""
+
+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 object file, thus containing the output of the
+assembler in object format for the respective target. The JIT optimization
+pipeline and backend are skipped and only target specific post-processing is
+performed on the object file before it is loaded onto the device.
 
 LIBOMPTARGET_JIT_REPLACEMENT_MODULE
 """""""""""""""""""""""""""""""""""

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

@@ -275,7 +275,8 @@ JITEngine::backend(Module &M, const std::string &ComputeUnitKind,
     M.print(FD, nullptr);
   }
 
-  opt(TM.get(), &TLII, M, OptLevel);
+  if (!JITSkipOpt)
+    opt(TM.get(), &TLII, M, OptLevel);
 
   if (PostOptIRModuleFileName.isPresent()) {
     std::error_code EC;
@@ -296,23 +297,27 @@ JITEngine::backend(Module &M, const std::string &ComputeUnitKind,
   return MemoryBuffer::getMemBufferCopy(OS.str());
 }
 
-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);
+Expected<std::unique_ptr<MemoryBuffer>>
+JITEngine::getOrCreateObjFile(const __tgt_device_image &Image, LLVMContext &Ctx,
+                              const std::string &ComputeUnitKind) {
 
-  // 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;
+  // Check if the user replaces the module at runtime with a finished object.
+  if (ReplacementObjectFileName.isPresent()) {
+    auto MBOrErr =
+        MemoryBuffer::getFileOrSTDIN(ReplacementObjectFileName.get());
+    if (!MBOrErr)
+      return createStringError(MBOrErr.getError(),
+                               "Could not read replacement obj from %s\n",
+                               ReplacementModuleFileName.get().c_str());
+    return std::move(*MBOrErr);
+  }
 
   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, CUI.Context);
+    auto ModOrErr = createModuleFromImage(Image, Ctx);
     if (!ModOrErr)
       return ModOrErr.takeError();
     Mod = ModOrErr->release();
@@ -323,17 +328,31 @@ JITEngine::compile(const __tgt_device_image &Image,
       return createStringError(MBOrErr.getError(),
                                "Could not read replacement module from %s\n",
                                ReplacementModuleFileName.get().c_str());
-    auto ModOrErr = createModuleFromMemoryBuffer(MBOrErr.get(), CUI.Context);
+    auto ModOrErr = createModuleFromMemoryBuffer(MBOrErr.get(), Ctx);
     if (!ModOrErr)
       return ModOrErr.takeError();
     Mod = ModOrErr->release();
   }
 
-  auto MBOrError = backend(*Mod, ComputeUnitKind, JITOptLevel);
-  if (!MBOrError)
-    return MBOrError.takeError();
+  return backend(*Mod, ComputeUnitKind, JITOptLevel);
+}
+
+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;
+
+  auto ObjMBOrErr = getOrCreateObjFile(Image, CUI.Context, ComputeUnitKind);
+  if (!ObjMBOrErr)
+    return ObjMBOrErr.takeError();
 
-  auto ImageMBOrErr = PostProcessing(std::move(*MBOrError));
+  auto ImageMBOrErr = PostProcessing(std::move(*ObjMBOrErr));
   if (!ImageMBOrErr)
     return ImageMBOrErr.takeError();
 

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

@@ -17,6 +17,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Target/TargetMachine.h"
@@ -67,6 +68,12 @@ struct JITEngine {
   compile(const __tgt_device_image &Image, const std::string &ComputeUnitKind,
           PostProcessingFn PostProcessing);
 
+  /// Create or retrieve the object image file from the file system or via
+  /// compilation of the \p Image.
+  Expected<std::unique_ptr<MemoryBuffer>>
+  getOrCreateObjFile(const __tgt_device_image &Image, LLVMContext &Ctx,
+                     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);
@@ -99,6 +106,8 @@ struct JITEngine {
   std::mutex ComputeUnitMapMutex;
 
   /// Control environment variables.
+  target::StringEnvar ReplacementObjectFileName =
+      target::StringEnvar("LIBOMPTARGET_JIT_REPLACEMENT_OBJECT");
   target::StringEnvar ReplacementModuleFileName =
       target::StringEnvar("LIBOMPTARGET_JIT_REPLACEMENT_MODULE");
   target::StringEnvar PreOptIRModuleFileName =
@@ -107,6 +116,8 @@ struct JITEngine {
       target::StringEnvar("LIBOMPTARGET_JIT_POST_OPT_IR_MODULE");
   target::UInt32Envar JITOptLevel =
       target::UInt32Envar("LIBOMPTARGET_JIT_OPT_LEVEL", 3);
+  target::BoolEnvar JITSkipOpt =
+      target::BoolEnvar("LIBOMPTARGET_JIT_SKIP_OPT", false);
 };
 
 } // namespace target