Properly implement the gpu-kernel ABI for amdgpu

[Flakebi]

View all comments

Add support to pass structs, arrays and vectors to amdgpu kernels.
Scalars and vectors are taken by value, aggregates are passed by byref
pointers. Structs containing a single scalar/vector are handled like
a scalar.

Judging from clang tests, nvptx seems to do somewhat the same, just
using byval instead of byref: https://github.com/llvm/llvm-project/blob/e4e18dba3d77f4a3eea58bcc9ccae5a5498ede7c/clang/test/CodeGenCUDA/kernel-args.cu

I tested a couple of the lit test signatures on real hardware and it
seems to work fine. Given the relatively simple implementation, I hope
this amount of testing is enough (the C calling convention seems like
a worse fit for Rust’s current ABI code, it’s still giving me headaches).

This adds two members to PassMode::Indirect.

address_space specifies the address space of an on_stack/byval or
by_ref pointer argument.

by_ref translates to LLVM’s byref, which is similar to on_stack/byval,
however, there is no extra copy made, the pointer may not point to the
stack but can point to some other address space, and the passed argument
should not be modified.

Both are used by the amdgpu target to implement the gpu-kernel
ABI.

Tracking issue for the gpu-kernel ABI: #135467
Tracking issue for the amdgpu target: #135024

If I read it correctly, I can’t notify the gpu-target group, so cc @kjetilkjeka, @kulst, @ZuseZ4, @workingjubilee

[rustbot]

rustc_codegen_cranelift is developed in its own repository. If possible, consider making this change to rust-lang/rustc_codegen_cranelift instead.

cc @bjorn3

This PR changes rustc_public

cc @oli-obk, @celinval, @ouz-a, @makai410

rustc_codegen_gcc is developed in its own repository. If possible, consider making this change to rust-lang/rustc_codegen_gcc instead.

cc @antoyo, @GuillaumeGomez

[rustbot]

r? @petrochenkov

rustbot has assigned @petrochenkov.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

Why was this reviewer chosen?

The reviewer was selected based on:

• Owners of files modified in this PR: compiler
• compiler expanded to 75 candidates
• Random selection from 21 candidates

[petrochenkov]

@rustbot reroll

[ZuseZ4]

I think this now sends our slices through the aggregate path, not the (Scalar)Pair anymore, can you add a test to confirm that? It would break Rust Offload, but for now we can make a PR to overwrite this change for functions that are a OFFLOAD_KERNEL. So if this otherwise more closely matches hip/clang, then I think that's still an improvement. With the planned change from libomptarget to Offload APIs Rust offload should also become more flexible and able to handle this directly, but I'll check.

[Flakebi]

Yes, scalar pair (and therefore all fat pointers, including slices) are handled as aggregates with this and passed as byref ptr.

Are you sure that this breaks Rust Offload?
The way arguments are passed from the CPU side is equivalent between direct values and byref ptr, it’s all directly stored in the argument buffer, without pointer indirections. So, no changes required on the CPU side to pass slices.
My guess/hope is that this does not break anything but only fixes things :)
(A case that breaks is if someone hackily passed real structs before by adding them as pointer in the CPU argument memory. I hope nobody did that and expects it to keep working…)

I will add a test that passes a slice.

[Flakebi]

Pre-committed the tests, fixed the now perma-link in the commit message and added a test taking a slice as argument. Total diff (just adding the slice test): https://github.com/rust-lang/rust/compare/371ff93ae5d9f8cbbc07c15451218ecbeab39a9c..b5b814cbbf6c93535e33ce4232f698570dec25ce

[Flakebi]

Sorry, one more force-push to fix the tests that failed in CI (amended in the first commit). Just --blessing them was enough to add the new PassMode::Indirect members in the stderr output. Diff: https://github.com/rust-lang/rust/compare/6a47e13cce5e4608e10e26d1bec4b5af6baba10f..ddd9a73295fbd2136470a8d0938a7fb565b6e0c5

[Flakebi]

And one more to fix the aarch64-only variant of the test. Diff: https://github.com/rust-lang/rust/compare/ddd9a73295fbd2136470a8d0938a7fb565b6e0c5..a90629b7e87febc01a7e6bd4c56bba2232472c05

[mejrs]

r? @ZuseZ4

[bjorn3]

Is byref ptr different from ptr in any way other than byref ptr having more UB? Also please make sure to add a copy of byref ptr parameters. Rustc assumes that it can write through ptr arguments, but byref ptr makes that UB.

[Flakebi]

Is byref ptr different from ptr in any way other than byref ptr having more UB?

Yes, byref ptr (and byval ptr) are both different from ptr.

On the calling side, a ptr would be passed in as a pointer.
However, a byref ptr is passed in as a value.

Arguably, there’s not much difference in the called gpu-kernel itself if you look just at the IR going into LLVM. The argument is handled like a pointer. Later in the backend part that lowers arguments to reads from the argument memory region, byref ptr is handled differently than ptr, either reading a pointer from the argument memory, or reading the value directly.

For gpu-kernel, this is observable in “user code”, i.e. outside the compiler, because the user assembles a memory region to pass as arguments when launching a gpu-kernel on the GPU through some API (cuda/hip/hsa/sycl).
A ptr argument means one needs to write a pointer to that argument memory region.
A byref ptr means one needs to write the passed value into the argument memory region (i.e. no double indirection).

So, if the function signature in Rust is extern "gpu-kernel" fn mykernel(p: *const Struct), the user should write a pointer into the argument memory region and the rustc should use ptr.
If the function signature in Rust is extern "gpu-kernel" fn mykernel(s: Struct), the user should write the struct into the argument memory region and the rustc should use byref ptr.

Also please make sure to add a copy of byref ptr parameters. Rustc assumes that it can write through ptr arguments, but byref ptr makes that UB.

I think the change in compiler/rustc_codegen_ssa/src/mir/mod.rs should take care of that

[bjorn3]

However, a byref ptr is passed in as a value.

That would be byval, right? On x86_64 there is absolutely no difference in emitted assembly between ptr and say ptr byref(<128 x i8>): https://rust.godbolt.org/z/T5joGfWWE

I think the change in compiler/rustc_codegen_ssa/src/mir/mod.rs should take care of that

👍

[rust-bors]

📌 Commit c498254 has been approved by bjorn3

It is now in the queue for this repository.

[JonathanBrouwer]

💔 I suspect this PR failed tests as part of a rollup
@bors r-

After fixing the problem, consider running a try job for the failed job before re-approving.

Link to failure: #163000 (comment)

[rust-bors]

This pull request was unapproved.

This PR was contained in a rollup (#163000), which was unapproved.

View changes since this unapproval

[rustbot]

This PR was rebased onto a different main commit. Here's a range-diff highlighting what actually changed.

Rebasing is a normal part of keeping PRs up to date, so no action is needed—this note is just to help reviewers.

[Flakebi]

Rebased, the failure was caused by new test configs being added between the last push here and the bors merge. The tests are running in GitHub CI, so I don’t think a try run is necessary to make sure these tests are working.

[bjorn3]

@bors r+

[rust-bors]

📌 Commit 09e54da has been approved by bjorn3

It is now in the queue for this repository.

[rust-bors]

☔ The latest upstream changes made this pull request unmergeable. Please resolve the merge conflicts by rebasing.

This pull request was unapproved.