[AArch64] Use NEON's tbl1 for 16xi8 and 8xi8 build vector with mask.

[lawben]

Summary

When using Clang's __builtin_shufflevector with a 16xi8 or 8xi8 source and runtime mask on an AArch64 target, LLVM currently generates 16 or 8 extract+and+insert operations. This patch replaces these inserts with (a vector AND +) NEON's tbl1 intruction.

Issue: https://github.com/llvm/llvm-project/issues/60515

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10218

Why limit this it v16i8?

10219

It looks like we are missing test coverage for most (all?) exits. It would be great if you could extend the test coverage.

[lawben]

Add tests for exit cases and extend to v8i8.

[lawben]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10218

Good point. I've added v8i8 too. I only had the v16i8 use case in mind when writing this. This also works for other vector sizes, but that requires some additional logic to multiply and and byte offsets, as TBL only supports v16i8 and v8i8 directly. I've seen some of that around the code and I think x86 instruction selection also does this. If you think this is valuable and the right place to do it, I'd be happy to implement that in a follow-up patch.

10219

I've added a few tests to cover the exits. Please let me know if you think something is missing. I also noticed that some of my assumptions were too pessimistic in the process, so I removed them.

[dmgreen]

Hi - I like the idea of generating tbl instructions from non-constant masks. Does the patch need an update to be based against main?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10218

This can test against the EVT directly: VT != MVT::v16i8 && .. .
Other types could be handled too, by breaking the larger type up into i8 lanes. i.e a i32 would add [idx*4,idx*4+1,idx*4+2,idx*4+3] for example, if you get my meaning.
I'm not sure if that is worth adding to this patch, or handling in another but it would be good to keep in mind.

10233–10234

This doesn't need the cast any more.

10260

Should something check that the inner extract element (MaskSource) has an index equal to i?

10286

You can concat with undef if all the needed elements are in the bottom half.

llvm/test/CodeGen/AArch64/neon-shuffle-vector-tbl.ll

12

It would be good to have a test case without this And mask, just using %mask directly.

[lawben]

Add a path without AND mask.

We now have a path that does not need an AND on the mask source. We now either expect all elements to have an AND or none.
Also addresses some review comments.

[lawben]

@dmgreen Thanks for the feedback. I've addressed the comments you made. I've renamed the method, as it does not require an AND anymore.

re: update to base against main. I tried to merge the current main branch into another patch (using git) and completely broke my setup, i.e., arc is doing really weird things, trying to lint hundreds of files, etc. So I'm not sure what the correct way is to do this. Do you maybe have a pointer for me? This is my first patch, so I'm new to this setup.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10218

I also had this in mind but was thinking about putting it in another patch. v16i8 and v8i8 directly map, so they are easier to generate. Other sizes have some more overhead, so it makes sense to split them up. I've seen the logic for this in X86 instruction selection, so I guess that could just be reused. I'll do that in a follow-up, once this is merged.

10260

Yes, good catch. I've added a check and a test.

llvm/test/CodeGen/AArch64/neon-shuffle-vector-tbl.ll

12

Makes sense. There is actually no reason to limit this to and-masked masks. So I added a path that either expects all elements to have an AND mask or none. If there are no ANDs, we can omit it and directly emit a tbl. I was looking at this from Clang's __builtin_shufflevector() (as this is my use) and that always adds an AND.