[lld][LTO] Add assembly output to LTO save-temps

[Pierre-vh]

Summary

When using -save-temps in Clang, we usually get a .s assembly file alongside the .o output.
However, when using LTO, Clang no longer produces those .s files because it "delegates" the responsibility of writing those temporaries to LLD, but LLD doesn't support writing out a .s assembly file alongside the final .o output with LTO enabled.

This patch adds a new "asm" save-temps value and adds the necessary logic to the LTO library to allow it to write out an assembly file alongside an object file when requested.

To implement this, we use the same trick as Clang: Codegen twice but with a different filetype for the second run.
As this is a debug feature, the performance shouldn't matter much

[Pierre-vh]

Cleanup diff so it can be reviewed

[MaskRay]

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

[Pierre-vh]

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

With your suggestion, does it also generate the final, linked output file? Or just the .S?
The idea is to be able to get both alongside each other. We'd also like to have comments in the assembly file about things like register usage - the ones you get when you codegen to assembly, not when you disassemble.

Though, I am surprised that you see this as adding significant complexity, I tried to not make it too intrusive. Is it the whole change that's problematic, or just some part of it?
If yes, maybe I can re-do the parts you find problematic to simplify them. The way I see it, this kind of patch should just be a small feature addition that shouldn't increase complexity too much. If you think complexity has increased significantly then I definitely failed somewhere and I can try to correct it before discarding the idea entirely.

[yaxunl]

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

We use lld to do device LTO in the HIP toolchain. clang driver extracts device bitcode and passes them to lld. When -save-temps is used, clang driver adds -save-temps to the options passed to lld. It is not convenient for the users to rerun lld command to get the assembly. It is better for lld to generate the assembly when -save-temps is used because: 1. it avoids redundant work, especially device LTO is usually the most time-consuming part 2. to keep the temporary file names consistent 3. it is a useful feature for lld itself

[MaskRay]

In D138560#3951169, @yaxunl wrote:

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

We use lld to do device LTO in the HIP toolchain. clang driver extracts device bitcode and passes them to lld. When -save-temps is used, clang driver adds -save-temps to the options passed to lld. It is not convenient for the users to rerun lld command to get the assembly. It is better for lld to generate the assembly when -save-temps is used because: 1. it avoids redundant work, especially device LTO is usually the most time-consuming part 2. to keep the temporary file names consistent 3. it is a useful feature for lld itself

clang --save-temps does not pass --save-temps to the linker. The user can use -Wl,--save-temps which is more orthogonal. I can use -Wl,--lto-emit-asm today to get assembly output.

In D138560#3950224, @Pierre-vh wrote:

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

With your suggestion, does it also generate the final, linked output file? Or just the .S?
The idea is to be able to get both alongside each other. We'd also like to have comments in the assembly file about things like register usage - the ones you get when you codegen to assembly, not when you disassemble.

There is an empty output file to satisfy some build tools' requirement that an output is already emitted. You can find assembly files in ${output}1. As mentioned, they are just not named *.S.

Though, I am surprised that you see this as adding significant complexity, I tried to not make it too intrusive. Is it the whole change that's problematic, or just some part of it?
If yes, maybe I can re-do the parts you find problematic to simplify them. The way I see it, this kind of patch should just be a small feature addition that shouldn't increase complexity too much. If you think complexity has increased significantly then I definitely failed somewhere and I can try to correct it before discarding the idea entirely.

This is significant complexity because you do codegen twice // Doing codegen twice may seem inefficient, but this is designed as a debug , which you probably would agree as well that this is inelegant.

Well, the real question is when --lto-emit-asm exists, why bother with another mode and duplicating the existing functionality. So far I fail to see a convincing argument that this new mode is useful.

[Pierre-vh]

In D138560#3951995, @MaskRay wrote:

In D138560#3951169, @yaxunl wrote:

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

We use lld to do device LTO in the HIP toolchain. clang driver extracts device bitcode and passes them to lld. When -save-temps is used, clang driver adds -save-temps to the options passed to lld. It is not convenient for the users to rerun lld command to get the assembly. It is better for lld to generate the assembly when -save-temps is used because: 1. it avoids redundant work, especially device LTO is usually the most time-consuming part 2. to keep the temporary file names consistent 3. it is a useful feature for lld itself

clang --save-temps does not pass --save-temps to the linker. The user can use -Wl,--save-temps which is more orthogonal. I can use -Wl,--lto-emit-asm today to get assembly output.

In D138560#3950224, @Pierre-vh wrote:

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

With your suggestion, does it also generate the final, linked output file? Or just the .S?
The idea is to be able to get both alongside each other. We'd also like to have comments in the assembly file about things like register usage - the ones you get when you codegen to assembly, not when you disassemble.

There is an empty output file to satisfy some build tools' requirement that an output is already emitted. You can find assembly files in ${output}1. As mentioned, they are just not named *.S.

Though, I am surprised that you see this as adding significant complexity, I tried to not make it too intrusive. Is it the whole change that's problematic, or just some part of it?
If yes, maybe I can re-do the parts you find problematic to simplify them. The way I see it, this kind of patch should just be a small feature addition that shouldn't increase complexity too much. If you think complexity has increased significantly then I definitely failed somewhere and I can try to correct it before discarding the idea entirely.

This is significant complexity because you do codegen twice // Doing codegen twice may seem inefficient, but this is designed as a debug , which you probably would agree as well that this is inelegant.

Well, the real question is when --lto-emit-asm exists, why bother with another mode and duplicating the existing functionality. So far I fail to see a convincing argument that this new mode is useful.

It does pass it for HIP code when using -fgpu-rdc, see HIPAMD.cpp:151. This would be the primary use case because when building the same HIP code without -fgpu-rdc, we get a .S temporary, but if we pass fgpu-rdc, that temporary is no longer there.

There is an empty output file to satisfy some build tools' requirement that an output is already emitted. You can find assembly files in ${output}1. As mentioned, they are just not named *.S.

This is not about build systems at all, it's for debuggability. An empty file is not useful to us, we need the assembly with the comments.

This is significant complexity because you do codegen twice // Doing codegen twice may seem inefficient, but this is designed as a debug , which you probably would agree as well that this is inelegant.

I know it's inefficient but:

1. Clang already uses the same "trick" to get a .S temporary. It codegens twice. Performance matters less for debug cases so it should be fine.
2. There is no alternative that I'm aware of. LLVM either emits a .o, a .s, but not both unlike GCC which emits a .s that's then assembled into the .o

I'm also unsure I'd call this "significant complexity", as long as documentation is good (if it's lacking, please tell me where) the idea itself is simple IMO

[MaskRay]

In D138560#3952769, @Pierre-vh wrote:

In D138560#3951995, @MaskRay wrote:

In D138560#3951169, @yaxunl wrote:

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

We use lld to do device LTO in the HIP toolchain. clang driver extracts device bitcode and passes them to lld. When -save-temps is used, clang driver adds -save-temps to the options passed to lld. It is not convenient for the users to rerun lld command to get the assembly. It is better for lld to generate the assembly when -save-temps is used because: 1. it avoids redundant work, especially device LTO is usually the most time-consuming part 2. to keep the temporary file names consistent 3. it is a useful feature for lld itself

clang --save-temps does not pass --save-temps to the linker. The user can use -Wl,--save-temps which is more orthogonal. I can use -Wl,--lto-emit-asm today to get assembly output.

In D138560#3950224, @Pierre-vh wrote:

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

With your suggestion, does it also generate the final, linked output file? Or just the .S?
The idea is to be able to get both alongside each other. We'd also like to have comments in the assembly file about things like register usage - the ones you get when you codegen to assembly, not when you disassemble.

There is an empty output file to satisfy some build tools' requirement that an output is already emitted. You can find assembly files in ${output}1. As mentioned, they are just not named *.S.

Though, I am surprised that you see this as adding significant complexity, I tried to not make it too intrusive. Is it the whole change that's problematic, or just some part of it?
If yes, maybe I can re-do the parts you find problematic to simplify them. The way I see it, this kind of patch should just be a small feature addition that shouldn't increase complexity too much. If you think complexity has increased significantly then I definitely failed somewhere and I can try to correct it before discarding the idea entirely.

This is significant complexity because you do codegen twice // Doing codegen twice may seem inefficient, but this is designed as a debug , which you probably would agree as well that this is inelegant.

Well, the real question is when --lto-emit-asm exists, why bother with another mode and duplicating the existing functionality. So far I fail to see a convincing argument that this new mode is useful.

It does pass it for HIP code when using -fgpu-rdc, see HIPAMD.cpp:151. This would be the primary use case because when building the same HIP code without -fgpu-rdc, we get a .S temporary, but if we pass fgpu-rdc, that temporary is no longer there.

It seems that the logic in HIPAMD.cpp:151 can simply be removed.

There is an empty output file to satisfy some build tools' requirement that an output is already emitted. You can find assembly files in ${output}1. As mentioned, they are just not named *.S.

This is not about build systems at all, it's for debuggability. An empty file is not useful to us, we need the assembly with the comments.

I think you misunderstand my comment. The assembly output of --lto-emit-asm is of course not empty. The object file output is empty which is fine for your case.
If you want the object file output as well, re-run LTO without --lto-emit-asm. It's similar to what your patch intends to do, but places less burden to libLTO.

[Pierre-vh]

In D138560#3958723, @MaskRay wrote:

In D138560#3952769, @Pierre-vh wrote:

In D138560#3951995, @MaskRay wrote:

In D138560#3951169, @yaxunl wrote:

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

We use lld to do device LTO in the HIP toolchain. clang driver extracts device bitcode and passes them to lld. When -save-temps is used, clang driver adds -save-temps to the options passed to lld. It is not convenient for the users to rerun lld command to get the assembly. It is better for lld to generate the assembly when -save-temps is used because: 1. it avoids redundant work, especially device LTO is usually the most time-consuming part 2. to keep the temporary file names consistent 3. it is a useful feature for lld itself

clang --save-temps does not pass --save-temps to the linker. The user can use -Wl,--save-temps which is more orthogonal. I can use -Wl,--lto-emit-asm today to get assembly output.

In D138560#3950224, @Pierre-vh wrote:

In D138560#3950141, @MaskRay wrote:

I am unsure I am convinced with the motivation. Say you have two bitcode files a.o and b.o, can you just use clang++ -fuse-ld=lld a.o b.o -Wl,--lto-emit-asm? This can be used with -Wl,--save-temps as well. The output is named in term of the -o output file name, instead of *.s, but the slight difference doesn't warrant adding significant more complexity to LTOBackend.cpp as this patch does.

With your suggestion, does it also generate the final, linked output file? Or just the .S?
The idea is to be able to get both alongside each other. We'd also like to have comments in the assembly file about things like register usage - the ones you get when you codegen to assembly, not when you disassemble.

There is an empty output file to satisfy some build tools' requirement that an output is already emitted. You can find assembly files in ${output}1. As mentioned, they are just not named *.S.

Though, I am surprised that you see this as adding significant complexity, I tried to not make it too intrusive. Is it the whole change that's problematic, or just some part of it?
If yes, maybe I can re-do the parts you find problematic to simplify them. The way I see it, this kind of patch should just be a small feature addition that shouldn't increase complexity too much. If you think complexity has increased significantly then I definitely failed somewhere and I can try to correct it before discarding the idea entirely.

This is significant complexity because you do codegen twice // Doing codegen twice may seem inefficient, but this is designed as a debug , which you probably would agree as well that this is inelegant.

Well, the real question is when --lto-emit-asm exists, why bother with another mode and duplicating the existing functionality. So far I fail to see a convincing argument that this new mode is useful.

It does pass it for HIP code when using -fgpu-rdc, see HIPAMD.cpp:151. This would be the primary use case because when building the same HIP code without -fgpu-rdc, we get a .S temporary, but if we pass fgpu-rdc, that temporary is no longer there.

It seems that the logic in HIPAMD.cpp:151 can simply be removed.

There is an empty output file to satisfy some build tools' requirement that an output is already emitted. You can find assembly files in ${output}1. As mentioned, they are just not named *.S.

This is not about build systems at all, it's for debuggability. An empty file is not useful to us, we need the assembly with the comments.

I think you misunderstand my comment. The assembly output of --lto-emit-asm is of course not empty. The object file output is empty which is fine for your case.
If you want the object file output as well, re-run LTO without --lto-emit-asm. It's similar to what your patch intends to do, but places less burden to libLTO.

It seems that the logic in HIPAMD.cpp:151 can simply be removed.

Why? How would it solve the issue?

If you want the object file output as well, re-run LTO without --lto-emit-asm. It's similar to what your patch intends to do, but places less burden to libLTO.

Are you suggesting that we do this at the driver level instead? So create two linker jobs there if --save-temps is passed to get both a .o and .s?
Or that in this specific use case we don't make any code change and just work around the issue?

[MaskRay]

Have you tried experimenting with -Wl,--lto-emit-asm? How does it not fulfill your requirement?

[Pierre-vh]

I checked -lto-emit-asm and it doesn't look like it does what I need. Ideally we want to get the .s file without needing to run the command twice, we'd expect -save-temps to always give us a .s file and that -fgpu-rdc wouldn't affect it
I think landing this is the simplest to make the behavior of using LTO/-fgpu-rdc consistent with non-LTO/no -fgpu-rdc builds. The next best thing would be a Driver change but I think it'll end up adding more complexity.

Do you still think this isn't the right way to go?

[Pierre-vh]

After thinking about it a bit more, I think if we need to codegen twice it's probably better to do it at the driver level instead of "silently" running codegen twice

[MaskRay]

(I was out of town for ~2 weeks.)

Thanks for abandoning this. Running passes twice in LLVMLTO is definitely weird, even if it appears to be the simplest approach satisfying your goal.
If you want both .s and .o, I wish that you think about doing this as whatever tool which wraps clang driver instead of doing it in the clang driver.
Regular non-GPU clang driver code doesn't allow emit both .s and .o in one run, either.
I don't track GPU changes that closely, so if you end up doing that, I might not see or object in the end.

lld/ELF/Driver.cpp

 // Report a warning for an unknown -z option.
 static void checkZOptions(opt::InputArgList &args) {
   for (auto *arg : args.filtered(OPT_z))
     if (!isKnownZFlag(arg->getValue()))
       warn("unknown -z value: " + StringRef(arg->getValue()));
 }
 
 constexpr const char *saveTempsValues[] = {
-    "resolution", "preopt",     "promote", "internalize",  "import",
-    "opt",        "precodegen", "prelink", "combinedindex"};
+    "resolution", "preopt",     "promote", "internalize", "import",
+    "opt",        "precodegen", "asm",     "prelink",     "combinedindex"};
 
 void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   ELFOptTable parser;
   opt::InputArgList args = parser.parse(argsArr.slice(1));
 
   // Interpret these flags early because error()/warn() depend on them.
   errorHandler().errorLimit = args::getInteger(args, OPT_error_limit, 20);
   errorHandler().fatalWarnings =

lld/ELF/Options.td

 def opt_remarks_passes: Separate<["--"], "opt-remarks-passes">,
   HelpText<"Regex for the passes that need to be serialized to the output file">;
 def opt_remarks_with_hotness: FF<"opt-remarks-with-hotness">,
   HelpText<"Include hotness information in the optimization remarks file">;
 def opt_remarks_format: Separate<["--"], "opt-remarks-format">,
   HelpText<"The format used for serializing remarks (default: YAML)">;
 def save_temps: F<"save-temps">, HelpText<"Save intermediate LTO compilation results">;
 def save_temps_eq: JJ<"save-temps=">, HelpText<"Save select intermediate LTO compilation results">,
-  Values<"resolution,preopt,promote,internalize,import,opt,precodegen,prelink,combinedindex">;
+  Values<"resolution,preopt,promote,internalize,import,opt,precodegen,asm,prelink,combinedindex">;
 def lto_basic_block_sections: JJ<"lto-basic-block-sections=">,
   HelpText<"Enable basic block sections for LTO">;
 defm lto_unique_basic_block_section_names: BB<"lto-unique-basic-block-section-names",
     "Give unique names to every basic block section for LTO",
     "Do not give unique names to every basic block section for LTO (default)">;
 defm shuffle_sections: EEq<"shuffle-sections",
   "Shuffle matched sections using the given seed before mapping them to the output sections. "
   "If -1, reverse the section order. If 0, use a random seed">,

lld/test/COFF/savetemps.ll

 ; REQUIRES: x86
 ; RUN: rm -fr %T/savetemps
 ; RUN: mkdir %T/savetemps
 ; RUN: llvm-as -o %T/savetemps/savetemps.obj %s
 ; RUN: lld-link /out:%T/savetemps/savetemps.exe /entry:main \
 ; RUN:     /subsystem:console %T/savetemps/savetemps.obj
 ; RUN: not llvm-dis -o - %T/savetemps/savetemps.exe.0.0.preopt.bc
 ; RUN: not llvm-dis -o - %T/savetemps/savetemps.exe.0.2.internalize.bc
 ; RUN: not llvm-dis -o - %T/savetemps/savetemps.exe.0.4.opt.bc
-; RUN: not llvm-dis -o - %T/savetemps/savetemps.exe.0.5.precodegen.bc
+; RUN: not llvm-dis -o - %T/savetemps/savetemps.exe.0.5.precodegen.bc.
 ; RUN: not llvm-objdump -s %T/savetemps/savetemps.exe.lto.obj
 ; RUN: lld-link /lldsavetemps /out:%T/savetemps/savetemps.exe /entry:main \
 ; RUN:     /subsystem:console %T/savetemps/savetemps.obj
 ; RUN: llvm-dis -o - %T/savetemps/savetemps.exe.0.0.preopt.bc | FileCheck %s
 ; RUN: llvm-dis -o - %T/savetemps/savetemps.exe.0.2.internalize.bc | FileCheck %s
 ; RUN: llvm-dis -o - %T/savetemps/savetemps.exe.0.4.opt.bc | FileCheck %s
 ; RUN: llvm-dis -o - %T/savetemps/savetemps.exe.0.5.precodegen.bc | FileCheck %s
+; RUN: cat %T/savetemps/savetemps.exe.0.6.asm.s | FileCheck %s --check-prefix=CHECK-ASM
 ; RUN: llvm-objdump -s %T/savetemps/savetemps.exe.lto.obj | \
 ; RUN:     FileCheck --check-prefix=CHECK-OBJDUMP %s
 
 ; CHECK: define i32 @main()
 ; CHECK-OBJDUMP: file format coff
+; CHECK-ASM: .globl main
 
 target datalayout = "e-m:w-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-pc-windows-msvc"
 
 define i32 @main() {
   ret i32 0
 }

lld/test/ELF/lto/save-temps-eq.ll

 ; RUN: ls | count 2
 
 ;; Check precodegen
 ; RUN: ld.lld main.o thin1.o --save-temps=precodegen
 ; RUN: cmp %t/all/a.out a.out && rm -f a.out
 ; RUN: mv *.5.precodegen* %t/all3
 ; RUN: ls | count 2
 
+;; Check asm
+; RUN: ld.lld main.o thin1.o --save-temps=asm
+; RUN: cmp %t/all/a.out a.out && rm -f a.out
+; RUN: mv *.6.asm* %t/all3
+; RUN: ls | count 2
+
 ;; Check combinedindex
 ; RUN: ld.lld main.o thin1.o --save-temps=combinedindex
 ; RUN: cmp %t/all/a.out a.out && rm -f a.out
 ; RUN: mv *.index.bc %t/all3
 ; RUN: mv *.index.dot %t/all3
 ; RUN: ls | count 2
 
 ;; Check prelink

llvm/include/llvm/LTO/Config.h

struct Config {
   /// 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;
 
+  // When this hook is non-null, and the output file type isn't an assembly file
+  // performs codegen a second time to emit an assembly file into
+  // the output stream returned by this hook.
+  std::function<std::unique_ptr<raw_fd_ostream>(const Module &, unsigned Task)>
+      GetAssemblyOutputStream;
+
   /// 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

llvm/lib/LTO/LTOBackend.cpp

 }
 
 [[noreturn]] static void reportOpenError(StringRef Path, Twine Msg) {
   errs() << "failed to open " << Path << ": " << Msg << '\n';
   errs().flush();
   exit(1);
 }
 
+/// Creates an output stream for temporary files based on the Module's name,
+/// the output file name and other information such as the desired file
+/// extension.
+std::unique_ptr<raw_fd_ostream>
+createTempFileOutputStream(const Module &M, std::string OutputFileName,
+                           bool UseInputModulePath, unsigned Task,
+                           std::string PathSuffix, std::string Extension) {
+  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;
+    if (Task != (unsigned)-1)
+      PathPrefix += utostr(Task) + ".";
+  } else
+    PathPrefix = M.getModuleIdentifier() + ".";
+  std::string Path = PathPrefix + PathSuffix;
+  if (!Extension.empty())
+    Path += "." + Extension;
+  std::error_code EC;
+  auto OS =
+      std::make_unique<raw_fd_ostream>(Path, EC, sys::fs::OpenFlags::OF_None);
+  // Because -save-temps is a debugging feature, we report the error
+  // directly and exit.
+  if (EC)
+    reportOpenError(Path, EC.message());
+
+  return OS;
+}
+
 Error Config::addSaveTemps(std::string OutputFileName, bool UseInputModulePath,
                            const DenseSet<StringRef> &SaveTempsArgs) {
   ShouldDiscardValueNames = false;
 
   std::error_code EC;
   if (SaveTempsArgs.empty() || SaveTempsArgs.contains("resolution")) {
     ResolutionFile =
         std::make_unique<raw_fd_ostream>(OutputFileName + "resolution.txt", EC,
                                          sys::fs::OpenFlags::OF_TextWithCRLF);
     if (EC) {
       ResolutionFile.reset();
       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;
-      // 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;
-        if (Task != (unsigned)-1)
-          PathPrefix += utostr(Task) + ".";
-      } else
-        PathPrefix = M.getModuleIdentifier() + ".";
-      std::string Path = PathPrefix + PathSuffix + ".bc";
-      std::error_code EC;
-      raw_fd_ostream OS(Path, EC, sys::fs::OpenFlags::OF_None);
-      // Because -save-temps is a debugging feature, we report the error
-      // directly and exit.
-      if (EC)
-        reportOpenError(Path, EC.message());
-      WriteBitcodeToFile(M, OS, /*ShouldPreserveUseListOrder=*/false);
+      std::unique_ptr<raw_fd_ostream> OS = createTempFileOutputStream(
+          M, OutputFileName, UseInputModulePath, Task, PathSuffix, "bc");
+      WriteBitcodeToFile(M, *OS, /*ShouldPreserveUseListOrder=*/false);
       return true;
     };
   };
 
+  auto setSaveAsmOutputHook = [&](std::string PathSuffix) {
+    GetAssemblyOutputStream = [=](const Module &M, unsigned TaskID) {
+      return createTempFileOutputStream(M, OutputFileName, UseInputModulePath,
+                                        TaskID, PathSuffix, "s");
+    };
+  };
+
   auto SaveCombinedIndex =
       [=](const ModuleSummaryIndex &Index,
           const DenseSet<GlobalValue::GUID> &GUIDPreservedSymbols) {
         std::string Path = OutputFileName + "index.bc";
         std::error_code EC;
         raw_fd_ostream OS(Path, EC, sys::fs::OpenFlags::OF_None);
         // Because -save-temps is a debugging feature, we report the error
         // directly and exit.
Error Config::addSaveTemps(std::string OutputFileName, bool UseInputModulePath,
 
   if (SaveTempsArgs.empty()) {
     setHook("0.preopt", PreOptModuleHook);
     setHook("1.promote", PostPromoteModuleHook);
     setHook("2.internalize", PostInternalizeModuleHook);
     setHook("3.import", PostImportModuleHook);
     setHook("4.opt", PostOptModuleHook);
     setHook("5.precodegen", PreCodeGenModuleHook);
+    setSaveAsmOutputHook("6.asm");
     CombinedIndexHook = SaveCombinedIndex;
   } else {
     if (SaveTempsArgs.contains("preopt"))
       setHook("0.preopt", PreOptModuleHook);
     if (SaveTempsArgs.contains("promote"))
       setHook("1.promote", PostPromoteModuleHook);
     if (SaveTempsArgs.contains("internalize"))
       setHook("2.internalize", PostInternalizeModuleHook);
     if (SaveTempsArgs.contains("import"))
       setHook("3.import", PostImportModuleHook);
     if (SaveTempsArgs.contains("opt"))
       setHook("4.opt", PostOptModuleHook);
     if (SaveTempsArgs.contains("precodegen"))
       setHook("5.precodegen", PreCodeGenModuleHook);
+    if (SaveTempsArgs.contains("asm"))
+      setSaveAsmOutputHook("6.asm");
     if (SaveTempsArgs.contains("combinedindex"))
       CombinedIndexHook = SaveCombinedIndex;
   }
 
   return Error::success();
 }
 
 #define HANDLE_EXTENSION(Ext)                                                  \
static void codegen(const Config &Conf, TargetMachine *TM,
 
   Expected<std::unique_ptr<CachedFileStream>> StreamOrErr =
       AddStream(Task, Mod.getModuleIdentifier());
   if (Error Err = StreamOrErr.takeError())
     report_fatal_error(std::move(Err));
   std::unique_ptr<CachedFileStream> &Stream = *StreamOrErr;
   TM->Options.ObjectFilenameForDebug = Stream->ObjectPathName;
 
+  const auto GetPassManager = [&]() {
     legacy::PassManager CodeGenPasses;
     TargetLibraryInfoImpl TLII(Triple(Mod.getTargetTriple()));
     CodeGenPasses.add(new TargetLibraryInfoWrapperPass(TLII));
     CodeGenPasses.add(
         createImmutableModuleSummaryIndexWrapperPass(&CombinedIndex));
     if (Conf.PreCodeGenPassesHook)
       Conf.PreCodeGenPassesHook(CodeGenPasses);
+    return CodeGenPasses;
+  };
+
+  {
+    legacy::PassManager CodeGenPasses = GetPassManager();
     if (TM->addPassesToEmitFile(CodeGenPasses, *Stream->OS,
                                 DwoOut ? &DwoOut->os() : nullptr,
                                 Conf.CGFileType))
       report_fatal_error("Failed to setup codegen");
     CodeGenPasses.run(Mod);
+  }
+
+  // If we didn't output an assembly file, but one is desired, re-codegen
+  // the module to create the file in the desired output stream.
+  // Doing codegen twice may seem inefficient, but this is designed as a debug
+  // option (primarily for --save-temps) so performance isn't critical.
+  if (Conf.CGFileType != CGFT_AssemblyFile && Conf.GetAssemblyOutputStream) {
+    std::unique_ptr<raw_fd_ostream> AsmOutputStream =
+        Conf.GetAssemblyOutputStream(Mod, Task);
+    legacy::PassManager CodeGenPasses = GetPassManager();
+    if (TM->addPassesToEmitFile(CodeGenPasses, *AsmOutputStream, nullptr,
+                                CGFT_AssemblyFile))
+      report_fatal_error("Failed to setup codegen");
+    CodeGenPasses.run(Mod);
+  }
 
   if (DwoOut)
     DwoOut->keep();
 }
 
 static void splitCodeGen(const Config &C, TargetMachine *TM,
                          AddStreamFn AddStream,
                          unsigned ParallelCodeGenParallelismLevel, Module &Mod,

llvm/test/ThinLTO/X86/selective-save-temps.ll

 ; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
 ; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
 ; RUN:    -select-save-temps=precodegen
 ; RUN: cmp all/a.out.1 build/a.out.1 && rm -f build/a.out.1
 ; RUN: cmp all/a.out.2 build/a.out.2 && rm -f build/a.out.2
 ; RUN: mv build/*.5.precodegen* all2
 ; RUN: ls build | count 0
 
+;; Check asm
+; RUN: llvm-lto2 run 1.bc 2.bc -o build/a.out \
+; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
+; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
+; RUN:    -select-save-temps=asm
+; RUN: cmp all/a.out.1 build/a.out.1 && rm -f build/a.out.1
+; RUN: cmp all/a.out.2 build/a.out.2 && rm -f build/a.out.2
+; RUN: mv build/*.6.asm* all2
+; RUN: ls build | count 0
+
 ;; Check combinedindex
 ; RUN: llvm-lto2 run 1.bc 2.bc -o build/a.out \
 ; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
 ; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
 ; RUN:    -select-save-temps=combinedindex
 ; RUN: cmp all/a.out.1 build/a.out.1 && rm -f build/a.out.1
 ; RUN: cmp all/a.out.2 build/a.out.2 && rm -f build/a.out.2
 ; RUN: cp build/*.index.bc subset2

llvm/tools/llvm-lto2/llvm-lto2.cpp

static cl::opt<std::string> AAPipeline("aa-pipeline",
                                        cl::value_desc("aapipeline"));
 
 static cl::opt<bool> SaveTemps("save-temps", cl::desc("Save temporary files"));
 
 static cl::list<std::string> SelectSaveTemps(
     "select-save-temps",
     cl::value_desc("One, or multiple of: "
                    "resolution,preopt,promote,internalize,import,opt,precodegen"
-                   ",combinedindex"),
+                   ",asm,combinedindex"),
     cl::desc("Save selected temporary files. Cannot be specified together with "
              "-save-temps"),
     cl::CommaSeparated);
 
 constexpr const char *SaveTempsValues[] = {
-    "resolution", "preopt", "promote",    "internalize",
-    "import",     "opt",    "precodegen", "combinedindex"};
+    "resolution", "preopt",     "promote", "internalize",  "import",
+    "opt",        "precodegen", "asm",     "combinedindex"};
 
 static cl::opt<bool>
     ThinLTODistributedIndexes("thinlto-distributed-indexes",
                               cl::desc("Write out individual index and "
                                        "import files for the "
                                        "distributed backend case"));
 
 static cl::opt<bool>