[ORC] Add more utilities to aid debugging output.

(1) A const accessor for the LLVMContext held by a ThreadSafeContext.

(2) A const accessor for the ThreadSafeModules held by an IRMaterializationUnit.

(3) A const MaterializationResponsibility reference to IRTransformLayer2's
    transform function. This makes IRTransformLayer2 useful for JIT debugging
    (since it can inspect JIT state through the responsibility argument) as well
    as program transformations.

llvm-svn: 343365

5 files changed, 23 insertions, 10 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h
index 1d4d33e5df4..d5f91cef359 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h
@@ -25,8 +25,8 @@ namespace orc {
 
 class IRTransformLayer2 : public IRLayer {
 public:
-  using TransformFunction =
-      std::function<Expected<ThreadSafeModule>(ThreadSafeModule)>;
+  using TransformFunction = std::function<Expected<ThreadSafeModule>(
+      ThreadSafeModule, const MaterializationResponsibility &R)>;
 
   IRTransformLayer2(ExecutionSession &ES, IRLayer &BaseLayer,
                     TransformFunction Transform = identityTransform);
@@ -38,7 +38,9 @@ public:
   void emit(MaterializationResponsibility R, VModuleKey K,
             ThreadSafeModule TSM) override;
 
-  static ThreadSafeModule identityTransform(ThreadSafeModule TSM) {
+  static ThreadSafeModule
+  identityTransform(ThreadSafeModule TSM,
+                    const MaterializationResponsibility &R) {
     return TSM;
   }
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Layer.h b/llvm/include/llvm/ExecutionEngine/Orc/Layer.h
index b3e14b755e9..0d3de74093e 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Layer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Layer.h
@@ -89,6 +89,8 @@ public:
   /// Return the ModuleIdentifier as the name for this MaterializationUnit.
   StringRef getName() const override;
 
+  const ThreadSafeModule &getModule() const { return TSM; }
+
 protected:
   ThreadSafeModule TSM;
   SymbolNameToDefinitionMap SymbolToDefinition;
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h b/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
index a12f19c30a6..bf946de532d 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
@@ -66,6 +66,10 @@ public:
   /// instance, or null if the instance was default constructed.
   LLVMContext *getContext() { return S ? S->Ctx.get() : nullptr; }
 
+  /// Returns a pointer to the LLVMContext that was used to construct this
+  /// instance, or null if the instance was default constructed.
+  const LLVMContext *getContext() const { return S ? S->Ctx.get() : nullptr; }
+
   Lock getLock() {
     assert(S && "Can not lock an empty ThreadSafeContext");
     return Lock(S);
diff --git a/llvm/lib/ExecutionEngine/Orc/IRTransformLayer.cpp b/llvm/lib/ExecutionEngine/Orc/IRTransformLayer.cpp
index ddd5c4a10c1..7a79a382d8d 100644
--- a/llvm/lib/ExecutionEngine/Orc/IRTransformLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/IRTransformLayer.cpp
@@ -22,7 +22,7 @@ void IRTransformLayer2::emit(MaterializationResponsibility R, VModuleKey K,
                              ThreadSafeModule TSM) {
   assert(TSM.getModule() && "Module must not be null");
 
-  if (auto TransformedTSM = Transform(std::move(TSM)))
+  if (auto TransformedTSM = Transform(std::move(TSM), R))
     BaseLayer.emit(std::move(R), std::move(K), std::move(*TransformedTSM));
   else {
     R.failMaterialization();
diff --git a/llvm/tools/lli/lli.cpp b/llvm/tools/lli/lli.cpp
index c9fc11b36b4..fab7b3505b3 100644
--- a/llvm/tools/lli/lli.cpp
+++ b/llvm/tools/lli/lli.cpp
@@ -702,10 +702,12 @@ int main(int argc, char **argv, char * const *envp) {
 static orc::IRTransformLayer2::TransformFunction createDebugDumper() {
   switch (OrcDumpKind) {
   case DumpKind::NoDump:
-    return [](orc::ThreadSafeModule TSM) { return TSM; };
+    return [](orc::ThreadSafeModule TSM,
+              const orc::MaterializationResponsibility &R) { return TSM; };
 
   case DumpKind::DumpFuncsToStdOut:
-    return [](orc::ThreadSafeModule TSM) {
+    return [](orc::ThreadSafeModule TSM,
+              const orc::MaterializationResponsibility &R) {
       printf("[ ");
 
       for (const auto &F : *TSM.getModule()) {
@@ -724,7 +726,8 @@ static orc::IRTransformLayer2::TransformFunction createDebugDumper() {
     };
 
   case DumpKind::DumpModsToStdOut:
-    return [](orc::ThreadSafeModule TSM) {
+    return [](orc::ThreadSafeModule TSM,
+              const orc::MaterializationResponsibility &R) {
       outs() << "----- Module Start -----\n"
              << *TSM.getModule() << "----- Module End -----\n";
 
@@ -732,7 +735,8 @@ static orc::IRTransformLayer2::TransformFunction createDebugDumper() {
     };
 
   case DumpKind::DumpModsToDisk:
-    return [](orc::ThreadSafeModule TSM) {
+    return [](orc::ThreadSafeModule TSM,
+              const orc::MaterializationResponsibility &R) {
       std::error_code EC;
       raw_fd_ostream Out(TSM.getModule()->getModuleIdentifier() + ".ll", EC,
                          sys::fs::F_Text);
@@ -792,12 +796,13 @@ int runOrcLazyJIT(const char *ProgName) {
 
   auto Dump = createDebugDumper();
 
-  J->setLazyCompileTransform([&](orc::ThreadSafeModule TSM) {
+  J->setLazyCompileTransform([&](orc::ThreadSafeModule TSM,
+                                 const orc::MaterializationResponsibility &R) {
     if (verifyModule(*TSM.getModule(), &dbgs())) {
       dbgs() << "Bad module: " << *TSM.getModule() << "\n";
       exit(1);
     }
-    return Dump(std::move(TSM));
+    return Dump(std::move(TSM), R);
   });
   J->getMainJITDylib().setFallbackDefinitionGenerator(
       orc::DynamicLibraryFallbackGenerator(