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 "Shared/EnvironmentVar.h" 15 #include "Shared/Utils.h" 16 17 #include "llvm/ADT/StringMap.h" 18 #include "llvm/ADT/StringRef.h" 19 #include "llvm/Analysis/TargetLibraryInfo.h" 20 #include "llvm/IR/LLVMContext.h" 21 #include "llvm/IR/Module.h" 22 #include "llvm/Support/Error.h" 23 #include "llvm/Target/TargetMachine.h" 24 #include "llvm/TargetParser/Triple.h" 25 26 #include <functional> 27 #include <memory> 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 private: 59 /// Compile the bitcode image \p Image and generate the binary image that can 60 /// be loaded to the target device of the triple \p Triple architecture \p 61 /// MCpu. \p PostProcessing will be called after codegen to handle cases such 62 /// as assember as an external tool. 63 Expected<const __tgt_device_image *> 64 compile(const __tgt_device_image &Image, const std::string &ComputeUnitKind, 65 PostProcessingFn PostProcessing); 66 67 /// Create or retrieve the object image file from the file system or via 68 /// compilation of the \p Image. 69 Expected<std::unique_ptr<MemoryBuffer>> 70 getOrCreateObjFile(const __tgt_device_image &Image, LLVMContext &Ctx, 71 const std::string &ComputeUnitKind); 72 73 /// Run backend, which contains optimization and code generation. 74 Expected<std::unique_ptr<MemoryBuffer>> 75 backend(Module &M, const std::string &ComputeUnitKind, unsigned OptLevel); 76 77 /// Run optimization pipeline. 78 void opt(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M, 79 unsigned OptLevel); 80 81 /// Run code generation. 82 void codegen(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M, 83 raw_pwrite_stream &OS); 84 85 /// The target triple used by the JIT. 86 const Triple TT; 87 88 struct ComputeUnitInfo { 89 /// LLVM Context in which the modules will be constructed. 90 LLVMContext Context; 91 92 /// Output images generated from LLVM backend. 93 SmallVector<std::unique_ptr<MemoryBuffer>, 4> JITImages; 94 95 /// A map of embedded IR images to JITed images. 96 DenseMap<const __tgt_device_image *, __tgt_device_image *> TgtImageMap; 97 }; 98 99 /// Map from (march) "CPUs" (e.g., sm_80, or gfx90a), which we call compute 100 /// units as they are not CPUs, to the image information we cached for them. 101 StringMap<ComputeUnitInfo> ComputeUnitMap; 102 std::mutex ComputeUnitMapMutex; 103 104 /// Control environment variables. 105 StringEnvar ReplacementObjectFileName = 106 StringEnvar("LIBOMPTARGET_JIT_REPLACEMENT_OBJECT"); 107 StringEnvar ReplacementModuleFileName = 108 StringEnvar("LIBOMPTARGET_JIT_REPLACEMENT_MODULE"); 109 StringEnvar PreOptIRModuleFileName = 110 StringEnvar("LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE"); 111 StringEnvar PostOptIRModuleFileName = 112 StringEnvar("LIBOMPTARGET_JIT_POST_OPT_IR_MODULE"); 113 UInt32Envar JITOptLevel = UInt32Envar("LIBOMPTARGET_JIT_OPT_LEVEL", 3); 114 BoolEnvar JITSkipOpt = BoolEnvar("LIBOMPTARGET_JIT_SKIP_OPT", false); 115 }; 116 117 } // namespace target 118 } // namespace omp 119 } // namespace llvm 120 121 #endif // OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H 122