llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.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/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 }