Frontend: support dumping the JIT state

Add a debugging mechanism that enables the JIT to dump the LLVM IR and
object files to enable debugging the JIT.  This makes it easier to debug
the JIT mode failures.  The idea was from Lang Hames!

Author: compnerd

include/swift/AST/IRGenOptions.h

@@ -139,6 +139,12 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
   PointerAuthSchema ResilientClassStubInitCallbacks;
 };
 
+enum class JITDebugArtifact : unsigned {
+  None,   ///< None
+  LLVMIR, ///< LLVM IR
+  Object, ///< Object File
+};
+
 /// The set of options supported by IR generation.
 class IRGenOptions {
 public:
@@ -326,6 +332,8 @@ class IRGenOptions {
   Optional<llvm::VersionTuple> AutolinkRuntimeCompatibilityLibraryVersion;
   Optional<llvm::VersionTuple> AutolinkRuntimeCompatibilityDynamicReplacementLibraryVersion;
 
+  JITDebugArtifact DumpJIT = JITDebugArtifact::None;
+
   IRGenOptions()
       : DWARFVersion(2), OutputKind(IRGenOutputKind::LLVMAssembly),
         Verify(true), OptMode(OptimizationMode::NotSet),

include/swift/Option/FrontendOptions.td

@@ -329,6 +329,9 @@ def disable_llvm_verify : Flag<["-"], "disable-llvm-verify">,
 def disable_llvm_value_names : Flag<["-"], "disable-llvm-value-names">,
   HelpText<"Don't add names to local values in LLVM IR">;
 
+def dump_jit : JoinedOrSeparate<["-"], "dump-jit">,
+  HelpText<"Dump JIT contents">;
+
 def enable_llvm_value_names : Flag<["-"], "enable-llvm-value-names">,
   HelpText<"Add names to local values in LLVM IR">;
 

lib/Frontend/CompilerInvocation.cpp

@@ -1448,9 +1448,23 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
   if (Args.hasArg(OPT_disable_concrete_type_metadata_mangled_name_accessors))
     Opts.DisableConcreteTypeMetadataMangledNameAccessors = true;
 
-  if (Args.hasArg(OPT_use_jit))
+  if (Args.hasArg(OPT_use_jit)) {
     Opts.UseJIT = true;
-  
+    if (const Arg *A = Args.getLastArg(OPT_dump_jit)) {
+      llvm::Optional<swift::JITDebugArtifact> artifact =
+          llvm::StringSwitch<llvm::Optional<swift::JITDebugArtifact>>(A->getValue())
+              .Case("llvm-ir", JITDebugArtifact::LLVMIR)
+              .Case("object", JITDebugArtifact::Object)
+              .Default(None);
+      if (!artifact) {
+        Diags.diagnose(SourceLoc(), diag::error_invalid_arg_value,
+                       A->getOption().getName(), A->getValue());
+        return true;
+      }
+      Opts.DumpJIT = *artifact;
+    }
+  }
+
   for (const Arg *A : Args.filtered(OPT_verify_type_layout)) {
     Opts.VerifyTypeLayoutNames.push_back(A->getValue());
   }

lib/Immediate/Immediate.cpp

@@ -30,6 +30,7 @@
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Config/config.h"
+#include "llvm/ExecutionEngine/Orc/DebugUtils.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/IR/LLVMContext.h"
@@ -75,6 +76,18 @@ static void *loadRuntimeLib(StringRef sharedLibName,
   return nullptr;
 }
 
+static void DumpLLVMIR(const llvm::Module &M) {
+  std::string path = (M.getName() + ".ll").str();
+  for (size_t count = 0; llvm::sys::fs::exists(path); )
+    path = (M.getName() + llvm::utostr(count++) + ".ll").str();
+
+  std::error_code error;
+  llvm::raw_fd_ostream stream(path, error);
+  if (error)
+    return;
+  M.print(stream, /*AssemblyAnnotationWriter=*/nullptr);
+}
+
 void *swift::immediate::loadSwiftRuntime(ArrayRef<std::string>
                                          runtimeLibPaths) {
 #if defined(_WIN32)
@@ -290,6 +303,18 @@ int swift::RunImmediately(CompilerInstance &CI,
   }
 
   auto Module = GenModule.getModule();
+
+  switch (IRGenOpts.DumpJIT) {
+  case JITDebugArtifact::None:
+    break;
+  case JITDebugArtifact::LLVMIR:
+    DumpLLVMIR(*Module);
+    break;
+  case JITDebugArtifact::Object:
+    JIT->getObjTransformLayer().setTransform(llvm::orc::DumpObjects());
+    break;
+  }
+
   {
     // Get a generator for the process symbols and attach it to the main
     // JITDylib.

test/IRGen/jit-debugging.swift

@@ -0,0 +1,15 @@
+// %empty-directory(%t)
+// RUN: not %target-swift-frontend -use-jit -dump-jit invalid -interpret %s 2>&1 | %FileCheck -check-prefix CHECK-INVALID %s
+// CHECK-INVALID: error: invalid value 'invalid' in 'dump-jit'
+
+// RUN: %empty-directory(%t)
+// RUN: cd %t && %target-swift-frontend -use-jit -dump-jit llvm-ir -interpret %s
+// RUN: %FileCheck -check-prefix CHECK-LLIR %s < %t/main.ll
+// CHECK-LLIR: ; ModuleID = 'main'
+
+// RUN: %empty-directory(%t)
+// RUN: cd %t && %target-swift-frontend -use-jit -dump-jit object -interpret %s
+// RUN: %llvm-nm --defined-only --extern-only %t/main-jitted-objectbuffer.o | %FileCheck -check-prefix CHECK-OBJ %s
+// CHECK-OBJ: T {{_?}}main
+
+let zero = 0