nvptx abi adjustment may keep PassMode::Direct for small aggregates in conv extern "C"

[kjetilkjeka]

          @kjetilkjeka I had to disable this test on nvptx64 since the `extern "C"` ABI uses `Direct` pass mode in invalid ways. I think that's caused by this code:

rust/compiler/rustc_target/src/abi/call/nvptx64.rs

Lines 10 to 14 in b75b3b3

	fn classify_arg<Ty>(arg: &mut ArgAbi<'_, Ty>) {
	if arg.layout.is_aggregate() && arg.layout.size.bits() > 64 {
	arg.make_indirect();
	}
	}

That's invalid since it doesn't say what to do for smaller aggregates, and they default to the (bad) Direct. Instead you have to say which register they are supposed to be passed in. You can check what the other targets are doing. Targets are expected to set an explicit pass mode for all aggregate arguments.

This should be easy to reproduce by having a function like

#[repr(C)]
struct ReprC1<T: ?Sized>(T);

extern "C" fn myfn(x: ReprC1<i32>) {}

which will likely ICE on the current compiler already.

Originally posted by @RalfJung in #117351 (comment)

[kjetilkjeka]

@rustbot label +O-NVPTX

[kjetilkjeka]

Thanks for pointing this out @RalfJung

As a note I had to use #[no_mangle] to avoid the compiler from removing the symbol before the assert was hit.

Do you know why this test started failing with your changes in #117351 and not before that? The reason I'm asking is that I'm curious to how to make the tests catch it if we have a regression and start producing PassMode::Direct after this issue is resolved

[RalfJung]

Do you know why this test started failing with your changes in #117351 and not before that?

We used to only run the assertion in codegen, i.e. when a function with such argument types actually gets generated or called. The abi compatibility test works on function pointers, it computes their ABI but never calls them so we never enter the part of the codegen backend that had the assertion.

I'm going to try to find a way to move this check to the ABI computation logic, but that turns out to be hard since we do compute some ABIs that we don't actually support...

[RalfJung]

I'm curious to how to make the tests catch it if we have a regression and start producing PassMode::Direct after this issue is resolved

#117500 does ensure we'd catch this.

[RalfJung]

This should not have been closed without enabling the relevant tests for nvtpx:
#125694.

[RalfJung]

Actually the issue is still open -- the test fails:

`PassMode::Direct` for aggregates only allowed for "unadjusted" and "ptx-kernel" functions and on wasm
Problematic type: TyAndLayout {
    ty: ReprC1<*const i32>,
    layout: Layout {
        size: Size(8 bytes),
        align: AbiAndPrefAlign {
            abi: Align(8 bytes),
            pref: Align(8 bytes),
        },
        abi: Aggregate {
            sized: true,
        },
        fields: Arbitrary {
            offsets: [
                Size(0 bytes),
            ],
            memory_index: [
                0,
            ],
        },
        largest_niche: None,
        variants: Single {
            index: 0,
        },
        max_repr_align: None,
        unadjusted_abi_align: Align(8 bytes),
    },
}

[RalfJung]

#117671 only changed classify_arg, shouldn't it change classify_ret as well?

[kjetilkjeka]

This comment explains the problem but is essentially this.

The only two PassModes that give the correct signature is Direct and Indirect{on_stack: true}. For arguments it's possible to use Indirect{on_stack: true} but that is not the case for returns.

I unfortunately don't think it's possible to have a correct C abi for ptx without accepting Direct or implementing the new ABI handling.

[RalfJung]

Hm, I only understand half of that comment.^^ I cannot parse the PTX signature. What does it all mean? How does the PTX signature look like when the struct has multiple fields with padding between them?

I also don't know much about on_stack. The fact that that is part of Indirect makes no sense to me as there's not actually a pointer being passed -- this is probably some LLVM stuff leaking through. Is there a way to tell LLVM "return this as a blob of N bytes" that then gets the right signature in PTX?

This is all completely terrible. LLVM's types are already ill-suited to express function call ABIs so LLVM backends have to use heuristics to reconstruct the actual ABI (in particular for targets like PTX that aren't actually assembly but still IRs), and then in Rust we use the PassMode to describe which LLVM types to generate but we can't actually describe all of them, urgh.

That said... I am a bit surprised that Cast doesn't work. You say

PassMode::Cast casted to an uniform does the correct thing for many types, but unfortunately not for single element structs (it passes like it was the underlying type, not the struct itself).

what exactly does that look like? cast will return things in an LLVM struct if prefix contains something, so I would expect that for the struct with a u8 field, you want Cast where prefix[0] is Some(Reg { kind: integer, size: 1 }) and everything else is None. (If you only set the rest it will use an array, or the underlying type for single-element arrays.

For more complicated types this does sound quite tricky to do correctly, though.

[kjetilkjeka]

Hm, I only understand half of that comment.^^ I cannot parse the PTX signature. What does it all mean? How does the PTX signature look like when the struct has multiple fields with padding between them?

Ptx is an intermediate format where passing something is basically just listing a named argument of the type it is passed as. The C ABI in ptx is will list a byte array of elements with a specific alignment requirement when passing structs.

When passing struct SingleU8{a: u8} it is listed as .param .align 1 .b8 <name>[1]

Primitive integer types are mainly passed as the actual type they are, except when they are of less size than 32 bits. All integers of less than 32bits will be padded to a 32bit integer. This is somewhat unfortunate, but also just something we have to deal with.

When passing a u8 it will then be listed as .param .b32 <name>. No explicit alignment is required as the datatype alignment is used.

what exactly does that look like? cast will return things in an LLVM struct if prefix contains something, so I would expect that for the struct with a u8 field, you want Cast where prefix[0] is Some(Reg { kind: integer, size: 1 }) and everything else is None. (If you only set the rest it will use an array, or the underlying type for single-element arrays.

When I try to cast to a single 8-bit register it passes it as a u8 instead of the SingleU8. That is, as you say, due to it passing it as the underlying type of single-element arrays.

---

I also don't know much about on_stack. The fact that that is part of Indirect makes no sense to me as there's not actually a pointer being passed -- this is probably some LLVM stuff leaking through. Is there a way to tell LLVM "return this as a blob of N bytes" that then gets the right signature in PTX?

I have had great problems with figuring out how Indirect and on_stack fits together. My partial conclusion was that when on_stack is set it's really not Indirect anymore. I would have to investigate more to understand how it maps to LLVM.

I think my next point of action in this task will be to look into the LLVM output from clang when compiling these functions. Perhaps that can guide what we need to output to LLVM from Rust. I guess a problem arise if that is exactly what we do generate when using PassMode::Direct. I wouldn't be entirely surprised if that was the case.

[RalfJung]

.param .align 1 .b8 <name>[1]

So if this was a struct with an i32 field it would be [4] at the end to indicate an array of length 4?

When I try to cast to a single 8-bit register it passes it as a u8 instead of the SingleU8. That is, as you say, due to it passing it as the underlying type of single-element arrays.

How are you setting up the cast?
When the rest is a u8 then it will indeed become a single u8. But when one prefix element is a u8, that should become a struct for LLVM, which is exactly the same as with Direct.

[kjetilkjeka]

So if this was a struct with an i32 field it would be [4] at the end to indicate an array of length 4?

Yes, this is correct.

How are you setting up the cast?
When the rest is a u8 then it will indeed become a single u8. But when one prefix element is a u8, that should become a struct for LLVM, which is exactly the same as with Direct.

Ahh, that is not intuitive at all for me, but it works indeed. I made the return tests pass with this code.

   if ret.layout.is_aggregate() {
        let align_bytes = ret.layout.align.abi.bytes();
        let size = ret.layout.size;

        let reg = match align_bytes {
            1 => Reg::i8(),
            2 => Reg::i16(),
            4 => Reg::i32(),
            8 => Reg::i64(),
            16 => Reg::i128(),
            _ => unreachable!("Align is given as power of 2 no larger than 16 bytes"),
        };

        if align_bytes == size.bytes() {
            ret.cast_to(CastTarget {
                prefix: [Some(reg), None, None, None, None, None, None, None],
                rest: Uniform::new(Reg::i8(), Size::from_bytes(0)),
                attrs: ArgAttributes {
                    regular: ArgAttribute::default(),
                    arg_ext: ArgExtension::None,
                    pointee_size: Size::ZERO,
                    pointee_align: None,
                },
            });
        } else {
            ret.cast_to(Uniform::new(reg, size));
        }
    } else if ret.layout.size.bits() < 32 {
        ret.extend_integer_width_to(32);
    }

I will most likely get a chance to create a PR on Friday. Thanks for that trick.

[RalfJung]

Yeah to be clear I have no idea how CastTarget is intended to be used, I just read the code for how it gets translated into an LLVM type, and then found a way to make that generate the LLVM type you want, for the one case you mentioned. I have no idea how to turn that into a general principle that works for all types.

You seem to make a special case for align_bytes == size.bytes(), but doesn't that mean that for a type with two i32 fields we'll get the wrong encoding? Specifically, currently that will be { i32, i32 } (a struct with two i32 fields), with your patch that will be [2 x i32] (a two-element array of i32). Maybe those generate the same PTX output, maybe not -- please carefully check that.

[kjetilkjeka]

You seem to make a special case for align_bytes == size.bytes(), but doesn't that mean that for a type with two i32 fields we'll get the wrong encoding? Specifically, currently that will be { i32, i32 } (a struct with two i32 fields), with your patch that will be [2 x i32] (a two-element array of i32). Maybe those generate the same PTX output, maybe not -- please carefully check that.

I have specifically added a test for this type in #125980. It seems to produce correct results.