llvm/unittests/ExecutionEngine/MCJIT/MCJITObjectCacheTest.cpp

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