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