From: Lang Hames Date: Thu, 19 Feb 2015 01:32:43 +0000 (+0000) Subject: [Orc][Kaleidoscope] Make the 'fully lazy' orc kaleidoscope tutorial lazier still. X-Git-Url: http://plrg.eecs.uci.edu/git/?p=oota-llvm.git;a=commitdiff_plain;h=32f3cd95bb06dddf550730b69c332d5e31bcebc2 [Orc][Kaleidoscope] Make the 'fully lazy' orc kaleidoscope tutorial lazier still. The new JIT doesn't IRGen stubs until they're referenced. git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@229807 91177308-0d34-0410-b5e6-96231b3b80d8 --- diff --git a/examples/Kaleidoscope/Orc/fully_lazy/toy.cpp b/examples/Kaleidoscope/Orc/fully_lazy/toy.cpp index 9210dd1f3bb..92f4ba22982 100644 --- a/examples/Kaleidoscope/Orc/fully_lazy/toy.cpp +++ b/examples/Kaleidoscope/Orc/fully_lazy/toy.cpp @@ -1184,27 +1184,14 @@ public: // We need a memory manager to allocate memory and resolve symbols for this // new module. Create one that resolves symbols by looking back into the JIT. auto MM = createLookasideRTDyldMM( - [&](const std::string &Name) -> uint64_t { + [&](const std::string &Name) { // First try to find 'Name' within the JIT. if (auto Symbol = findMangledSymbol(Name)) return Symbol.getAddress(); - // If we don't find 'Name' in the JIT, see if we have some AST - // for it. - auto DefI = Session.FunctionDefs.find(Name); - if (DefI == Session.FunctionDefs.end()) - return 0; - - // We have AST for 'Name'. IRGen it, add it to the JIT, and - // return the address for it. - // FIXME: What happens if IRGen fails? - addModule(IRGen(Session, *DefI->second)); - - // Remove the function definition's AST now that we've - // finished with it. - Session.FunctionDefs.erase(DefI); - - return findMangledSymbol(Name).getAddress(); + // If we don't already have a definition of 'Name' then search + // the ASTs. + return searchUncompiledASTs(Name); }, [](const std::string &S) { return 0; } ); @@ -1223,7 +1210,7 @@ public: JITSymbol findMangledSymbolIn(LazyEmitLayerT::ModuleSetHandleT H, const std::string &Name) { - return LazyEmitLayer.findSymbolIn(H, Name, true); + return LazyEmitLayer.findSymbolIn(H, Name, true); } JITSymbol findSymbolIn(LazyEmitLayerT::ModuleSetHandleT H, @@ -1232,19 +1219,49 @@ public: } void addFunctionDefinition(std::unique_ptr FnAST) { - // Step 1) IRGen a prototype for this function: + FunctionDefs[Mangle(FnAST->Proto->Name)] = std::move(FnAST); + } + +private: + + // This method searches the FunctionDefs map for a definition of 'Name'. If it + // finds one it generates a stub for it and returns the address of the stub. + TargetAddress searchUncompiledASTs(const std::string &Name) { + auto DefI = FunctionDefs.find(Name); + if (DefI == FunctionDefs.end()) + return 0; + + // We have AST for 'Name'. IRGen a stub for it and add it to the JIT. + // FIXME: What happens if IRGen fails? + auto H = irGenStub(std::move(DefI->second)); + + // Remove the map entry now that we're done with it. + Session.FunctionDefs.erase(DefI); + + // Return the address of the stub. + return findMangledSymbolIn(H, Name).getAddress(); + } + + // This method will take the AST for a function definition and IR-gen a stub + // for that function that will, on first call, IR-gen the action body of the + // function. + ModuleHandleT irGenStub(std::unique_ptr FnAST) { + // Step 1) IRGen a prototype for the stub. This will have the same type as + // the function. IRGenContext C(Session); Function *F = FnAST->Proto->IRGen(C); C.getM().setDataLayout(TM->getDataLayout()); - // Step 2) Create a compile callback that will be used to compile this - // function when it is first called. + // Step 2) Get a compile callback that can be used to compile the body of + // the function. The resulting CallbackInfo type will let us set the + // compile and update actions for the callback, and get a pointer to + // the jit trampoline that we need to call to trigger those actions. auto CallbackInfo = CompileCallbacks.getCompileCallback(*F->getFunctionType()); // Step 3) Create a stub that will indirectly call the body of this // function once it is compiled. Initially, set the function - // pointer for the indirection to point at the compile callback. + // pointer for the indirection to point at the trampoline. std::string BodyPtrName = (F->getName() + "$address").str(); GlobalVariable *FunctionBodyPointer = createImplPointer(*F, BodyPtrName, CallbackInfo.getAddress()); @@ -1253,7 +1270,7 @@ public: // Step 4) Add the module containing the stub to the JIT. auto H = addModule(C.takeM()); - // Step 5) Set the compile and update actions for the callback. + // Step 5) Set the compile and update actions. // // The compile action will IRGen the function and add it to the JIT, then // request its address, which will trigger codegen. Since we don't need the @@ -1264,15 +1281,15 @@ public: // The update action will update FunctionBodyPointer to point at the newly // compiled function. CallbackInfo.setCompileAction( - [this,Fn = std::shared_ptr(std::move(FnAST))](){ + [this, Fn = std::shared_ptr(std::move(FnAST))](){ auto H = addModule(IRGen(Session, *Fn)); return findSymbolIn(H, Fn->Proto->Name).getAddress(); }); CallbackInfo.setUpdateAction( CompileCallbacks.getLocalFPUpdater(H, Mangle(BodyPtrName))); - } -private: + return H; + } std::unique_ptr TM; Mangler Mang; @@ -1283,6 +1300,8 @@ private: LazyEmitLayerT LazyEmitLayer; JITCompileCallbackManager CompileCallbacks; + + std::map> FunctionDefs; }; static void HandleDefinition(SessionContext &S, KaleidoscopeJIT &J) {