From d9415cd0241ebd7d426d0827ee0c65649018f90b Mon Sep 17 00:00:00 2001 From: Johannes Doerfert Date: Sun, 15 Jan 2023 11:28:03 -0800 Subject: [PATCH] [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 | 18 +++++ .../common/PluginInterface/JIT.cpp | 73 ++++++++++++------- .../common/PluginInterface/JIT.h | 11 +++ 3 files changed, 75 insertions(+), 27 deletions(-) diff --git a/openmp/docs/design/Runtimes.rst b/openmp/docs/design/Runtimes.rst index 4697780c207d..2db14aaeddb4 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=`` * ``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= (object file)`` * ``LIBOMPTARGET_JIT_REPLACEMENT_MODULE= (LLVM-IR file)`` * ``LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE= (LLVM-IR file)`` * ``LIBOMPTARGET_JIT_POST_OPT_IR_MODULE= (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 4382135ff675..b7551279277e 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,6 +297,46 @@ JITEngine::backend(Module &M, const std::string &ComputeUnitKind, return MemoryBuffer::getMemBufferCopy(OS.str()); } +Expected> +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 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(); diff --git a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.h b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/JIT.h index 0c51810690c8..9eb45a387c67 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 @@ 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> + getOrCreateObjFile(const __tgt_device_image &Image, LLVMContext &Ctx, + const std::string &ComputeUnitKind); + /// Run backend, which contains optimization and code generation. Expected> 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