[NVPTX] Incomplete and inconsistent mechanisms for lowering vectors loads/stores with sub-32-bit values

[dakersnar]

I started this discussion here: #67322 (comment), and it was requested that I open a new issue. Let me know if there is anything I need to adjust (labels, formatting, etc), I'm relatively new to open source LLVM.

CC @Artem-B, here is a simple repro that demonstrates the issue: https://godbolt.org/z/no5fEGo3h.

There are two major pieces to this issue that I'm hoping to solve:

1. Support for efficiently lowering load/store v8i8 and store v16i8 is missing from the current backend. These vectors currently get split into two or four scalar loads/stores.
2. Existing lowering for sub-32-bit load/store vectors is handled inconsistently, load v16xi8 is handled by a DagCombine ([NVPTX] Preserve v16i8 vector loads when legalizing #67322), while lowering for `load/store v8x[i16,f16,bf16] is handled via the type legalizer (https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp#L6322).

Given that, I think the goal should be to unify the approaches into a single mechanism, whichever of the two we prefer, and extend it to support the missing vector types.

A clear demonstration of a valid use case, and a patch demonstrating improved code generation would be welcome, and we can discuss the details there. Without actually implementing a prototype and seeing how clean/messy it is, it's hard to tell ahead of time what would work best here.

I actually have two very rough prototypes working with both the DagCombine approach and the Type Legalizer approach. The Dag Combine approach is very slightly cleaner, and the generated PTX is near identical in both cases (with the only real difference being whether the generated vector is made up of u32s or b32s). However, it generates a slightly different SelectionDAG pattern post type-legalization, so it would represent a bit more churn if we were to transition the existing 16-bit vector handling over to it. I was hoping to lean on @Artem-B's experience in this space to see if you had any gut feelings for which approach would be preferred, and also whether you share my thinking that unifying them is a good idea at all.

As a quick example, here is the post-type-legalizer diff of the selectiondag for the approaches, when the input is the "generic_8xi16" function from the compiler explorer link.

Left is with the current Type Legalizer handling and right is with the DagCombine handling (I just changed it for the load, the store is untouched). The PTX ends up near-identical in this simple case, but I'm not sure whether the left side is enabling specific optimizations for v2x[i/f/bf]16 in other cases, as I know there was a lot of work done in that space that I haven't dug super deep into yet, e.g. 620db1f#diff-6fda74ef428299644e9f49a2b0994c0d850a760b89828f655030a114060d075a

Anyway, I can put up a patch for one of the approaches if that would be helpful, or possibly both, if you don't lean in one direction. I'll also note that the Dag Combine approach could be rewritten to match the output of the Type Legalizer approach, and vice versa, so there might be two questions at play here: 1) where do we prefer this lives and 2) what do we prefer the type-legalizer output to look like.

[Artem-B]

CC @Artem-B, here is a simple repro that demonstrates the issue: https://godbolt.org/z/no5fEGo3h.

Here's a somewhat better example with alignment info added on loads/stores: https://godbolt.org/z/v9sve64Yj

There are two major pieces to this issue that I'm hoping to solve:

Support for efficiently lowering load/store v8i8 and store v16i8 is missing from the current backend. These vectors currently get split into two or four scalar loads/stores.

v8i8 is indeed split into two 32-bit loads.

        ld.param.u64    %rd1, [generic_8xi8_param_0];
        ld.u32  %r1, [%rd1];
        ld.u32  %r2, [%rd1+4];

v16i8 is fine, once alignment is set correctly:

        ld.param.u64    %rd1, [generic_16xi8_param_0];
        ld.v4.u32       {%r1, %r2, %r3, %r4}, [%rd1];

Existing lowering for sub-32-bit load/store vectors is handled inconsistently, load v16xi8 is handled by a DagCombine (#67322), while lowering for `load/store v8x[i16,f16,bf16] is handled via the type legalizer (https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp#L6322).

There's more than one way to do it.

The easiest way would probably be to extend NVPTXISelLowering.cpp where we already deal with all the load/store quirkiness.
Doung it all in DagCombiner is possible, but then we'd have to move all existing vectorized load/store handling there, and Im not sure if that will buy us much.

[dakersnar]

v16i8 is fine, once alignment is set correctly:

load v16i8 is actually already fine in my link, with or without the alignment, due to the DagCombine optimization:

ld.v4.u32 {%r1, %r2, %r3, %r4}, [%rd1];

However, even in your link, the store v16i8 is currently not being handled well, because the store version of the optimization was reverted due to an alignment issue (#73646 (review)):

        st.u32  [%rd1+12], %r52;
        st.u32  [%rd1+8], %r41;
        st.u32  [%rd1+4], %r30;
        st.u32  [%rd1], %r19;

The easiest way would probably be to extend NVPTXISelLowering.cpp where we already deal with all the load/store quirkiness.
Doung it all in DagCombiner is possible, but then we'd have to move all existing vectorized load/store handling there, and Im not sure if that will buy us much.

To clarify, both existing versions of this handling live in NVPTXISelLowering.cpp, but when I refer to the "Dag Combine approach" I'm talking about PerformLOADCombine which gets called from NVPTXTargetLowering::PerformDAGCombine in the DagCombine phase, whereas the "type legalzier approach" is referring to ReplaceLoadVector and LowerSTOREVector, which get called from NVPTXTargetLowering::LowerOperation and NVPTXTargetLowering::ReplaceNodeResults which get called in the Type Legalization phase. Does that make sense?

Given that, are you suggesting we extend the "type legalizer approach" with v8i8 and v16i8, deleting PerformLOADCombine? I think that's a reasonable call to minimize churn, and I can put up a PR to demonstrate what that would look like.

[dakersnar]

load v16i8 is actually already fine in my link, with or without the alignment, due to the DagCombine optimization

ld.v4.u32 {%r1, %r2, %r3, %r4}, [%rd1];

Side note, is this correct behavior to lower it as a vector when the alignment is not specified? That's the current behavior, you can see the differences with no align, align 8, and align 16: https://godbolt.org/z/Pr46Ks4a6. Only align 8 gets lowered as scalars.

Edit: I refreshed my memory on align, and yes, this is correct. v4i32 also behaves the same way: https://godbolt.org/z/Phq75ehrb

[Artem-B]

the store v16i8 is currently not being handled well, because the store version of the optimization was reverted due to an alignment issue (#73646 (review)):

        st.u32  [%rd1+12], %r52;
        st.u32  [%rd1+8], %r41;
        st.u32  [%rd1+4], %r30;
        st.u32  [%rd1], %r19;

This should indeed be fixed.

As it happens LowerSTOREVector currently does not handle v8i8 or v16i8 and thus we end up with suboptimal stores. Previously the problem was that we didn't have register-sized i8 vector type we could use the same way we use (e.g. how we use f16x2 x 4 for lowering v8f16 loads/stores), but right now we should be able to do it via v4i8 so we'll do loads/stores as i32/v4i8 which can already handle now.

when I refer to the "Dag Combine approach" I'm talking about PerformLOADCombine which gets called from NVPTXTargetLowering::PerformDAGCombine in the DagCombine phase,

That's exactly what I mean that you can massage the graph into a shape that could be lowered to a better set of instructions, but if you do it here, you're replicating existing functionality we have in the LowerSTOREVector.
Conceptually both approaches are fine, I just want to keep the magic in one place, if we can.

is this correct behavior to lower it as a vector when the alignment is not specified?

LLVM IR spec says "An omitted align argument means that the operation has the ABI alignment for the target.", so whatever is the expected alignment for the type we used. In other words, yes, I believe we're OK to use vectorized loads/stores in this case.

you can see the differences with no align, align 8, and align 16: https://godbolt.org/z/Pr46Ks4a6. Only align 8 gets lowered as scalars.

We do have gaps. i8 vectors were rarely used until fairly recently -- they are almost never generated from CUDA sources, and for IR-based JIT, the generated IR often already has nicely structured loads/stores. Over the years we've patched things piecemeal, to address specific use cases, but the upside of general solution is fairly limited -- the best we can do is a 128-bit load/store with minimum overhead for conversion between types. In practice, I've only seen folks complaining about, essentially, copy-only kernels, that were copying using long vi8 types, and ended up with byte-wize loads/stores, or with vi32 loads/stores (good) and bunch of split into bytes/recombine into i32, then store nonsense. If we do need to split i32 into bytes, the overhead of doing that would probably eat all the benefits of loading all the data using ld.v4.i8 32 bits at a time, as we will no longer be bottlenecked by how fast we can issue load/store instructions.

I guess, what I'm saying is that overhauling the vectorized load/store lowering is probably not worth the effort. Filling in the few remaining gaps would be sufficient.

That said, if you have time to re-do it -- go for it. We've certainly accumulated a lot of organic growth over the years that could use cleaning up. That would be welcome, too.

[dakersnar]

Previously the problem was that we didn't have register-sized i8 vector type we could use the same way we use (e.g. how we use f16x2 x 4 for lowering v8f16 loads/stores), but right now we should be able to do it via v4i8 so we'll do loads/stores as i32/v4i8 which can already handle now.

This is music to my ears; it means I've stared at the codebase and git history enough this week to understand why things are the way they are 😀.

Conceptually both approaches are fine, I just want to keep the magic in one place, if we can.

Yep, I'm 100% with you there! With v4i8 existing now, we should be able to keep everything in one place.

If we do need to split i32 into bytes, the overhead of doing that would probably eat all the benefits of loading all the data using ld.v4.i8 32 bits at a time, as we will no longer be bottlenecked by how fast we can issue load/store instructions.

Fair enough, I believe you that this is potentially true, although I think it is outside of the scope of what I'm proposing. If we find that the heuristic of "minimize the raw number of loads/stores as much as possible" is not a good idea in these edge cases where 8-bit data is being extracted and manipulated, maybe that's more of a LoadStoreVectorizer problem to intelligently avoid generating those vectors in the first place.

I guess, what I'm saying is that overhauling the vectorized load/store lowering is probably not worth the effort. Filling in the few remaining gaps would be sufficient.

I agree, I think the solution I'm working on fills in the gaps and cleans things up without needing to re-do everything. I'll share a PR when it is ready in the next few days, but I definitely feel more confident in the direction after this back and forth, thank you!

[dakersnar]

Shall I close this issue?