115 lines · c
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-exception6//7//===----------------------------------------------------------------------===//8//9//===----------------------------------------------------------------------===//10 11#ifndef OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H12#define OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H13 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 30struct __tgt_device_image;31 32namespace llvm {33class MemoryBuffer;34 35namespace omp {36namespace target {37namespace plugin {38struct GenericDeviceTy;39} // namespace plugin40 41/// The JIT infrastructure and caching mechanism.42struct JITEngine {43 /// Function type for a callback that will be called after the backend is44 /// 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 simply52 /// return \p Image. It is expected to return a memory buffer containing the53 /// generated device image that could be loaded to the device directly.54 Expected<std::unique_ptr<MemoryBuffer>>55 process(StringRef Image, target::plugin::GenericDeviceTy &Device);56 57private:58 /// Compile the bitcode image \p Image and generate the binary image that can59 /// be loaded to the target device of the triple \p Triple architecture \p60 /// MCpu. \p PostProcessing will be called after codegen to handle cases such61 /// as assembler as an external tool.62 Expected<std::unique_ptr<MemoryBuffer>>63 compile(StringRef Image, const std::string &ComputeUnitKind,64 PostProcessingFn PostProcessing);65 66 /// Create or retrieve the object image file from the file system or via67 /// compilation of the \p Image.68 Expected<std::unique_ptr<MemoryBuffer>>69 getOrCreateObjFile(StringRef Image, LLVMContext &Ctx,70 const std::string &ComputeUnitKind);71 72 /// Run backend, which contains optimization and code generation.73 Expected<std::unique_ptr<MemoryBuffer>>74 backend(Module &M, const std::string &ComputeUnitKind, unsigned OptLevel);75 76 /// Run optimization pipeline.77 void opt(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,78 unsigned OptLevel);79 80 /// Run code generation.81 void codegen(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,82 raw_pwrite_stream &OS);83 84 /// The target triple used by the JIT.85 const Triple TT;86 87 struct ComputeUnitInfo {88 /// LLVM Context in which the modules will be constructed.89 LLVMContext Context;90 };91 92 /// Map from (march) "CPUs" (e.g., sm_80, or gfx90a), which we call compute93 /// units as they are not CPUs, to the image information we cached for them.94 StringMap<ComputeUnitInfo> ComputeUnitMap;95 std::mutex ComputeUnitMapMutex;96 97 /// Control environment variables.98 StringEnvar ReplacementObjectFileName =99 StringEnvar("LIBOMPTARGET_JIT_REPLACEMENT_OBJECT");100 StringEnvar ReplacementModuleFileName =101 StringEnvar("LIBOMPTARGET_JIT_REPLACEMENT_MODULE");102 StringEnvar PreOptIRModuleFileName =103 StringEnvar("LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE");104 StringEnvar PostOptIRModuleFileName =105 StringEnvar("LIBOMPTARGET_JIT_POST_OPT_IR_MODULE");106 UInt32Envar JITOptLevel = UInt32Envar("LIBOMPTARGET_JIT_OPT_LEVEL", 3);107 BoolEnvar JITSkipOpt = BoolEnvar("LIBOMPTARGET_JIT_SKIP_OPT", false);108};109 110} // namespace target111} // namespace omp112} // namespace llvm113 114#endif // OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_JIT_H115