[LTO] Add option to emit assembly and bitcode from LTO CodeGen

[tobiasvk]

Summary

LTOCodeGenerator and the llvm-lto utility can currently only emit object code.
With this patch, we gain the ability to output assembly and bitcode if desired.
This is primarily useful for testing (e.g. lit tests of LTO CodeGen). Our
linker also has a -save-temps option to save assembly during LTO which can be
implemented on top of this functionality.

[pcc]

I don't think this is necessary. There is already a LTOCodeGenerator::writeMergedModules() function that you can use to write the bitcode. If you would like to test the module splitting functionality, you can use llvm-split on that bitcode file.

[tobiasvk]

Ah I wasn't aware of the llvm-split utility - makes sense. It would still be useful to get bitcode out of llvm-lto for testing. But that could be built on top of writeMergedModule, as you mentioned.

The part about emitting assembly remains, though.

I'll update the patch.

[mehdi_amini]

The bitcode->assembly part is available with llc though, not sure why it's not enough?

[tobiasvk]

In D14554#286654, @joker.eph wrote:

The bitcode->assembly part is available with llc though, not sure why it's not enough?

Yes... but I'd rather not have our linker depend on llc or, alternatively, set up its own CodeGen pass manager just to emit assembly. Also, if we're using module splitting, there'd be an additional step: bitcode -> llvm-split -> llc. The option in LTOCodeGenerator seems a lot cleaner to me.

[mehdi_amini]

Sorry, I missed the fact that the point was to enable an interface for your linker, I just thought about the llvm-lto case. If you need this functionality in your linker, then it makes sense.

[tobiasvk]

Updated patch with Peter and Mehdi's comments - thank you both!

The patch now only adds the ability to emit assembly to LTOCodeGenerator, and a new option -save-merged-module to llvm-lto to write out the merged module before CodeGen for testing.

[pcc]

LGTM

test/LTO/X86/llvm-lto-output.ll

12

Will the asm look like this on every architecture? I think so, but it may be safest to specify a target explicitly.

[mehdi_amini]

The llvm-lto changes seems unrelated to the ability to write ASM from LTOCodeGenerator right? This could be committed separately.

test/LTO/X86/llvm-lto-output.ll

12

Especially since the test is in the "X86" folder. I'd expect an X86 triple at least.

[mehdi_amini]

I meant the "SaveModuleFile" option in llvm-lto.

include/llvm/CodeGen/ParallelCG.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_PARALLELCG_H
 #define LLVM_CODEGEN_PARALLELCG_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/CodeGen.h"
+#include "llvm/Target/TargetMachine.h"
 
 namespace llvm {
 
 class Module;
 class TargetOptions;
 class raw_pwrite_stream;
 
 /// Split M into OSs.size() partitions, and generate code for each. Writes
-/// OSs.size() object files to the output streams in OSs. The resulting object
-/// files if linked together are intended to be equivalent to the single object
+/// OSs.size() output files to the output streams in OSs. The resulting output
+/// files if linked together are intended to be equivalent to the single output
 /// file that would have been code generated from M.
 ///
 /// \returns M if OSs.size() == 1, otherwise returns std::unique_ptr<Module>().
 std::unique_ptr<Module>
 splitCodeGen(std::unique_ptr<Module> M, ArrayRef<raw_pwrite_stream *> OSs,
              StringRef CPU, StringRef Features, const TargetOptions &Options,
              Reloc::Model RM = Reloc::Default,
              CodeModel::Model CM = CodeModel::Default,
-             CodeGenOpt::Level OL = CodeGenOpt::Default);
+             CodeGenOpt::Level OL = CodeGenOpt::Default,
+             TargetMachine::CodeGenFileType FT = TargetMachine::CGFT_ObjectFile);
 
 } // namespace llvm
 
 #endif

include/llvm/LTO/LTOCodeGenerator.h

 #ifndef LLVM_LTO_LTOCODEGENERATOR_H
 #define LLVM_LTO_LTOCODEGENERATOR_H
 
 #include "llvm-c/lto.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Linker/Linker.h"
+#include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include <string>
 #include <vector>
 
 namespace llvm {
   class LLVMContext;
   class DiagnosticInfo;
   class GlobalValue;
struct LTOCodeGenerator {
   bool addModule(struct LTOModule *);
 
   /// Set the destination module.
   void setModule(std::unique_ptr<LTOModule> M);
 
   void setTargetOptions(TargetOptions Options);
   void setDebugInfo(lto_debug_model);
   void setCodePICModel(Reloc::Model Model) { RelocModel = Model; }
   
+  /// Set the file type to be emitted (assembly or object code).
+  /// The default is TargetMachine::CGFT_ObjectFile. 
+  void setFileType(TargetMachine::CodeGenFileType FT) { FileType = FT; }
+
   void setCpu(const char *MCpu) { this->MCpu = MCpu; }
   void setAttr(const char *MAttr) { this->MAttr = MAttr; }
   void setOptLevel(unsigned OptLevel);
 
   void setShouldInternalize(bool Value) { ShouldInternalize = Value; }
   void setShouldEmbedUselists(bool Value) { ShouldEmbedUselists = Value; }
 
   void addMustPreserveSymbol(StringRef Sym) { MustPreserveSymbols[Sym] = 1; }
struct LTOCodeGenerator {
   /// LTOCodeGenerator::compilexxx() and
   /// LTOCodeGenerator::writeMergedModules().
   void parseCodeGenDebugOptions();
 
   /// Write the merged module to the file specified by the given path.  Return
   /// true on success.
   bool writeMergedModules(const char *Path, std::string &ErrMsg);
 
-  /// Compile the merged module into a *single* object file; the path to object
+  /// Compile the merged module into a *single* output file; the path to output
   /// file is returned to the caller via argument "name". Return true on
   /// success.
   ///
-  /// \note It is up to the linker to remove the intermediate object file.  Do
+  /// \note It is up to the linker to remove the intermediate output file.  Do
   /// not try to remove the object file in LTOCodeGenerator's destructor as we
-  /// don't who (LTOCodeGenerator or the obj file) will last longer.
+  /// don't who (LTOCodeGenerator or the output file) will last longer.
   bool compile_to_file(const char **Name, bool DisableVerify,
                        bool DisableInline, bool DisableGVNLoadPRE,
                        bool DisableVectorization, std::string &ErrMsg);
 
   /// As with compile_to_file(), this function compiles the merged module into
-  /// single object file. Instead of returning the object-file-path to the
-  /// caller (linker), it brings the object to a buffer, and return the buffer
-  /// to the caller. This function should delete intermediate object file once
+  /// single output file. Instead of returning the output file path to the
+  /// caller (linker), it brings the output to a buffer, and returns the buffer
+  /// to the caller. This function should delete the intermediate file once
   /// its content is brought to memory. Return NULL if the compilation was not
   /// successful.
   std::unique_ptr<MemoryBuffer> compile(bool DisableVerify, bool DisableInline,
                                         bool DisableGVNLoadPRE,
                                         bool DisableVectorization,
                                         std::string &errMsg);
 
   /// Optimizes the merged module.  Returns true on success.
   bool optimize(bool DisableVerify, bool DisableInline, bool DisableGVNLoadPRE,
                 bool DisableVectorization, std::string &ErrMsg);
 
-  /// Compiles the merged optimized module into a single object file. It brings
-  /// the object to a buffer, and returns the buffer to the caller. Return NULL
+  /// Compiles the merged optimized module into a single output file. It brings
+  /// the output to a buffer, and returns the buffer to the caller. Return NULL
   /// if the compilation was not successful.
   std::unique_ptr<MemoryBuffer> compileOptimized(std::string &ErrMsg);
 
-  /// Compile the merged optimized module into out.size() object files each
+  /// Compile the merged optimized module into out.size() output files each
   /// representing a linkable partition of the module. If out contains more
   /// than one element, code generation is done in parallel with out.size()
-  /// threads.  Object files will be written to members of out. Returns true on
+  /// threads.  Output files will be written to members of out. Returns true on
   /// success.
   bool compileOptimized(ArrayRef<raw_pwrite_stream *> Out, std::string &ErrMsg);
 
   void setDiagnosticHandler(lto_diagnostic_handler_t, void *);
 
   LLVMContext &getContext() { return Context; }
 
 private:
private:
   std::string NativeObjectPath;
   TargetOptions Options;
   CodeGenOpt::Level CGOptLevel = CodeGenOpt::Default;
   unsigned OptLevel = 2;
   lto_diagnostic_handler_t DiagHandler = nullptr;
   void *DiagContext = nullptr;
   bool ShouldInternalize = true;
   bool ShouldEmbedUselists = false;
+  TargetMachine::CodeGenFileType FileType = TargetMachine::CGFT_ObjectFile;
 };
 }
 #endif

lib/CodeGen/ParallelCG.cpp

 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/SplitModule.h"
 
 using namespace llvm;
 
 static void codegen(Module *M, llvm::raw_pwrite_stream &OS,
                     const Target *TheTarget, StringRef CPU, StringRef Features,
                     const TargetOptions &Options, Reloc::Model RM,
-                    CodeModel::Model CM, CodeGenOpt::Level OL) {
+                    CodeModel::Model CM, CodeGenOpt::Level OL,
+                    TargetMachine::CodeGenFileType FileType) {
   std::unique_ptr<TargetMachine> TM(TheTarget->createTargetMachine(
       M->getTargetTriple(), CPU, Features, Options, RM, CM, OL));
 
   legacy::PassManager CodeGenPasses;
-  if (TM->addPassesToEmitFile(CodeGenPasses, OS,
-                              TargetMachine::CGFT_ObjectFile))
+  if (TM->addPassesToEmitFile(CodeGenPasses, OS, FileType))
     report_fatal_error("Failed to setup codegen");
   CodeGenPasses.run(*M);
 }
 
 std::unique_ptr<Module>
 llvm::splitCodeGen(std::unique_ptr<Module> M,
                    ArrayRef<llvm::raw_pwrite_stream *> OSs, StringRef CPU,
                    StringRef Features, const TargetOptions &Options,
-                   Reloc::Model RM, CodeModel::Model CM, CodeGenOpt::Level OL) {
+                   Reloc::Model RM, CodeModel::Model CM, CodeGenOpt::Level OL,
+                   TargetMachine::CodeGenFileType FileType) {
   StringRef TripleStr = M->getTargetTriple();
   std::string ErrMsg;
   const Target *TheTarget = TargetRegistry::lookupTarget(TripleStr, ErrMsg);
   if (!TheTarget)
     report_fatal_error(Twine("Target not found: ") + ErrMsg);
 
   if (OSs.size() == 1) {
     codegen(M.get(), *OSs[0], TheTarget, CPU, Features, Options, RM, CM,
-            OL);
+            OL, FileType);
     return M;
   }
 
   std::vector<thread> Threads;
   SplitModule(std::move(M), OSs.size(), [&](std::unique_ptr<Module> MPart) {
     // We want to clone the module in a new context to multi-thread the codegen.
     // We do it by serializing partition modules to bitcode (while still on the
     // main thread, in order to avoid data races) and spinning up new threads
     // which deserialize the partitions into separate contexts.
     // FIXME: Provide a more direct way to do this in LLVM.
     SmallVector<char, 0> BC;
     raw_svector_ostream BCOS(BC);
     WriteBitcodeToFile(MPart.get(), BCOS);
 
     llvm::raw_pwrite_stream *ThreadOS = OSs[Threads.size()];
     Threads.emplace_back(
-        [TheTarget, CPU, Features, Options, RM, CM, OL,
+        [TheTarget, CPU, Features, Options, RM, CM, OL, FileType,
          ThreadOS](const SmallVector<char, 0> &BC) {
           LLVMContext Ctx;
           ErrorOr<std::unique_ptr<Module>> MOrErr =
               parseBitcodeFile(MemoryBufferRef(StringRef(BC.data(), BC.size()),
                                                "<split-module>"),
                                Ctx);
           if (!MOrErr)
             report_fatal_error("Failed to read bitcode");
           std::unique_ptr<Module> MPartInCtx = std::move(MOrErr.get());
 
           codegen(MPartInCtx.get(), *ThreadOS, TheTarget, CPU, Features,
-                  Options, RM, CM, OL);
+                  Options, RM, CM, OL, FileType);
         },
         // Pass BC using std::move to ensure that it get moved rather than
         // copied into the thread's context.
         std::move(BC));
   });
 
   for (thread &T : Threads)
     T.join();
 
   return {};
 }

lib/LTO/LTOCodeGenerator.cpp

bool LTOCodeGenerator::writeMergedModules(const char *Path,
   }
 
   Out.keep();
   return true;
 }
 
 bool LTOCodeGenerator::compileOptimizedToFile(const char **Name,
                                               std::string &ErrMsg) {
-  // make unique temp .o file to put generated object file
+  // make unique temp output file to put generated code
   SmallString<128> Filename;
   int FD;
+
+  const char *Extension =
+      (FileType == TargetMachine::CGFT_AssemblyFile ? "s" : "o");
+
   std::error_code EC =
-      sys::fs::createTemporaryFile("lto-llvm", "o", FD, Filename);
+      sys::fs::createTemporaryFile("lto-llvm", Extension, FD, Filename);
   if (EC) {
     ErrMsg = EC.message();
     return false;
   }
 
   // generate object file
   tool_output_file objFile(Filename.c_str(), FD);
 
bool LTOCodeGenerator::compileOptimized(ArrayRef<raw_pwrite_stream *> Out,
 
   // Do code generation. We need to preserve the module in case the client calls
   // writeMergedModules() after compilation, but we only need to allow this at
   // parallelism level 1. This is achieved by having splitCodeGen return the
   // original module at parallelism level 1 which we then assign back to
   // MergedModule.
   MergedModule =
       splitCodeGen(std::move(MergedModule), Out, MCpu, FeatureStr, Options,
-                   RelocModel, CodeModel::Default, CGOptLevel);
+                   RelocModel, CodeModel::Default, CGOptLevel, FileType);
 
   return true;
 }
 
 /// setCodeGenDebugOptions - Set codegen debugging options to aid in debugging
 /// LTO problems.
 void LTOCodeGenerator::setCodeGenDebugOptions(const char *Options) {
   for (std::pair<StringRef, StringRef> o = getToken(Options); !o.first.empty();

test/LTO/X86/llvm-lto-output.ll

+; Test the various output formats of the llvm-lto utility
+;
+; RUN: llvm-as < %s > %t1
+;
+; RUN: llvm-lto -exported-symbol=main -save-merged-module -filetype=asm -o %t2 %t1
+; RUN: llvm-dis -o - %t2.merged.bc | FileCheck %s
+; CHECK: @main()
+
+; RUN: FileCheck --check-prefix=ASM %s < %t2
+; RUN: llvm-lto -exported-symbol=main -filetype=obj -o %t2 %t1
+; RUN: llvm-objdump -d %t2 | FileCheck --check-prefix=ASM %s
+; ASM: main:
+;
+
+define i32 @main() {
+entry:
+  ret i32 23
+}
+

tools/llvm-lto/llvm-lto.cpp

 static cl::opt<bool>
 UseDiagnosticHandler("use-diagnostic-handler", cl::init(false),
   cl::desc("Use a diagnostic handler to test the handler interface"));
 
 static cl::opt<bool>
     ThinLTO("thinlto", cl::init(false),
             cl::desc("Only write combined global index for ThinLTO backends"));
 
+static cl::opt<bool>
+SaveModuleFile("save-merged-module", cl::init(false),
+               cl::desc("Write merged LTO module to file before CodeGen"));
+
 static cl::list<std::string>
 InputFilenames(cl::Positional, cl::OneOrMore,
   cl::desc("<input bitcode files>"));
 
 static cl::opt<std::string>
 OutputFilename("o", cl::init(""),
   cl::desc("Override output filename"),
   cl::value_desc("filename"));
for (unsigned i = 0; i < MAttrs.size(); ++i) {
     if (i > 0)
       attrs.append(",");
     attrs.append(MAttrs[i]);
   }
 
   if (!attrs.empty())
     CodeGen.setAttr(attrs.c_str());
 
+  if (FileType.getNumOccurrences())
+    CodeGen.setFileType(FileType);
+
   if (!OutputFilename.empty()) {
     std::string ErrorInfo;
     if (!CodeGen.optimize(DisableVerify, DisableInline, DisableGVNLoadPRE,
                           DisableLTOVectorization, ErrorInfo)) {
       errs() << argv[0] << ": error optimizing the code: " << ErrorInfo << "\n";
       return 1;
     }
 
+    if (SaveModuleFile) {
+      std::string ModuleFilename = OutputFilename;
+      ModuleFilename += ".merged.bc";
+      std::string ErrMsg;
+
+      if (!CodeGen.writeMergedModules(ModuleFilename.c_str(), ErrMsg)) {
+        errs() << argv[0] << ": " << ErrMsg << "\n";
+        return 1;
+      }
+    }
+
     std::list<tool_output_file> OSs;
     std::vector<raw_pwrite_stream *> OSPtrs;
     for (unsigned I = 0; I != Parallelism; ++I) {
       std::string PartFilename = OutputFilename;
       if (Parallelism != 1)
         PartFilename += "." + utostr(I);
       std::error_code EC;
       OSs.emplace_back(PartFilename, EC, sys::fs::F_None);
if (!OutputFilename.empty()) {
     for (tool_output_file &OS : OSs)
       OS.keep();
   } else {
     if (Parallelism != 1) {
       errs() << argv[0] << ": -j must be specified together with -o\n";
       return 1;
     }
 
+    if (SaveModuleFile) {
+      errs() << argv[0] << ": -save-merged-module must be specified with -o\n";
+      return 1;
+    }
+
     std::string ErrorInfo;
     const char *OutputName = nullptr;
     if (!CodeGen.compile_to_file(&OutputName, DisableVerify, DisableInline,
                                  DisableGVNLoadPRE, DisableLTOVectorization,
                                  ErrorInfo)) {
       errs() << argv[0]
              << ": error compiling the code: " << ErrorInfo
              << "\n";
       return 1;
     }
 
     outs() << "Wrote native object file '" << OutputName << "'\n";
   }
 
   return 0;
 }