[llvm] [AArch64][SVE] Support lowering masked loads/stores of <4 x bf16> and <8 x bf16> (PR #208744)

[Shanzhi Chen via llvm-commits]

chenshanzhi wrote:

Notes for review:
1. For masked store of <4 x bf16> and  <8 x bf16>, `Action` in `SelectionDAGLegalize::LegalizeOp` is `Expand` rather than `Custom`, which makes `ISD::MSTORE` of `MVT::v4bf16` / `MVT::v8bf16` is not lowered.
2. After adding `setOperationAction` for `ISD::MLOAD`/`ISD::MSTORE` of `MVT::v4bf16` / `MVT::v8bf16` in `AArch64TargetLowering::AArch64TargetLowering`, the masked stores are lowered correctly, but the masked loads are not. This is because the different checks in `LowerMLOAD` and `LowerMSTORE`:

In `LowerMLOAD`:
```c++
  if (useSVEForFixedLengthVectorVT(VT, /*OverrideNEON=*/true))
    return LowerFixedLengthVectorMLoadToSVE(Op, DAG);
```
In `LowerMSTORE`:
```c++
  if (VT.isFixedLengthVector())
    return LowerFixedLengthVectorMStoreToSVE(Op, DAG);
```
And `useSVEForFixedLengthVectorVT` returns false for all vectors of bf16.

3. Simply add `bf16` among `f16` in the `switch` statement of `useSVEForFixedLengthVectorVT` would cause a lot of test failures. This is firstly because vectors of bf16 would be passed to `addTypeForFixedLengthSVE` and the following code in this function does not handle vectors of bf16 properly.
```c++
  // Mark floating-point truncating stores/extending loads as having custom
  // lowering
  if (VT.isFloatingPoint()) {
    MVT InnerVT = VT.changeVectorElementType(MVT::f16);
    while (InnerVT != VT) {
      setTruncStoreAction(VT, InnerVT, Custom);
      setLoadExtAction(ISD::EXTLOAD, VT, InnerVT, Default);
      InnerVT = InnerVT.changeVectorElementType(
          MVT::getFloatingPointVT(2 * InnerVT.getScalarSizeInBits()));
    }
  }
```
4. After handlin vectors of bf16 in `addTypeForFixedLengthSVE`, there're still 5 failed tests.
```
Failed Tests (5):
  LLVM :: CodeGen/AArch64/sve-fixed-length-extract-subvector.ll
  LLVM :: CodeGen/AArch64/sve-fixed-length-fp-compares.ll
  LLVM :: CodeGen/AArch64/sve-fixed-length-splat-segment.ll
  LLVM :: CodeGen/AArch64/sve-indexed-arithmetic.ll
  LLVM :: CodeGen/AArch64/sve2p1-vector-shuffles.ll
```
After some analysis, this PR choose to only process the `OverrideNEON` cases for v4bf16/v8bf16 in `useSVEForFixedLengthVectorVT`. This is because not all SVE operations are supported for vectors of bf16. For example, comparing v16bf16 cannot be lowered directly to a machine. This is why `CodeGen/AArch64/sve-fixed-length-fp-compares.ll` fails.

5. Another solution is keeping `useSVEForFixedLengthVectorVT` unchanged and adding more checks for v4bf16 and v8bf16 in this `if` in `LowerMLOAD`. But I think that would be a little customized.
```c++
  if (useSVEForFixedLengthVectorVT(VT, /*OverrideNEON=*/true))
    return LowerFixedLengthVectorMLoadToSVE(Op, DAG);
```

https://github.com/llvm/llvm-project/pull/208744