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 // SystemZ is not supported yet. 47 if (Triple.isSystemZ()) 48 GTEST_SKIP(); 49 50 // 32-bit X86 is not supported yet. 51 if (Triple.isX86() && Triple.isArch32Bit()) 52 GTEST_SKIP(); 53 54 if (Triple.isPPC()) 55 GTEST_SKIP(); 56 57 auto EPC = SelfExecutorProcessControl::Create(); 58 if (!EPC) { 59 consumeError(EPC.takeError()); 60 GTEST_SKIP(); 61 } 62 63 auto DLOrErr = JTMB->getDefaultDataLayoutForTarget(); 64 if (!DLOrErr) { 65 consumeError(DLOrErr.takeError()); 66 GTEST_SKIP(); 67 } 68 ES = std::make_unique<ExecutionSession>(std::move(*EPC)); 69 JD = &ES->createBareJITDylib("main"); 70 ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>( 71 *ES, std::make_unique<InProcessMemoryManager>(16384)); 72 DL = std::make_unique<DataLayout>(std::move(*DLOrErr)); 73 74 auto TM = JTMB->createTargetMachine(); 75 if (!TM) { 76 consumeError(TM.takeError()); 77 GTEST_SKIP(); 78 } 79 auto CompileFunction = 80 std::make_unique<TMOwningSimpleCompiler>(std::move(*TM)); 81 CompileLayer = std::make_unique<IRCompileLayer>(*ES, *ObjLinkingLayer, 82 std::move(CompileFunction)); 83 } 84 85 Error addIRModule(ResourceTrackerSP RT, ThreadSafeModule TSM) { 86 assert(TSM && "Can not add null module"); 87 88 TSM.withModuleDo([&](Module &M) { M.setDataLayout(*DL); }); 89 90 return ROLayer->add(std::move(RT), std::move(TSM)); 91 } 92 93 JITDylib *JD{nullptr}; 94 std::unique_ptr<ExecutionSession> ES; 95 std::unique_ptr<ObjectLinkingLayer> ObjLinkingLayer; 96 std::unique_ptr<IRCompileLayer> CompileLayer; 97 std::unique_ptr<ReOptimizeLayer> ROLayer; 98 std::unique_ptr<DataLayout> DL; 99 }; 100 101 static Function *createRetFunction(Module *M, StringRef Name, 102 uint32_t ReturnCode) { 103 Function *Result = Function::Create( 104 FunctionType::get(Type::getInt32Ty(M->getContext()), {}, false), 105 GlobalValue::ExternalLinkage, Name, M); 106 107 BasicBlock *BB = BasicBlock::Create(M->getContext(), Name, Result); 108 IRBuilder<> Builder(M->getContext()); 109 Builder.SetInsertPoint(BB); 110 111 Value *RetValue = ConstantInt::get(M->getContext(), APInt(32, ReturnCode)); 112 Builder.CreateRet(RetValue); 113 return Result; 114 } 115 116 TEST_F(ReOptimizeLayerTest, BasicReOptimization) { 117 MangleAndInterner Mangle(*ES, *DL); 118 119 auto &EPC = ES->getExecutorProcessControl(); 120 EXPECT_THAT_ERROR(JD->define(absoluteSymbols( 121 {{Mangle("__orc_rt_jit_dispatch"), 122 {EPC.getJITDispatchInfo().JITDispatchFunction, 123 JITSymbolFlags::Exported}}, 124 {Mangle("__orc_rt_jit_dispatch_ctx"), 125 {EPC.getJITDispatchInfo().JITDispatchContext, 126 JITSymbolFlags::Exported}}, 127 {Mangle("__orc_rt_reoptimize_tag"), 128 {ExecutorAddr(), JITSymbolFlags::Exported}}})), 129 Succeeded()); 130 131 auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer, *JD); 132 EXPECT_THAT_ERROR(RM.takeError(), Succeeded()); 133 134 ROLayer = std::make_unique<ReOptimizeLayer>(*ES, *DL, *CompileLayer, **RM); 135 ROLayer->setReoptimizeFunc( 136 [&](ReOptimizeLayer &Parent, 137 ReOptimizeLayer::ReOptMaterializationUnitID MUID, unsigned CurVerison, 138 ResourceTrackerSP OldRT, ThreadSafeModule &TSM) { 139 TSM.withModuleDo([&](Module &M) { 140 for (auto &F : M) { 141 if (F.isDeclaration()) 142 continue; 143 for (auto &B : F) { 144 for (auto &I : B) { 145 if (ReturnInst *Ret = dyn_cast<ReturnInst>(&I)) { 146 Value *RetValue = 147 ConstantInt::get(M.getContext(), APInt(32, 53)); 148 Ret->setOperand(0, RetValue); 149 } 150 } 151 } 152 } 153 }); 154 return Error::success(); 155 }); 156 EXPECT_THAT_ERROR(ROLayer->reigsterRuntimeFunctions(*JD), Succeeded()); 157 158 ThreadSafeContext Ctx(std::make_unique<LLVMContext>()); 159 auto M = std::make_unique<Module>("<main>", *Ctx.getContext()); 160 M->setTargetTriple(sys::getProcessTriple()); 161 162 (void)createRetFunction(M.get(), "main", 42); 163 164 EXPECT_THAT_ERROR(addIRModule(JD->getDefaultResourceTracker(), 165 ThreadSafeModule(std::move(M), std::move(Ctx))), 166 Succeeded()); 167 168 auto Result = cantFail(ES->lookup({JD}, Mangle("main"))); 169 auto FuncPtr = Result.getAddress().toPtr<int (*)()>(); 170 for (size_t I = 0; I <= ReOptimizeLayer::CallCountThreshold; I++) 171 EXPECT_EQ(FuncPtr(), 42); 172 EXPECT_EQ(FuncPtr(), 53); 173 } 174