unittests/ExecutionEngine/MCJIT/MCJITObjectCacheTest.cpp

   1 //===- MCJITObjectCacheTest.cpp - Unit tests for MCJIT object caching -----===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9
  10 #include "MCJITTestBase.h"
  11 #include "llvm/ADT/SmallVector.h"
  12 #include "llvm/ADT/StringMap.h"
  13 #include "llvm/ADT/StringSet.h"
  14 #include "llvm/ExecutionEngine/MCJIT.h"
  15 #include "llvm/ExecutionEngine/ObjectCache.h"
  16 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
  17 #include "gtest/gtest.h"
  18
  19 using namespace llvm;
  20
  21 namespace {
  22
  23 class TestObjectCache : public ObjectCache {
  24 public:
  25   TestObjectCache() : DuplicateInserted(false) { }
  26
  27   void notifyObjectCompiled(const Module *M, MemoryBufferRef Obj) override {
  28     // If we've seen this module before, note that.
  29     const std::string ModuleID = M->getModuleIdentifier();
  30     if (ObjMap.find(ModuleID) != ObjMap.end())
  31       DuplicateInserted = true;
  32     // Store a copy of the buffer in our map.
  33     ObjMap[ModuleID] = copyBuffer(Obj);
  34   }
  35
  36   std::unique_ptr<MemoryBuffer> getObject(const Module *M) override {
  37     const MemoryBuffer* BufferFound = getObjectInternal(M);
  38     ModulesLookedUp.insert(M->getModuleIdentifier());
  39     if (!BufferFound)
  40       return nullptr;
  41     // Our test cache wants to maintain ownership of its object buffers
  42     // so we make a copy here for the execution engine.
  43     return MemoryBuffer::getMemBufferCopy(BufferFound->getBuffer());
  44   }
  45
  46   // Test-harness-specific functions
  47   bool wereDuplicatesInserted() { return DuplicateInserted; }
  48
  49   bool wasModuleLookedUp(const Module *M) {
  50     return ModulesLookedUp.find(M->getModuleIdentifier())
  51                                       != ModulesLookedUp.end();
  52   }
  53
  54   const MemoryBuffer* getObjectInternal(const Module* M) {
  55     // Look for the module in our map.
  56     const std::string ModuleID = M->getModuleIdentifier();
  57     StringMap<const MemoryBuffer *>::iterator it = ObjMap.find(ModuleID);
  58     if (it == ObjMap.end())
  59       return nullptr;
  60     return it->second;
  61   }
  62
  63 private:
  64   MemoryBuffer *copyBuffer(MemoryBufferRef Buf) {
  65     // Create a local copy of the buffer.
  66     std::unique_ptr<MemoryBuffer> NewBuffer =
  67         MemoryBuffer::getMemBufferCopy(Buf.getBuffer());
  68     MemoryBuffer *Ret = NewBuffer.get();
  69     AllocatedBuffers.push_back(std::move(NewBuffer));
  70     return Ret;
  71   }
  72
  73   StringMap<const MemoryBuffer *> ObjMap;
  74   StringSet<>                     ModulesLookedUp;
  75   SmallVector<std::unique_ptr<MemoryBuffer>, 2> AllocatedBuffers;
  76   bool                            DuplicateInserted;
  77 };
  78
  79 class MCJITObjectCacheTest : public testing::Test, public MCJITTestBase {
  80 protected:
  81   enum {
  82     OriginalRC = 6,
  83     ReplacementRC = 7
  84   };
  85
  86   void SetUp() override {
  87     M.reset(createEmptyModule("<main>"));
  88     Main = insertMainFunction(M.get(), OriginalRC);
  89   }
  90
  91   void compileAndRun(int ExpectedRC = OriginalRC) {
  92     // This function shouldn't be called until after SetUp.
  93     ASSERT_TRUE(bool(TheJIT));
  94     ASSERT_TRUE(nullptr != Main);
  95
  96     // We may be using a null cache, so ensure compilation is valid.
  97     TheJIT->finalizeObject();
  98     void *vPtr = TheJIT->getPointerToFunction(Main);
  99
 100     EXPECT_TRUE(nullptr != vPtr)
 101       << "Unable to get pointer to main() from JIT";
 102
 103     int (*FuncPtr)() = (int(*)())(intptr_t)vPtr;
 104     int returnCode = FuncPtr();
 105     EXPECT_EQ(returnCode, ExpectedRC);
 106   }
 107
 108   Function *Main;
 109 };
 110
 111 TEST_F(MCJITObjectCacheTest, SetNullObjectCache) {
 112   SKIP_UNSUPPORTED_PLATFORM;
 113
 114   createJIT(std::move(M));
 115
 116   TheJIT->setObjectCache(nullptr);
 117
 118   compileAndRun();
 119 }
 120
 121 TEST_F(MCJITObjectCacheTest, VerifyBasicObjectCaching) {
 122   SKIP_UNSUPPORTED_PLATFORM;
 123
 124   std::unique_ptr<TestObjectCache> Cache(new TestObjectCache);
 125
 126   // Save a copy of the module pointer before handing it off to MCJIT.
 127   const Module * SavedModulePointer = M.get();
 128
 129   createJIT(std::move(M));
 130
 131   TheJIT->setObjectCache(Cache.get());
 132
 133   // Verify that our object cache does not contain the module yet.
 134   const MemoryBuffer *ObjBuffer = Cache->getObjectInternal(SavedModulePointer);
 135   EXPECT_EQ(nullptr, ObjBuffer);
 136
 137   compileAndRun();
 138
 139   // Verify that MCJIT tried to look-up this module in the cache.
 140   EXPECT_TRUE(Cache->wasModuleLookedUp(SavedModulePointer));
 141
 142   // Verify that our object cache now contains the module.
 143   ObjBuffer = Cache->getObjectInternal(SavedModulePointer);
 144   EXPECT_TRUE(nullptr != ObjBuffer);
 145
 146   // Verify that the cache was only notified once.
 147   EXPECT_FALSE(Cache->wereDuplicatesInserted());
 148 }
 149
 150 TEST_F(MCJITObjectCacheTest, VerifyLoadFromCache) {
 151   SKIP_UNSUPPORTED_PLATFORM;
 152
 153   std::unique_ptr<TestObjectCache> Cache(new TestObjectCache);
 154
 155   // Compile this module with an MCJIT engine
 156   createJIT(std::move(M));
 157   TheJIT->setObjectCache(Cache.get());
 158   TheJIT->finalizeObject();
 159
 160   // Destroy the MCJIT engine we just used
 161   TheJIT.reset();
 162
 163   // Create a new memory manager.
 164   MM.reset(new SectionMemoryManager());
 165
 166   // Create a new module and save it. Use a different return code so we can
 167   // tell if MCJIT compiled this module or used the cache.
 168   M.reset(createEmptyModule("<main>"));
 169   Main = insertMainFunction(M.get(), ReplacementRC);
 170   const Module * SecondModulePointer = M.get();
 171
 172   // Create a new MCJIT instance to load this module then execute it.
 173   createJIT(std::move(M));
 174   TheJIT->setObjectCache(Cache.get());
 175   compileAndRun();
 176
 177   // Verify that MCJIT tried to look-up this module in the cache.
 178   EXPECT_TRUE(Cache->wasModuleLookedUp(SecondModulePointer));
 179
 180   // Verify that MCJIT didn't try to cache this again.
 181   EXPECT_FALSE(Cache->wereDuplicatesInserted());
 182 }
 183
 184 TEST_F(MCJITObjectCacheTest, VerifyNonLoadFromCache) {
 185   SKIP_UNSUPPORTED_PLATFORM;
 186
 187   std::unique_ptr<TestObjectCache> Cache(new TestObjectCache);
 188
 189   // Compile this module with an MCJIT engine
 190   createJIT(std::move(M));
 191   TheJIT->setObjectCache(Cache.get());
 192   TheJIT->finalizeObject();
 193
 194   // Destroy the MCJIT engine we just used
 195   TheJIT.reset();
 196
 197   // Create a new memory manager.
 198   MM.reset(new SectionMemoryManager());
 199
 200   // Create a new module and save it. Use a different return code so we can
 201   // tell if MCJIT compiled this module or used the cache. Note that we use
 202   // a new module name here so the module shouldn't be found in the cache.
 203   M.reset(createEmptyModule("<not-main>"));
 204   Main = insertMainFunction(M.get(), ReplacementRC);
 205   const Module * SecondModulePointer = M.get();
 206
 207   // Create a new MCJIT instance to load this module then execute it.
 208   createJIT(std::move(M));
 209   TheJIT->setObjectCache(Cache.get());
 210
 211   // Verify that our object cache does not contain the module yet.
 212   const MemoryBuffer *ObjBuffer = Cache->getObjectInternal(SecondModulePointer);
 213   EXPECT_EQ(nullptr, ObjBuffer);
 214
 215   // Run the function and look for the replacement return code.
 216   compileAndRun(ReplacementRC);
 217
 218   // Verify that MCJIT tried to look-up this module in the cache.
 219   EXPECT_TRUE(Cache->wasModuleLookedUp(SecondModulePointer));
 220
 221   // Verify that our object cache now contains the module.
 222   ObjBuffer = Cache->getObjectInternal(SecondModulePointer);
 223   EXPECT_TRUE(nullptr != ObjBuffer);
 224
 225   // Verify that MCJIT didn't try to cache this again.
 226   EXPECT_FALSE(Cache->wereDuplicatesInserted());
 227 }
 228
 229 } // end anonymous namespace