Merge pull request swiftlang#33301 from compnerd/jit-state-of-mind

Frontend: support dumping the JIT state

Author: CodaFi

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