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