[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,6 +297,46 @@ JITEngine::backend(Module &M, const std::string &ComputeUnitKind,
   return MemoryBuffer::getMemBufferCopy(OS.str());
 }
 
+Expected<std::unique_ptr<MemoryBuffer>>
+JITEngine::getOrCreateObjFile(const __tgt_device_image &Image, LLVMContext &Ctx,
+                              const std::string &ComputeUnitKind) {
+
+  // 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, Ctx);
+    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(), Ctx);
+    if (!ModOrErr)
+      return ModOrErr.takeError();
+    Mod = ModOrErr->release();
+  }
+
+  return backend(*Mod, ComputeUnitKind, JITOptLevel);
+}
+
 Expected<const __tgt_device_image *>
 JITEngine::compile(const __tgt_device_image &Image,
                    const std::string &ComputeUnitKind,
@@ -307,33 +348,11 @@ JITEngine::compile(const __tgt_device_image &Image,
   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, 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(), CUI.Context);
-    if (!ModOrErr)
-      return ModOrErr.takeError();
-    Mod = ModOrErr->release();
-  }
+  auto ObjMBOrErr = getOrCreateObjFile(Image, CUI.Context, ComputeUnitKind);
+  if (!ObjMBOrErr)
+    return ObjMBOrErr.takeError();
 
-  auto MBOrError = backend(*Mod, ComputeUnitKind, JITOptLevel);
-  if (!MBOrError)
-    return MBOrError.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 @@ private:
   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 @@ private:
   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 @@ private:
       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