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