195 lines · cpp
1#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"2#include "OrcTestCommon.h"3#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"4#include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h"5#include "llvm/ExecutionEngine/Orc/CompileUtils.h"6#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"7#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"8#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"9#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"10#include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"11#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"12#include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"13#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"14#include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"15#include "llvm/ExecutionEngine/Orc/SelfExecutorProcessControl.h"16#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"17#include "llvm/IR/IRBuilder.h"18#include "llvm/Support/CodeGen.h"19#include "llvm/TargetParser/Host.h"20#include "llvm/Testing/Support/Error.h"21#include "gtest/gtest.h"22 23using namespace llvm;24using namespace llvm::orc;25using namespace llvm::jitlink;26 27class ReOptimizeLayerTest : public testing::Test {28public:29 ~ReOptimizeLayerTest() override {30 if (ES)31 if (auto Err = ES->endSession())32 ES->reportError(std::move(Err));33 }34 35protected:36 void SetUp() override {37 auto JTMB = JITTargetMachineBuilder::detectHost();38 // Bail out if we can not detect the host.39 if (!JTMB) {40 consumeError(JTMB.takeError());41 GTEST_SKIP();42 }43 44 // COFF-ARM64 is not supported yet45 auto Triple = JTMB->getTargetTriple();46 if (Triple.isOSBinFormatCOFF())47 GTEST_SKIP();48 49 // SystemZ is not supported yet.50 if (Triple.isSystemZ())51 GTEST_SKIP();52 53 // 32-bit X86 is not supported yet.54 if (Triple.isX86() && Triple.isArch32Bit())55 GTEST_SKIP();56 57 if (Triple.isPPC())58 GTEST_SKIP();59 60 // RISC-V is not supported yet61 if (Triple.isRISCV())62 GTEST_SKIP();63 64 // ARM is not supported yet.65 if (Triple.isARM())66 GTEST_SKIP();67 68 auto EPC = SelfExecutorProcessControl::Create();69 if (!EPC) {70 consumeError(EPC.takeError());71 GTEST_SKIP();72 }73 74 auto DLOrErr = JTMB->getDefaultDataLayoutForTarget();75 if (!DLOrErr) {76 consumeError(DLOrErr.takeError());77 GTEST_SKIP();78 }79 80 auto PageSize = sys::Process::getPageSize();81 if (!PageSize) {82 consumeError(PageSize.takeError());83 GTEST_SKIP();84 }85 86 ES = std::make_unique<ExecutionSession>(std::move(*EPC));87 JD = &ES->createBareJITDylib("main");88 89 ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(90 *ES, std::make_unique<MapperJITLinkMemoryManager>(91 10 * 1024 * 1024,92 std::make_unique<InProcessMemoryMapper>(*PageSize)));93 DL = std::make_unique<DataLayout>(std::move(*DLOrErr));94 95 auto TM = JTMB->createTargetMachine();96 if (!TM) {97 consumeError(TM.takeError());98 GTEST_SKIP();99 }100 auto CompileFunction =101 std::make_unique<TMOwningSimpleCompiler>(std::move(*TM));102 CompileLayer = std::make_unique<IRCompileLayer>(*ES, *ObjLinkingLayer,103 std::move(CompileFunction));104 }105 106 Error addIRModule(ResourceTrackerSP RT, ThreadSafeModule TSM) {107 assert(TSM && "Can not add null module");108 109 TSM.withModuleDo([&](Module &M) { M.setDataLayout(*DL); });110 111 return ROLayer->add(std::move(RT), std::move(TSM));112 }113 114 JITDylib *JD{nullptr};115 std::unique_ptr<ExecutionSession> ES;116 std::unique_ptr<ObjectLinkingLayer> ObjLinkingLayer;117 std::unique_ptr<IRCompileLayer> CompileLayer;118 std::unique_ptr<ReOptimizeLayer> ROLayer;119 std::unique_ptr<DataLayout> DL;120};121 122static Function *createRetFunction(Module *M, StringRef Name,123 uint32_t ReturnCode) {124 Function *Result = Function::Create(125 FunctionType::get(Type::getInt32Ty(M->getContext()), {}, false),126 GlobalValue::ExternalLinkage, Name, M);127 128 BasicBlock *BB = BasicBlock::Create(M->getContext(), Name, Result);129 IRBuilder<> Builder(M->getContext());130 Builder.SetInsertPoint(BB);131 132 Value *RetValue = ConstantInt::get(M->getContext(), APInt(32, ReturnCode));133 Builder.CreateRet(RetValue);134 return Result;135}136 137TEST_F(ReOptimizeLayerTest, BasicReOptimization) {138 MangleAndInterner Mangle(*ES, *DL);139 140 auto &EPC = ES->getExecutorProcessControl();141 EXPECT_THAT_ERROR(JD->define(absoluteSymbols(142 {{Mangle("__orc_rt_jit_dispatch"),143 {EPC.getJITDispatchInfo().JITDispatchFunction,144 JITSymbolFlags::Exported}},145 {Mangle("__orc_rt_jit_dispatch_ctx"),146 {EPC.getJITDispatchInfo().JITDispatchContext,147 JITSymbolFlags::Exported}},148 {Mangle("__orc_rt_reoptimize_tag"),149 {ExecutorAddr(), JITSymbolFlags::Exported}}})),150 Succeeded());151 152 auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer);153 EXPECT_THAT_ERROR(RM.takeError(), Succeeded());154 155 ROLayer = std::make_unique<ReOptimizeLayer>(*ES, *DL, *CompileLayer, **RM);156 ROLayer->setReoptimizeFunc(157 [&](ReOptimizeLayer &Parent,158 ReOptimizeLayer::ReOptMaterializationUnitID MUID, unsigned CurVerison,159 ResourceTrackerSP OldRT, ThreadSafeModule &TSM) {160 TSM.withModuleDo([&](Module &M) {161 for (auto &F : M) {162 if (F.isDeclaration())163 continue;164 for (auto &B : F) {165 for (auto &I : B) {166 if (ReturnInst *Ret = dyn_cast<ReturnInst>(&I)) {167 Value *RetValue =168 ConstantInt::get(M.getContext(), APInt(32, 53));169 Ret->setOperand(0, RetValue);170 }171 }172 }173 }174 });175 return Error::success();176 });177 EXPECT_THAT_ERROR(ROLayer->reigsterRuntimeFunctions(*JD), Succeeded());178 179 auto Ctx = std::make_unique<LLVMContext>();180 auto M = std::make_unique<Module>("<main>", *Ctx);181 M->setTargetTriple(Triple(sys::getProcessTriple()));182 183 (void)createRetFunction(M.get(), "main", 42);184 185 EXPECT_THAT_ERROR(addIRModule(JD->getDefaultResourceTracker(),186 ThreadSafeModule(std::move(M), std::move(Ctx))),187 Succeeded());188 189 auto Result = cantFail(ES->lookup({JD}, Mangle("main")));190 auto FuncPtr = Result.getAddress().toPtr<int (*)()>();191 for (size_t I = 0; I <= ReOptimizeLayer::CallCountThreshold; I++)192 EXPECT_EQ(FuncPtr(), 42);193 EXPECT_EQ(FuncPtr(), 53);194}195