Properly implement the gpu-kernel ABI for amdgpu - #162177

Flakebi · 2026-09-02T09:01:16Z

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 · 2026-09-02T09:01:20Z

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 · 2026-09-02T09:01:23Z

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 · 2026-09-02T10:24:22Z

@rustbot reroll

ZuseZ4 · 2026-09-02T12:03:29Z

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 · 2026-09-02T15:21:11Z

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 · 2026-09-03T07:27:18Z

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 · 2026-09-03T07:35:25Z

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 · 2026-09-03T08:27:23Z

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

mejrs · 2026-09-06T21:22:06Z

r? @ZuseZ4

bjorn3 · 2026-09-09T09:52:54Z

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 · 2026-09-09T11:17:38Z

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 · 2026-09-09T11:40:12Z

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

👍

Flakebi · 2026-09-09T16:38:43Z

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>): rust.godbolt.org/z/T5joGfWWE

x86 doesn’t really have a use for byref, so yeah, there it is probably just handled like byval.

The gpu-kernel ABI for amdgpu is probably the only case where byref makes sense. Not sure if there will be another GPU backend that uses it, nvptx currently uses byval – probably because ptx is just another IR and their backend does more stuff with it. I don’t know what sycl/intel uses, it’s much less mature, so wasn’t included in the clang lit tests I looked at.

byval is defined as the value being placed on the stack, which would be addrspace(5) on amdgpu and has 32-bit pointers. However, kernel arguments are not passed on the stack, but a constant memory region, which is addrspace(4) on amdgpu and has 64-bit pointers.
Therefore LLVM introduced byref, which allows non-stack address spaces to be used (+the other changes in the definition).

bjorn3 · 2026-09-09T18:38:49Z

Does the amdgpu kernel ABI pass pointers into the constant memory region as arguments for byref? Or are they at constant offsets from the start of the constant memory region? If the former, ptr addrspace(4) would work, right? If the latter, shouldn't LLVM have used something like ptr addrspace(4) byval(...)?

Flakebi · 2026-09-10T08:57:00Z

Does the amdgpu kernel ABI pass pointers into the constant memory region as arguments for byref? Or are they at constant offsets from the start of the constant memory region?

It’s… complicated, as there are no calls to gpu-kernel functions (at least no call in the LLVM IR call sense).
Instead, the user assembles an argument memory region as described before and passes a pointer and size to a vendor API (e.g. cuda or hip), and also tells the API which kernel to launch (basically by giving it the compiled GPU binary and a symbol name).
Internally, behind that API, the memory region is then uploaded to GPU memory (VRAM) and a command is sent to the GPU hardware, telling it to start n instances if the kernel at address x, giving it the memory address to the uploaded arguments.

The kernel is compiled in a way to read from that memory. There is a pass in the LLVM backend that takes the IR arguments on an amdgpu_kernel function and replaces them with argument_ptr + offset where argumpnt_ptr is the memory address x passed by the hardware that started the kernel and offset is a constant.

If the former, ptr addrspace(4) would work, right? If the latter, shouldn't LLVM have used something like ptr addrspace(4) byval(...)?

Not quite sure I understand the question completely, but I think “the latter” is the correct answer.

LLVM/clang does use ptr addrspace(4) byval(…).
And with this PR, Rust does as well. From the tests added here:

// CHECK: define amdgpu_kernel void @kernel_struct_arg(ptr addrspace(4) noalias nofree noundef readnone byref([12 x i8]) align 4 captures(none) dereferenceable(12) {{%.+}})
#[no_mangle]
pub extern "gpu-kernel" fn kernel_struct_arg(_: StructArg) {}

If it wasn’t clear, Rust should have always used byref to pass struct arguments to gpu-kernels on amdgpu. It just was broken so far.

bjorn3 · 2026-09-10T09:30:47Z

Instead, the user assembles an argument memory region as described before and passes a pointer and size to a vendor API (e.g. cuda or hip), and also tells the API which kernel to launch (basically by giving it the compiled GPU binary and a symbol name).

The exact way the user is supposed to assemble this memory region and launch the kernel is part of the calling convention.

The kernel is compiled in a way to read from that memory. There is a pass in the LLVM backend that takes the IR arguments on an amdgpu_kernel function and replaces them with argument_ptr + offset where argumpnt_ptr is the memory address x passed by the hardware that started the kernel and offset is a constant.

For byref arguments does it also load a pointer into the constant memory region from argument_ptr + offset or does it expect the argument to be stored at a fixed offset from the start of the constant memory region?

And with this PR, Rust does as well. From the tests added here:

That is byref, not byval. byval normally refers to a fixed offset from the stack pointer, but my suggestion was to use an address space modifier to make it refer to a fixed offset from the constant memory region. You could kind of treat the constant memory region as a secondary stack to pass argument, right? Except that you can only pass arguments to extern "gpu-kernel" in this way and thus there is no need to have a stack pointer for this "secondary stack" and instead a fixed address 0 can be used as value for the stack pointer.

I do understand that it isn't trivial to change in LLVM, but if it would have made sense to model it as byval with explicit address space in LLVM, perhaps we shouldn't be calling the rustc side construct ByRef, but something like InGpuConstantMemory or AmdgpuConstantMemory?

rustbot · 2026-09-10T10:00:40Z

Some changes occurred to MIR optimizations

cc @rust-lang/wg-mir-opt

Flakebi · 2026-09-10T10:01:28Z

For byref arguments does it also load a pointer into the constant memory region from argument_ptr + offset or does it expect the argument to be stored at a fixed offset from the start of the constant memory region?

It expects byref arguments to be stored at a fixed offset from the start of the constant memory region.

You could kind of treat the constant memory region as a secondary stack to pass argument, right?

Yes, apart from it being not a stack ;)
So, in rustc, we could theoretically re-purpose on_stack and just change some places to: if on_stack + gpu_kernel on amdgpu, handle this like we need to handle byref.
But I’m not sure this is a good idea given the slightly different requirements (cannot be modified, so needs to be copied in the callee).

I’m open to rename ByRef to something else. Maybe AmdgpuKernelArg?

I just pushed the change to use an enum (only the first commit changed). Mostly a mechanical change, though I added a bug! in compiler/rustc_codegen_ssa/src/mir/block.rs if byref is used in a call.
(diff: https://github.com/rust-lang/rust/compare/a90629b7e87febc01a7e6bd4c56bba2232472c05..c0e1ed53274aed32d5ee2004ae57a5311dd518f2)

bjorn3 · 2026-09-10T10:07:33Z

Maybe AmdgpuKernelArg?

Sounds fine to me.

Both will be used by the amdgpu target to implement the `gpu-kernel` ABI. `address_space` specifies the address space of an indirect argument. `AmdgpuKernelArg` 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. byval and byref are mutually exclusive, so change on_stack to an enum with the new states, Pointer (none), OnStack and AmdgpuKernelArg.

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/3a8affeef4da19d39191aac316e189eca3214a8c/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).

Flakebi · 2026-09-11T08:34:17Z

Renamed the enum variant to AmdgpuKernelArg (diff).

bjorn3 · 2026-09-11T09:12:28Z

+            match size {
+                1 => Some(Uniform::new(Reg::i8(), field.layout.size)),
+                2 => Some(Uniform::new(Reg::i16(), field.layout.size)),
+                _ => Some(Uniform::new(Reg::i32(), field.layout.size)),


Is a plain i64 also supposed to be represented as [2 x i32]?

View changes since the review

bjorn3 · 2026-09-11T09:12:46Z

+    match field.backend_repr {
+        BackendRepr::Scalar(_)
+        | BackendRepr::SimdVector { .. }
+        | BackendRepr::SimdScalableVector { .. } => {


You can probably bug!() on SimdScalableVector.

View changes since the review

bjorn3 · 2026-09-11T09:17:19Z

            continue;
        }
-        classify_arg(cx, arg);
+        if fn_abi.conv == CanonAbi::GpuKernel {


You are no longer respecting pass_indirectly_in_non_rustic_abis for extern "C". Also is this target supposed to be able to link against existing C code?

View changes since the review

Also is this target supposed to be able to link against existing C code?

No, currently not.
I started implementing towards that but that is a much larger problem than the gpu-kernel ABI and will likely require larger additions to rustc’s ABI handling (concretely, passing a repr(C) struct by value needs to be an LLVM IR struct that is passed by value and I think we need the correct types there, so using the current cast does not work).

bjorn3 · 2026-09-11T09:19:28Z

+        meta_attrs: Option<ArgAttributes>,
+        address_space: Option<AddressSpace>,
+        mode: IndirectMode,
+    },


Maybe unsized args could be split into a new PassMode::IndirectUnsized variant. Best left for a follow up PR. I can do it.

View changes since the review

rustbot assigned petrochenkov Sep 2, 2026

This comment has been minimized.

Sign in to view

bjorn3 reviewed Sep 2, 2026

View reviewed changes

Comment thread compiler/rustc_target/src/callconv/mod.rs Outdated

rustbot assigned mejrs and unassigned petrochenkov Sep 2, 2026

Flakebi force-pushed the amdgpu-kernel-cc branch from 371ff93 to b5b814c Compare September 3, 2026 07:24

This comment has been minimized.

Sign in to view

Flakebi force-pushed the amdgpu-kernel-cc branch from b5b814c to 6a47e13 Compare September 3, 2026 07:26

Flakebi force-pushed the amdgpu-kernel-cc branch from 6a47e13 to ddd9a73 Compare September 3, 2026 07:33

This comment has been minimized.

Sign in to view

Flakebi force-pushed the amdgpu-kernel-cc branch from ddd9a73 to a90629b Compare September 3, 2026 08:26

rustbot assigned ZuseZ4 and unassigned mejrs Sep 6, 2026

bjorn3 reviewed Sep 9, 2026

View reviewed changes

Comment thread compiler/rustc_codegen_ssa/src/mir/mod.rs

Flakebi force-pushed the amdgpu-kernel-cc branch from a90629b to c0e1ed5 Compare September 10, 2026 10:00

Flakebi added 3 commits September 10, 2026 23:21

Pre-commit amdgpu gpu-kernel ABI test

b142d07

Flakebi force-pushed the amdgpu-kernel-cc branch from c0e1ed5 to 552358c Compare September 11, 2026 08:32

rustbot added the F-explicit_tail_calls `#![feature(explicit_tail_calls)]` label Sep 11, 2026

bjorn3 reviewed Sep 11, 2026

View reviewed changes

Uh oh!

Uh oh!

Conversation

Flakebi commented Sep 2, 2026 • edited by rustbot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rustbot commented Sep 2, 2026

Uh oh!

rustbot commented Sep 2, 2026

Uh oh!

This comment has been minimized.

This comment has been minimized.

Uh oh!

petrochenkov commented Sep 2, 2026

Uh oh!

ZuseZ4 commented Sep 2, 2026

Uh oh!

Flakebi commented Sep 2, 2026

Uh oh!

This comment has been minimized.

Flakebi commented Sep 3, 2026

Uh oh!

Flakebi commented Sep 3, 2026

Uh oh!

This comment has been minimized.

Flakebi commented Sep 3, 2026

Uh oh!

mejrs commented Sep 6, 2026

Uh oh!

bjorn3 commented Sep 9, 2026

Uh oh!

Flakebi commented Sep 9, 2026

Uh oh!

Uh oh!

bjorn3 commented Sep 9, 2026

Uh oh!

Flakebi commented Sep 9, 2026

Uh oh!

bjorn3 commented Sep 9, 2026

Uh oh!

Flakebi commented Sep 10, 2026

Uh oh!

bjorn3 commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rustbot commented Sep 10, 2026

Uh oh!

Flakebi commented Sep 10, 2026

Uh oh!

bjorn3 commented Sep 10, 2026

Uh oh!

Flakebi commented Sep 11, 2026

Uh oh!

bjorn3 Sep 11, 2026 • edited by rustbot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bjorn3 Sep 11, 2026 • edited by rustbot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bjorn3 Sep 11, 2026 • edited by rustbot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Flakebi Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

bjorn3 Sep 11, 2026 • edited by rustbot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Flakebi commented Sep 2, 2026 •

edited by rustbot

Loading

bjorn3 commented Sep 10, 2026 •

edited

Loading

bjorn3 Sep 11, 2026 •

edited by rustbot

Loading

bjorn3 Sep 11, 2026 •

edited by rustbot

Loading

bjorn3 Sep 11, 2026 •

edited by rustbot

Loading

bjorn3 Sep 11, 2026 •

edited by rustbot

Loading