[Openmp-commits] [openmp] d9415cd - [OpenMP][JIT] Introduce more debugging configuration options

Johannes Doerfert via Openmp-commits openmp-commits at lists.llvm.org
Sun Jan 15 11:44:38 PST 2023


Author: Johannes Doerfert
Date: 2023-01-15T11:44:10-08:00
New Revision: d9415cd0241ebd7d426d0827ee0c65649018f90b

URL: https://github.com/llvm/llvm-project/commit/d9415cd0241ebd7d426d0827ee0c65649018f90b
DIFF: https://github.com/llvm/llvm-project/commit/d9415cd0241ebd7d426d0827ee0c65649018f90b.diff

LOG: [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.

Added: 
    

Modified: 
    openmp/docs/design/Runtimes.rst
    openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.cpp
    openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.h

Removed: 
    


################################################################################
diff  --git a/openmp/docs/design/Runtimes.rst b/openmp/docs/design/Runtimes.rst
index 4697780c207de..2db14aaeddb44 100644
--- a/openmp/docs/design/Runtimes.rst
+++ b/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
 """""""""""""""""""""""""""""""""""

diff  --git a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.cpp b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.cpp
index 4382135ff6758..b7551279277e3 100644
--- a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.cpp
+++ b/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();
 

diff  --git a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.h b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.h
index 0c51810690c8f..9eb45a387c67a 100644
--- a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.h
+++ b/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


        


More information about the Openmp-commits mailing list