[LTO] Added flag to generate assembly file with after LTO passes

[bunty2020]

Summary

save temps option currently doesn't generate assembly file with -flto flag.
To generate assembly file after LTO passes use -Wl,-plugin-opt=-save-temps

[tejohnson]

Thanks for the patch. Looks like there wasn't a way to dump assembly before codegen even before the restructuring of this out of gold-plugin.

Your patch does not follow LLVM coding conventions. Please clang-format it.

Please also add a test.

lib/LTO/LTOBackend.cpp

45

Just "lto-print-asm" (drop leading 'f' since this is an internal option)

46

s/in//

205

The -Wl,-plugin-opt is gold specific, and this LTO support is used by multiple clients (with more coming). Just remove the "-Wl,-plugin-opt=" (and update option name.

206

This should be a static function. But in any case I am suggesting a different approach below where this would go away.

211

Better to set this up in addSaveTemps (but keep optional) and use the same file naming used there (and remove the FLTOAsmFilename option). You could then do it via a new variant of setHook (e.g. setAsmHook) and create a new config hook like PrintAsmHook.

E.g. something like

auto setAsmHook = [&](ModuleHookFn &Hook) {

 // Keep track of the hook created earlier, which also needs to run.
 ModuleHookFn ExistingHook = Hook;
 Hook = [=](unsigned Task, const Module &M, const TargetMachine *TM) {
... create the filename as in setHook but with .s extension and invoke addPassesToEmitFile...

}
...
setAsmHook(PrintAsmHook);

Then you would replace where you are calling genAssemblyFile with an invocation of that new hook.

[mehdi_amini]

lib/LTO/LTOBackend.cpp

49

I rather not have any cl::opt controlling the behavior of the API.

211

I agree that this would be better, however usual hooks don't have access to the targetMachine, and more important, the CodeGen actually modifies the Module, which forces it to be cloned entirely before dumping the ASM.

[bunty2020]

I have modified save temps to generate assembly file too. But, I had to clone the module for assembly generation. I would like to know if there is a work around to generate assembly without cloning module.

[mehdi_amini]

In D24020#531164, @bunty2020 wrote:

I have modified save temps to generate assembly file too. But, I had to clone the module for assembly generation. I would like to know if there is a work around to generate assembly without cloning module.

I already mentioned it in my previous comments: the codegen is modifying the IR. So running it two times on the same module is likely to give different results (Depending on your use-case, it may or may not matter to you, but it can be surprising).

include/llvm/LTO/Config.h

150

Comment says "before", name is "after".

Also you shouldn't describe what a hook is actually doing, this is up to the client.

lib/LTO/LTOBackend.cpp

123

This will be "costly", we don't want this by default with "addSaveTemps".
And since it hasn't been needed till now, it is likely a "rare" need. Most of the time (i.e. always till now), using clang -O2 -disable-llvm-optzns output.lto.bc -S is enough to reproduce the codegen and produce an assembly file out of the save-temps bitcode if needed.

[tejohnson]

include/llvm/LTO/Config.h

150

Part of the issue is that there is already a PreCodeGenModuleHook. That can't be used here though because of the need for passing the TM. So either name it something specific to asm generation like PreCodeGenAsmHook, or specific to taking the TM like PreCodeGenTMHook. Mehdi - WDYT?

248

Remove the whitespace change.

lib/LTO/LTOBackend.cpp

91

The block below here that computes the name and opens the file can be refactored out into a named lambda or helper function.

123

Right, I had suggested keeping this optional in my original suggestion about moving this to a hook. I think it should go back under the option from the original patch.

230

Remove whitespace change

include/llvm/LTO/Config.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LTO_CONFIG_H
 #define LLVM_LTO_CONFIG_H
 
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/Support/CodeGen.h"
+#include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 
 #include <functional>
 
 namespace llvm {
 
 class Error;
 class Module;
struct Config {
   /// Module hooks must be thread safe with respect to the linker's internal
   /// data structures. A module hook will never be called concurrently from
   /// multiple threads with the same task ID, or the same module.
   ///
   /// Note that in out-of-process backend scenarios, none of the hooks will be
   /// called for ThinLTO tasks.
   typedef std::function<bool(unsigned Task, const Module &)> ModuleHookFn;
 
+  typedef std::function<bool(unsigned Task, const Module &, TargetMachine *)>
+      ModuleAsmHookFn;
+
   /// This module hook is called after linking (regular LTO) or loading
   /// (ThinLTO) the module, before modifying it.
   ModuleHookFn PreOptModuleHook;
 
   /// This hook is called after promoting any internal functions
   /// (ThinLTO-specific).
   ModuleHookFn PostPromoteModuleHook;
 
   /// This hook is called after internalizing the module.
   ModuleHookFn PostInternalizeModuleHook;
 
   /// This hook is called after importing from other modules (ThinLTO-specific).
   ModuleHookFn PostImportModuleHook;
 
   /// This module hook is called after optimization is complete.
   ModuleHookFn PostOptModuleHook;
 
   /// This module hook is called before code generation. It is similar to the
   /// PostOptModuleHook, but for parallel code generation it is called after
   /// splitting the module.
   ModuleHookFn PreCodeGenModuleHook;
 
+  /// This module hook is called before code generation. It generates
+  /// assembly file by running all the codegen passes on module
+  ModuleAsmHookFn PostCodeGenModuleHook;
+
   /// A combined index hook is called after all per-module indexes have been
   /// combined (ThinLTO-specific). It can be used to implement -save-temps for
   /// the combined index.
   ///
   /// If this function returns false, any further processing for ThinLTO tasks
   /// is aborted.
   ///
   /// It is called regardless of whether the backend is in-process, although it
Config(Config &&X)
         DiagHandler(std::move(X.DiagHandler)),
         ResolutionFile(std::move(X.ResolutionFile)),
         PreOptModuleHook(std::move(X.PreOptModuleHook)),
         PostPromoteModuleHook(std::move(X.PostPromoteModuleHook)),
         PostInternalizeModuleHook(std::move(X.PostInternalizeModuleHook)),
         PostImportModuleHook(std::move(X.PostImportModuleHook)),
         PostOptModuleHook(std::move(X.PostOptModuleHook)),
         PreCodeGenModuleHook(std::move(X.PreCodeGenModuleHook)),
+        PostCodeGenModuleHook(std::move(X.PostCodeGenModuleHook)),
         CombinedIndexHook(std::move(X.CombinedIndexHook)) {}
   // FIXME: Remove once MSVC can synthesize move ops.
   Config &operator=(Config &&X) {
     CPU = std::move(X.CPU);
     Features = std::move(X.Features);
     Options = std::move(X.Options);
     MAttrs = std::move(X.MAttrs);
     RelocModel = std::move(X.RelocModel);
struct LTOLLVMContext : LLVMContext {
 
   LTOLLVMContext(const Config &C) : DiagHandler(C.DiagHandler) {
     setDiscardValueNames(C.ShouldDiscardValueNames);
     enableDebugTypeODRUniquing();
     setDiagnosticHandler(funcDiagHandler, &DiagHandler, true);
   }
   DiagnosticHandlerFunction DiagHandler;
 };
-
 }
 }
 
 #endif

lib/LTO/LTOBackend.cpp

 #include "llvm/MC/SubtargetFeature.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/FunctionImportUtils.h"
 #include "llvm/Transforms/Utils/SplitModule.h"
 
 using namespace llvm;
 using namespace lto;
 
 Error Config::addSaveTemps(std::string OutputFileName,
                            bool UseInputModulePath) {
   ShouldDiscardValueNames = false;
 
   std::error_code EC;
   ResolutionFile = llvm::make_unique<raw_fd_ostream>(
       OutputFileName + "resolution.txt", EC, sys::fs::OpenFlags::F_Text);
   if (EC)
     return errorCodeToError(EC);
 
   auto setHook = [&](std::string PathSuffix, ModuleHookFn &Hook) {
     // Keep track of the hook provided by the linker, which also needs to run.
     ModuleHookFn LinkerHook = Hook;
     Hook = [=](unsigned Task, const Module &M) {
       // If the linker's hook returned false, we need to pass that result
       // through.
       if (LinkerHook && !LinkerHook(Task, M))
         return false;
 
       std::string PathPrefix;
Hook = [=](unsigned Task, const Module &M) {
                      << '\n';
         exit(1);
       }
       WriteBitcodeToFile(&M, OS, /*ShouldPreserveUseListOrder=*/false);
       return true;
     };
   };
 
+  auto setAsmHook = [&](std::string PathSuffix, ModuleAsmHookFn &Hook) {
+    // Keep track of the hook provided by the linker, which also needs to run.
+    ModuleAsmHookFn LinkerHook = Hook;
+    Hook = [=](unsigned Task, const Module &M, TargetMachine *TM) {
+      //// If the linker's hook returned false, we need to pass that result
+      //// through.
+      // if (LinkerHook && !LinkerHook(Task, M))
+      //  return false;
+
+      auto AsmGenModule = CloneModule(&M);
+      std::string PathPrefix;
+      // If this is the combined module (not a ThinLTO backend compile) or the
+      // user hasn't requested using the input module's path, emit to a file
+      // named from the provided OutputFileName with the Task ID appended.
+      if (M.getModuleIdentifier() == "ld-temp.o" || !UseInputModulePath) {
+        PathPrefix = OutputFileName + utostr(Task);
+      } else
+        PathPrefix = M.getModuleIdentifier();
+      std::string Path = PathPrefix + "." + PathSuffix + ".s";
+      std::error_code EC;
+      int FD;
+      EC = sys::fs::openFileForWrite(StringRef(Path), FD, sys::fs::F_None);
+      if (EC) {
+        // Because -save-temps is a debugging feature, we report the error
+        // directly and exit.
+        llvm::errs() << "failed to open " << Path << ": " << EC.message()
+                     << '\n';
+        exit(1);
+      }
+
+      std::unique_ptr<raw_pwrite_stream> OS =
+          llvm::make_unique<llvm::raw_fd_ostream>(FD, true);
+
+      legacy::PassManager CodeGenPasses;
+
+      if (TM->addPassesToEmitFile(CodeGenPasses, *OS,
+                                  TargetMachine::CGFT_AssemblyFile))
+        report_fatal_error("Failed to setup assembly gen in codegen");
+
+      CodeGenPasses.run(*AsmGenModule);
+      return true;
+    };
+  };
+
   setHook("0.preopt", PreOptModuleHook);
   setHook("1.promote", PostPromoteModuleHook);
   setHook("2.internalize", PostInternalizeModuleHook);
   setHook("3.import", PostImportModuleHook);
   setHook("4.opt", PostOptModuleHook);
   setHook("5.precodegen", PreCodeGenModuleHook);
+  setAsmHook("6.postcodegen", PostCodeGenModuleHook);
 
   CombinedIndexHook = [=](const ModuleSummaryIndex &Index) {
     std::string Path = OutputFileName + "index.bc";
     std::error_code EC;
     raw_fd_ostream OS(Path, EC, sys::fs::OpenFlags::F_None);
     if (EC) {
       // Because -save-temps is a debugging feature, we report the error
       // directly and exit.
return [Task, &AddOutput](unsigned TaskId) {
     auto Output = AddOutput(Task);
     if (Output->isCachingEnabled() && Output->tryLoadFromCache(""))
       report_fatal_error("Cache hit without a valid key?");
     assert(Task == TaskId && "Unexpexted TaskId mismatch");
     return Output;
   };
 }
 
 void codegen(Config &C, TargetMachine *TM, AddOutputFn AddOutput, unsigned Task,
              Module &M) {
   if (C.PreCodeGenModuleHook && !C.PreCodeGenModuleHook(Task, M))
     return;
 
   auto Output = AddOutput(Task);
   std::unique_ptr<raw_pwrite_stream> OS = Output->getStream();
   legacy::PassManager CodeGenPasses;
+
+  // Generate assembly file
+  if (C.PostCodeGenModuleHook && !C.PostCodeGenModuleHook(Task, M, TM))
+    return;
+
   if (TM->addPassesToEmitFile(CodeGenPasses, *OS,
                               TargetMachine::CGFT_ObjectFile))
     report_fatal_error("Failed to setup codegen");
   CodeGenPasses.run(M);
 }
 
 void splitCodeGen(Config &C, TargetMachine *TM, AddOutputFn AddOutput,
                   unsigned ParallelCodeGenParallelismLevel,
else if (M.getTargetTriple().empty())
     M.setTargetTriple(C.DefaultTriple);
 
   std::string Msg;
   const Target *T = TargetRegistry::lookupTarget(M.getTargetTriple(), Msg);
   if (!T)
     return make_error<StringError>(Msg, inconvertibleErrorCode());
   return T;
 }
-
 }
 
 Error lto::backend(Config &C, AddOutputFn AddOutput,
                    unsigned ParallelCodeGenParallelismLevel,
                    std::unique_ptr<Module> M) {
   Expected<const Target *> TOrErr = initAndLookupTarget(C, *M);
   if (!TOrErr)
     return TOrErr.takeError();