NVPTX: Allow PassMode::Direct for ptx kernels for now

[kjetilkjeka]

Upgrading the nvptx toolchain to the newest nightly makes it hit the assert that links to #115666

It seems like most targets get around this by using PassMode::Indirect. That is impossible for the kernel as it's not a normal call, but instead the arguments are copied from CPU to GPU and the passed pointer would be invalid when it reached the GPU.

I also made an experiment with PassMode::Cast but at least the most simple version of this broke the assembly API tests.

I added fixing the pass mode in my unofficial tracking issue list (I do not have the necessary permissions to update to official one). #38788 (comment)

Since the ptx_abi is currently unstable and have been working with PassMode::Direct for more than a year now, the steps above is hopefully sufficient to enable it as an exception until I can prioritize to fix it. I'm currently looking at steps to enable the CI for nvptx64 again and would prefer to finish that first.

[rustbot]

r? @cjgillot

(rustbot has picked a reviewer for you, use r? to override)

[kjetilkjeka]

@rustbot label +O-NVPTX

[RalfJung]

Since the ptx_abi is currently unstable and have been working with PassMode::Direct for more than a year now,

Have you ever tried having different-but-ABI-compatible types on the caller and callee side of these functions? "Have been working" might just mean "we didn't hit the edge case that breaks it". You can use the wrong ABI for many years without realizing; the assertion helps to find those cases before they occur in production.

I admit I have no ideas how arguments for these kernel functions work. Is there a restriction on which types can be used as arguments? What happens for large array types, or structs with tons of fields (more than there are registers)? How does one even call these functions, since it sounds like it's not a regular function call?

[workingjubilee]

@RalfJung Do we have any test cases that we expect a target to pass in order to declare its ABI "good"? It sounds like you have some "reasonably easy to test but persnickety to lower" ones in mind. Does the test infra support such?

[workingjubilee]

I admit I have no ideas how arguments for these kernel functions work. Is there a restriction on which types can be used as arguments? What happens for large array types, or structs with tons of fields (more than there are registers)? How does one even call these functions, since it sounds like it's not a regular function call?

so, like, it technically IS a regular function call
...onto a different computer.
...so, RPC, with the "serializing data over the (tiny, local) network" that entails.

@kjetilkjeka Correct me if I am wrong:

The GPU, because it not only does not have access to the CPU stack, but has no concept of a stack at the PTX level, instead provides some space that the CPU is expected to write args into? And PassMode::Direct accomplishes that write, correct?

[kjetilkjeka]

...so, RPC, with the "serializing data over the (tiny, local) network" that entails.

I think this is a good model. Additionally ptx-kernels are not allowed to have any return values. Any pointers to CPU memory will of course not be valid on the GPU. There are no limits on size of structs or unions that can be used as args.

In practice the nvidia (as the GPU vendor) will provide you with a C API to call these kernels . The arguments to the kernels is given as a double pointer which is interpreted as an array of pointers to the arguments. The driver uses the PTX description to know the length of the argument list and the size and alignment of the arguments. The driver (the part running on the CPU) will then be responsible for copying this memory into the correct GPU location before the kernels are dispatched.

The GPU, because it not only does not have access to the CPU stack, but has no concept of a stack at the PTX level, instead provides some space that the CPU is expected to write args into? And PassMode::Direct accomplishes that write, correct?

To make it even clearer, how the kernel are called (from the CPU) is not something Rust can influence directly since that is the CUDA drivers job. The PassMode::DIrect makes it correct on the callee side of things.

Let me also try to add a little context/history. This ABI used to miscompile due to calling make_indirect() even for these kernels. As I mentioned, that will result in invalid code. A while back I replaced this with different logic and added the unit tests for the kernel ABI. There exists a interoperability guide but it is not 100% correct compared to what the driver actually does. The behavior and tests are therefore based on mimicking the behavior of Nvidias cuda compiler (nvcc) and clang.

If there already exists a sensible way to make it correct without using PassingMode::Direct that would be great. If not it is better to keep the current behavior than to use something that is impossible to match from the caller side.

(note: it's also technically possible to dispatch kernels from kernels but I have not tested it in Rust and believe it's unsupported. I don't think it will influence this decision)

[workingjubilee]

Excellent.

[workingjubilee]

This only affects code using an unstable feature, in any case, which means that if the ABI is imperfect, that is to be expected, unfortunately.

Moreover, the change appears to be correct.

[RalfJung]

Do we have any test cases that we expect a target to pass in order to declare its ABI "good"? It sounds like you have some "reasonably easy to test but persnickety to lower" ones in mind. Does the test infra support such?

We have tests/ui/abi/compatibility.rs, but that test assumes that PassMode::Direct is not ever used for aggregates. (Also that test only checks for a particular kind of ABI issues. ABIs can be broken in so many wonderful ways, I have no idea how to test all of them. For instance that check cannot find #115609.) The entire rustc ABI infrastructure is based on that assumption. So the assertion that this PR proposes to disable is part of the tests which are required to ensure that the ABI is "good"; without that assertion we are basically throwing random data at LLVM and hope it works. It's a violation of the design of PassMode itself. (I'm not blaming anyone, that design was never properly documented. But that doesn't make violations of the design any less problematic.)

---

Can I call a function with this ABI like a regular Rust function? If yes, then how are arguments passed for such calls?

The driver uses the PTX description to know the length of the argument list and the size and alignment of the arguments. The driver (the part running on the CPU) will then be responsible for copying this memory into the correct GPU location before the kernels are dispatched.

Okay so that sounds like arguments are never passed in registers and everything that is memory-layout-compatible is also ABI-compatible. That means PassMode::Direct for aggregates is likely harmless. It's still a hack though and in re-designing the ABI handling (rust-lang/compiler-team#672) we should find a better way.

[RalfJung]

Can you open an issue specifically for "nvptx uses Direct for Aggregates", and link that from here and from #38788?

[kjetilkjeka]

Issue is created and linked to in code comments.

I updated my suggestion post with the new link but I do not have the permissions to update the tracking issue itself. That means someone else will have to copy over the correct link if we want it in the "head post" of the tracking issue.

I also see that @workingjubilee copied my suggestion to updated text into the tracking issue in #38788 (Thanks!), but that was still with the old link.

[RalfJung]

Okay, thanks!

This important question is still open:

Can I call a function with this ABI like a regular Rust function -- directly, without passing it to the runtime? If yes, then how are arguments passed for such calls?

[kjetilkjeka]

This important question is still open:

Can I call a function with this ABI like a regular Rust function -- directly, without passing it to the runtime? If yes, then how are arguments passed for such calls?

You're not supposed to. Well not directly at least.

There exists mechanisms for dispatching the kernel from the GPU. But this is, to my knowledge, different than calling a function. I can only find documentation on the programming model (not the ISA level). When I attempt to compile it to ptx to inspect the code I observe that it calls some undocumented ptx intrinsics. I assume this notify the driver to dispatch it in a similar way to what is done from the CPU and the same rules apply here. This behavior might of course change in newer version of the ptx isa, we're at version 8 now..

There's a bug in the rust compiler where calling a ptx-kernel function will generate invalid ptx. For the short term I think this should give an error message. For the longer term it could be interesting looking into supporting dynamic parallelism. This was previously not mentioned in the tracking issue so I added it to my update proposal.

[RalfJung]

You're not supposed to.

That wasn't my question. :)

What happens if you declare a extern "ptx-kernel" fn myfun(x: (i32, i32, i32)) { ... }, and then just call it as myfun((0, 1, 2)) directly from main, i.e. in the "host" part of the code? If this does not lead to a compilation error, and even more so if it does not require unsafe, then we have to worry about the ABI used for such calls.

Maybe the answer is "that should be a hard error but currently isn't"; in that case we need to adjust comments and issues appropriately.

[kjetilkjeka]

If you do it from the host you're on a target where this ABI doesn't make any sense and will have a compilation error on your hands

error[E0570]: `"ptx-kernel"` is not a supported ABI for the current target

If you do it from normal device code (running on the GPU, compiling with --target nvptx64-nvidia-cuda) it will trigger the miscompilation I mentioned in my previous comment. An error will currently be given from the nvidia tooling when the ptx is lowered to binary code (cubin) suitable for a concrete GPU chip. This is not the desired behavior and at least before stabilizing this ABI it needs to at least give a proper compiler error.

Maybe the answer is "that should be a hard error but currently isn't"; in that case we need to adjust comments and issues appropriately.

Spot on for the device code. But this MR doesn't change any behavior related to this problem. Is it sufficient to flag it in the tracking issue (like I already did) or does this MR need to change anything related to this as well?

[RalfJung]

will have a compilation error

Great!

Yeah then I think I can see why this current "just works" and adding the exception seems fine. Thanks for the explanations!

@bors r=workingjubilee,RalfJung rollup

[bors]

📌 Commit bb45c81 has been approved by workingjubilee,RalfJung

It is now in the queue for this repository.

[kjetilkjeka]

Thanks to both of you for reviewing 🙂