Repository navigation
[clang-nvlink-wrapper] Keep LTO task outputs distinct - #238883
Conversation
|
Hello @0z5a 👋 Thank you for submitting a Pull Request (PR) to the LLVM Project. Since this is your first PR, here are a few useful links covering our main contribution policies and review practices.
Please reply to this message to confirm that you have read these policies, especially the LLVM AI Tool Use Policy, and that any AI tool usage has been noted in the PR description. Frequently asked questionsHow do I add reviewers? This PR will be automatically labeled, and the relevant teams will be notified. For some parts of the project, reviewers may also be added automatically. You can also add reviewers manually using the Reviewers section on this page. If you cannot use that section, it is probably because you do not have write permissions for the repository. In that case, you can request a review by tagging reviewers in a comment using What if there are no comments? If you have not received any comments on your PR after a week, you can request a review by pinging the PR with a comment such as “Ping”. The common courtesy ping rate is once a week. Please remember that you are asking for volunteer time from other developers. Are any special GitHub settings required to contribute to LLVM? We only require contributors to have a public email address associated with their GitHub commits, see this section of LLVM Developer Policy for details. If you have questions, feel free to leave a comment on this PR, or ask on LLVM Discord or LLVM Discourse. Thank you, |
llvmorg-github-actions
Bot
added
the
clang
|
@llvm/pr-subscribers-clang Author: 0z5a (0z5a) ChangesWith Fixes #238882 Testing:
3 Files Affected:
diff --git a/clang/test/OffloadTools/clang-nvlink-wrapper/lto-save-temps.ll b/clang/test/OffloadTools/clang-nvlink-wrapper/lto-save-temps.ll
new file mode 100644
index 00000000000000..4adadb5a144858
--- /dev/null
+++ b/clang/test/OffloadTools/clang-nvlink-wrapper/lto-save-temps.ll
@@ -0,0 +1,25 @@
+; REQUIRES: nvptx-registered-target
+;
+; RUN: llvm-as %s -o %t.bc
+; RUN: mkdir -p %t.dir
+; RUN: cd %t.dir && clang-nvlink-wrapper --dry-run --save-temps --plugin-opt=lto-partitions=2 -arch sm_90 %t.bc -o multi 2>&1 | FileCheck %s --check-prefix=MULTI
+; RUN: cd %t.dir && clang-nvlink-wrapper --dry-run --save-temps -arch sm_90 %t.bc -o single 2>&1 | FileCheck %s --check-prefix=SINGLE
+;
+; MULTI: ptxas{{.*}} -c multi.0.s {{.*}} -o multi.0.cubin
+; MULTI: ptxas{{.*}} -c multi.1.s {{.*}} -o multi.1.cubin
+; MULTI: nvlink{{.*}} -o multi multi.0.cubin multi.1.cubin
+; SINGLE: ptxas{{.*}} -c single.s {{.*}} -o single.cubin
+; SINGLE: nvlink{{.*}} -o single single.cubin
+
+target datalayout = "e-p6:32:32-i64:64-i128:128-i256:256-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64-nvidia-cuda"
+
+define i32 @foo(i32 %x) {
+ %r = add i32 %x, 1
+ ret i32 %r
+}
+
+define i32 @bar(i32 %x) {
+ %r = mul i32 %x, 3
+ ret i32 %r
+}
diff --git a/clang/test/OffloadTools/clang-nvlink-wrapper/thinlto-save-temps.c b/clang/test/OffloadTools/clang-nvlink-wrapper/thinlto-save-temps.c
new file mode 100644
index 00000000000000..a182700e5d9884
--- /dev/null
+++ b/clang/test/OffloadTools/clang-nvlink-wrapper/thinlto-save-temps.c
@@ -0,0 +1,10 @@
+// REQUIRES: nvptx-registered-target
+// RUN: %clang -cc1 -triple nvptx64-nvidia-cuda -flto=thin -funified-lto -emit-llvm-bc %s -o %t.bc
+// RUN: mkdir -p %t.dir
+// RUN: cd %t.dir && clang-nvlink-wrapper --dry-run --save-temps --plugin-opt=thinlto -arch sm_90 %t.bc -o thin 2>&1 | FileCheck %s
+// CHECK-NOT: ptxas{{.*}} -c -O
+// CHECK: ptxas{{.*}} -c thin.1.s {{.*}} -o thin.1.cubin
+// CHECK-NEXT: {{.*}}nvlink{{.*}} -o thin thin.1.cubin
+
+int foo(int x) { return x + 1; }
+int bar(int x) { return x * 3; }
diff --git a/clang/tools/clang-nvlink-wrapper/ClangNVLinkWrapper.cpp b/clang/tools/clang-nvlink-wrapper/ClangNVLinkWrapper.cpp
index 6550902a6a6551..36ea0a3768471e 100644
--- a/clang/tools/clang-nvlink-wrapper/ClangNVLinkWrapper.cpp
+++ b/clang/tools/clang-nvlink-wrapper/ClangNVLinkWrapper.cpp
@@ -48,6 +48,8 @@
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/WithColor.h"
+#include <mutex>
+
using namespace llvm;
using namespace llvm::opt;
using namespace llvm::object;
@@ -327,7 +329,8 @@ struct Symbol {
bool UsedInRegularObj;
};
-Expected<StringRef> runPTXAs(StringRef File, const ArgList &Args) {
+Expected<StringRef> runPTXAs(StringRef File, const ArgList &Args, size_t Task,
+ size_t MaxTasks) {
SmallVector<StringRef, 1> SearchPaths;
if (Arg *A = Args.getLastArg(OPT_cuda_path_EQ))
SearchPaths.push_back(Args.MakeArgString(A->getValue() + Twine("/bin")));
@@ -342,8 +345,11 @@ Expected<StringRef> runPTXAs(StringRef File, const ArgList &Args) {
return createStringError(
"must pass in an explicit nvptx64 gpu architecture to 'ptxas'");
- auto TempFileOrErr = createTempFile(
- Args, sys::path::stem(Args.getLastArgValue(OPT_o, "a.out")), "cubin");
+ std::string Prefix =
+ sys::path::stem(Args.getLastArgValue(OPT_o, "a.out")).str();
+ if (MaxTasks > 1)
+ Prefix += "." + std::to_string(Task);
+ auto TempFileOrErr = createTempFile(Args, Prefix, "cubin");
if (!TempFileOrErr)
return TempFileOrErr.takeError();
@@ -685,6 +691,7 @@ getInput(const cl::LibraryOptionsContext &LibraryOptions, const ArgList &Args) {
// Run the LTO job to compile the bitcode.
size_t MaxTasks = LTOBackend.getMaxTasks();
SmallVector<StringRef> LTOFiles(MaxTasks);
+ std::mutex TempFilesMutex;
auto AddStream =
[&](size_t Task,
const Twine &ModuleName) -> std::unique_ptr<CachedFileStream> {
@@ -693,8 +700,12 @@ getInput(const cl::LibraryOptionsContext &LibraryOptions, const ArgList &Args) {
if (Args.hasArg(OPT_lto_emit_asm))
TempFile = Args.getLastArgValue(OPT_o, "a.out");
else {
- auto TempFileOrErr = createTempFile(
- Args, sys::path::stem(Args.getLastArgValue(OPT_o, "a.out")), "s");
+ std::lock_guard<std::mutex> Lock(TempFilesMutex);
+ std::string Prefix =
+ sys::path::stem(Args.getLastArgValue(OPT_o, "a.out")).str();
+ if (MaxTasks > 1)
+ Prefix += "." + std::to_string(Task);
+ auto TempFileOrErr = createTempFile(Args, Prefix, "s");
if (!TempFileOrErr)
reportError(TempFileOrErr.takeError());
TempFile = Args.MakeArgString(*TempFileOrErr);
@@ -711,8 +722,10 @@ getInput(const cl::LibraryOptionsContext &LibraryOptions, const ArgList &Args) {
if (Args.hasArg(OPT_lto_emit_llvm) || Args.hasArg(OPT_lto_emit_asm))
return Files;
- for (StringRef LTOFile : LTOFiles) {
- auto FileOrErr = runPTXAs(LTOFile, Args);
+ for (auto [Task, LTOFile] : enumerate(LTOFiles)) {
+ if (LTOFile.empty())
+ continue;
+ auto FileOrErr = runPTXAs(LTOFile, Args, Task, MaxTasks);
if (!FileOrErr)
return FileOrErr.takeError();
Files.emplace_back(*FileOrErr);
|
|
I have read the contribution policies, including the LLVM AI Tool Use Policy, and the AI tool usage is noted in the PR description. |
| "must pass in an explicit nvptx64 gpu architecture to 'ptxas'"); | ||
|
|
||
| auto TempFileOrErr = createTempFile( | ||
| Args, sys::path::stem(Args.getLastArgValue(OPT_o, "a.out")), "cubin"); |
There was a problem hiding this comment.
This should be changed to instead derive its name from the input file (i.e. foo.s -> foo.cubin). That should fix this and make the signature changes unnecessary.
There was a problem hiding this comment.
Done in 56ceb02: runPTXAs names the cubin after its input (foo.s -> foo.cubin) and keeps its old signature.
| // Run the LTO job to compile the bitcode. | ||
| size_t MaxTasks = LTOBackend.getMaxTasks(); | ||
| SmallVector<StringRef> LTOFiles(MaxTasks); | ||
| std::mutex TempFilesMutex; |
There was a problem hiding this comment.
A mutex shouldn't be necessary. createTempFile is intended to be thread safe on its own and everything else should have an index in an array of size MaxTasks like above.
There was a problem hiding this comment.
In clang-linker-wrapper it is: createOutputFile takes TempFilesMutex and TempFiles is a std::list. The copy here appends to a SmallVector without a lock, and AddStream runs on the LTO codegen thread pool (splitCodeGen), so two tasks can push to TempFiles at once. I moved the lock into createTempFile and made TempFiles a std::list to match, and dropped the mutex at the call site.
| auto TempFileOrErr = createTempFile( | ||
| Args, sys::path::stem(Args.getLastArgValue(OPT_o, "a.out")), "s"); |
There was a problem hiding this comment.
Referenced this from the existing LTO handling in lld.
StringRef Stem = sys::path::stem(Args.getLastArgValue(OPT_o, "a.out"));
auto TempFileOrErr = createTempFile(
Args, Stem + (Task == 0 ? Twine("") : Twine('.') + Twine(Task)), "s")There was a problem hiding this comment.
Switched to the lld scheme: task 0 gets <stem>.s, task N gets <stem>.N.s. The tests now expect multi.s/multi.cubin and multi.1.s/multi.1.cubin.
Name each LTO task's assembly like lld (`<stem>.s`, then `<stem>.<task>.s`) and name the cubin after the assembly it comes from, so runPTXAs keeps its signature. Lock createTempFile itself, as clang-linker-wrapper does, and keep TempFiles in a std::list so returned references stay valid, instead of a mutex at the call site. Co-Authored-By: Claude Opus 5.5 <[email protected]>
|
@0z5a Congratulations on having your first Pull Request (PR) merged into the LLVM Project! Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR. Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues. How to do this, and the rest of the post-merge process, is covered in detail here. If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again. If you don't get any reports, no action is required from you. Your changes are working as expected, well done! |
With
--save-tempsand multiple LTO tasks, the wrapper reuses the same PTX and cubin names, so later tasks overwrite earlier results andnvlinkreceives duplicate inputs. Give each task a distinct output name, serialize creation of shared temporary-file bookkeeping, and skip inactive task slots.Fixes #238882
Testing:
FileCheckregressions; Unified ThinLTO regression for the unused task slot.ptxasandnvlink: successfully linked a two-kernel, two-partition image forsm_90andsm_103awith--save-temps. The unpatched current baseline fails with duplicate definitions.AI assistance
OpenAI Codex was used to prepare this change and its tests, and Claude Code (Claude Opus 5.5) was used for the follow-up that addresses review comments. I reviewed every changed line and the test results above, and I am responsible for the contribution.