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