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

   1 //===- JIT.h - Target independent JIT infrastructure ----------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 //===----------------------------------------------------------------------===//
  10
  11 #ifndef OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H
  12 #define OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H
  13
  14 #include "Utilities.h"
  15
  16 #include "llvm/ADT/StringMap.h"
  17 #include "llvm/ADT/StringRef.h"
  18 #include "llvm/Analysis/TargetLibraryInfo.h"
  19 #include "llvm/IR/LLVMContext.h"
  20 #include "llvm/IR/Module.h"
  21 #include "llvm/Support/Error.h"
  22 #include "llvm/Target/TargetMachine.h"
  23 #include "llvm/TargetParser/Triple.h"
  24
  25 #include <functional>
  26 #include <memory>
  27 #include <shared_mutex>
  28 #include <string>
  29
  30 struct __tgt_device_image;
  31
  32 namespace llvm {
  33 class MemoryBuffer;
  34
  35 namespace omp {
  36 namespace target {
  37 namespace plugin {
  38 struct GenericDeviceTy;
  39 } // namespace plugin
  40
  41 /// The JIT infrastructure and caching mechanism.
  42 struct JITEngine {
  43   /// Function type for a callback that will be called after the backend is
  44   /// called.
  45   using PostProcessingFn =
  46       std::function<Expected<std::unique_ptr<MemoryBuffer>>(
  47           std::unique_ptr<MemoryBuffer>)>;
  48
  49   JITEngine(Triple::ArchType TA);
  50
  51   /// Run jit compilation if \p Image is a bitcode image, otherwise simply
  52   /// return \p Image. It is expected to return a memory buffer containing the
  53   /// generated device image that could be loaded to the device directly.
  54   Expected<const __tgt_device_image *>
  55   process(const __tgt_device_image &Image,
  56           target::plugin::GenericDeviceTy &Device);
  57
  58   /// Return true if \p Image is a bitcode image that can be JITed for the given
  59   /// architecture.
  60   bool checkBitcodeImage(const __tgt_device_image &Image);
  61
  62 private:
  63   /// Compile the bitcode image \p Image and generate the binary image that can
  64   /// be loaded to the target device of the triple \p Triple architecture \p
  65   /// MCpu. \p PostProcessing will be called after codegen to handle cases such
  66   /// as assember as an external tool.
  67   Expected<const __tgt_device_image *>
  68   compile(const __tgt_device_image &Image, const std::string &ComputeUnitKind,
  69           PostProcessingFn PostProcessing);
  70
  71   /// Create or retrieve the object image file from the file system or via
  72   /// compilation of the \p Image.
  73   Expected<std::unique_ptr<MemoryBuffer>>
  74   getOrCreateObjFile(const __tgt_device_image &Image, LLVMContext &Ctx,
  75                      const std::string &ComputeUnitKind);
  76
  77   /// Run backend, which contains optimization and code generation.
  78   Expected<std::unique_ptr<MemoryBuffer>>
  79   backend(Module &M, const std::string &ComputeUnitKind, unsigned OptLevel);
  80
  81   /// Run optimization pipeline.
  82   void opt(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
  83            unsigned OptLevel);
  84
  85   /// Run code generation.
  86   void codegen(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
  87                raw_pwrite_stream &OS);
  88
  89   /// The target triple used by the JIT.
  90   const Triple TT;
  91
  92   struct ComputeUnitInfo {
  93     /// LLVM Context in which the modules will be constructed.
  94     LLVMContext Context;
  95
  96     /// Output images generated from LLVM backend.
  97     SmallVector<std::unique_ptr<MemoryBuffer>, 4> JITImages;
  98
  99     /// A map of embedded IR images to JITed images.
 100     DenseMap<const __tgt_device_image *, __tgt_device_image *> TgtImageMap;
 101   };
 102
 103   /// Map from (march) "CPUs" (e.g., sm_80, or gfx90a), which we call compute
 104   /// units as they are not CPUs, to the image information we cached for them.
 105   StringMap<ComputeUnitInfo> ComputeUnitMap;
 106   std::mutex ComputeUnitMapMutex;
 107
 108   /// Control environment variables.
 109   target::StringEnvar ReplacementObjectFileName =
 110       target::StringEnvar("LIBOMPTARGET_JIT_REPLACEMENT_OBJECT");
 111   target::StringEnvar ReplacementModuleFileName =
 112       target::StringEnvar("LIBOMPTARGET_JIT_REPLACEMENT_MODULE");
 113   target::StringEnvar PreOptIRModuleFileName =
 114       target::StringEnvar("LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE");
 115   target::StringEnvar PostOptIRModuleFileName =
 116       target::StringEnvar("LIBOMPTARGET_JIT_POST_OPT_IR_MODULE");
 117   target::UInt32Envar JITOptLevel =
 118       target::UInt32Envar("LIBOMPTARGET_JIT_OPT_LEVEL", 3);
 119   target::BoolEnvar JITSkipOpt =
 120       target::BoolEnvar("LIBOMPTARGET_JIT_SKIP_OPT", false);
 121 };
 122
 123 } // namespace target
 124 } // namespace omp
 125 } // namespace llvm
 126
 127 #endif // OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H