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