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