[AArch64][CodeGen] Restrict bfloat vector operations to what's actually supported

[stuij]

Summary

Previously in addTypeForNeon, we would set the operations for bfloat vectors
like other generic types. But as bfloat is a storage-only type a number of
operations shouldn't be set. This patch fixes that.

[dmgreen]

Quick question - what is the expected behaviour? Do we just never expect to see an bf16 add, and if we do it's a fatal error? Or is some form of automatic promotion expected to happen?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1021

These are integer, so I wouldn't expect then to cause much trouble.

1045

Why not leave these as expand?

[dnsampaio]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1045

And they are integers as well, no? Perhaps the correct tests we want here is !VT.isFloatingPoint() or alike?

[stuij]

In D85101#2190026, @dmgreen wrote:

Quick question - what is the expected behaviour? Do we just never expect to see an bf16 add, and if we do it's a fatal error? Or is some form of automatic promotion expected to happen?

Yes, I don't think we should see explicit bfloat adds. It's a storage type, and addition might not be what the user intended. So we should error.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1045

As I understand, as bfloat is a storage type, it should be converted explicitly to a different type if you want to do arithmetic operations. It makes sense to me to explicitly not allow this.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1045

The frem is the only one in this list that would make sense for bf16. I believe the default would be Legal(?) So Expand seems like a sensible enough option. At least, it would seem a little odd to specifically mark them differently for bf16. We don't do anything special for fadd/fdiv/etc as far as I understand.

[stuij]

addressing review comments

[dmgreen]

Thanks. Sounds good to me.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 //===-- AArch64ISelLowering.cpp - AArch64 DAG Lowering Implementation  ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (Subtarget->hasNEON()) {
     // Someone set us up the NEON.
     addDRTypeForNEON(MVT::v2f32);
     addDRTypeForNEON(MVT::v8i8);
     addDRTypeForNEON(MVT::v4i16);
     addDRTypeForNEON(MVT::v2i32);
     addDRTypeForNEON(MVT::v1i64);
     addDRTypeForNEON(MVT::v1f64);
     addDRTypeForNEON(MVT::v4f16);
+    if (Subtarget->hasBF16())
       addDRTypeForNEON(MVT::v4bf16);
 
     addQRTypeForNEON(MVT::v4f32);
     addQRTypeForNEON(MVT::v2f64);
     addQRTypeForNEON(MVT::v16i8);
     addQRTypeForNEON(MVT::v8i16);
     addQRTypeForNEON(MVT::v4i32);
     addQRTypeForNEON(MVT::v2i64);
     addQRTypeForNEON(MVT::v8f16);
+    if (Subtarget->hasBF16())
       addQRTypeForNEON(MVT::v8bf16);
   }
 
   if (Subtarget->hasSVE()) {
     // Add legal sve predicate types
     addRegisterClass(MVT::nxv2i1, &AArch64::PPRRegClass);
     addRegisterClass(MVT::nxv4i1, &AArch64::PPRRegClass);
     addRegisterClass(MVT::nxv8i1, &AArch64::PPRRegClass);
     addRegisterClass(MVT::nxv16i1, &AArch64::PPRRegClass);
if (VT == MVT::v2f32 || VT == MVT::v4f32 || VT == MVT::v2f64) {
     setOperationAction(ISD::FCOPYSIGN, VT, Custom);
   }
 
   setOperationAction(ISD::EXTRACT_VECTOR_ELT, VT, Custom);
   setOperationAction(ISD::INSERT_VECTOR_ELT, VT, Custom);
   setOperationAction(ISD::BUILD_VECTOR, VT, Custom);
   setOperationAction(ISD::VECTOR_SHUFFLE, VT, Custom);
   setOperationAction(ISD::EXTRACT_SUBVECTOR, VT, Custom);
+
+  if (VT.getVectorElementType() != MVT::bf16) {
     setOperationAction(ISD::SRA, VT, Custom);
     setOperationAction(ISD::SRL, VT, Custom);
     setOperationAction(ISD::SHL, VT, Custom);
     setOperationAction(ISD::OR, VT, Custom);
+  }
+
   setOperationAction(ISD::SETCC, VT, Custom);
   setOperationAction(ISD::CONCAT_VECTORS, VT, Legal);
 
   setOperationAction(ISD::SELECT, VT, Expand);
   setOperationAction(ISD::SELECT_CC, VT, Expand);
   setOperationAction(ISD::VSELECT, VT, Expand);
   for (MVT InnerVT : MVT::all_valuetypes())
     setLoadExtAction(ISD::EXTLOAD, InnerVT, VT, Expand);
 
   // CNT supports only B element sizes, then use UADDLP to widen.
   if (VT != MVT::v8i8 && VT != MVT::v16i8)
     setOperationAction(ISD::CTPOP, VT, Custom);
 
+  if (VT.getVectorElementType() != MVT::bf16) {
     setOperationAction(ISD::UDIV, VT, Expand);
     setOperationAction(ISD::SDIV, VT, Expand);
     setOperationAction(ISD::UREM, VT, Expand);
     setOperationAction(ISD::SREM, VT, Expand);
     setOperationAction(ISD::FREM, VT, Expand);
+  }
 
   setOperationAction(ISD::FP_TO_SINT, VT, Custom);
   setOperationAction(ISD::FP_TO_UINT, VT, Custom);
 
   if (!VT.isFloatingPoint())
     setOperationAction(ISD::ABS, VT, Legal);
 
   // [SU][MIN|MAX] are available for all NEON types apart from i64.
   if (!VT.isFloatingPoint() && VT != MVT::v2i64 && VT != MVT::v1i64)
     for (unsigned Opcode : {ISD::SMIN, ISD::SMAX, ISD::UMIN, ISD::UMAX})
       setOperationAction(Opcode, VT, Legal);
 
   // F[MIN|MAX][NUM|NAN] are available for all FP NEON types.
   if (VT.isFloatingPoint() &&
+      VT.getVectorElementType() != MVT::bf16 &&
       (VT.getVectorElementType() != MVT::f16 || Subtarget->hasFullFP16()))
     for (unsigned Opcode :
          {ISD::FMINIMUM, ISD::FMAXIMUM, ISD::FMINNUM, ISD::FMAXNUM})
       setOperationAction(Opcode, VT, Legal);
 
   if (Subtarget->isLittleEndian()) {
     for (unsigned im = (unsigned)ISD::PRE_INC;
          im != (unsigned)ISD::LAST_INDEXED_MODE; ++im) {

llvm/test/CodeGen/AArch64/bf16-vector-bitcast.ll

-; RUN: llc < %s -asm-verbose=0 -mtriple=aarch64-none-eabi | FileCheck %s
+; RUN: llc < %s -asm-verbose=0 -mtriple=aarch64-none-eabi -mattr=+bf16 | FileCheck %s
 
 define <4 x i16> @v4bf16_to_v4i16(float, <4 x bfloat> %a) nounwind {
 ; CHECK-LABEL: v4bf16_to_v4i16:
 ; CHECK-NEXT: mov v0.16b, v1.16b
 ; CHECK-NEXT: ret
 entry:
   %1 = bitcast <4 x bfloat> %a to <4 x i16>
   ret <4 x i16> %1

llvm/test/CodeGen/AArch64/bf16-vector-shuffle.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc < %s -asm-verbose=0 -mtriple=aarch64-none-eabi | FileCheck %s
+; RUN: llc < %s -asm-verbose=0 -mtriple=aarch64-none-eabi -mattr=+bf16 | FileCheck %s
 
 ; bfloat16x4_t test_vcreate_bf16(uint64_t a) { return vcreate_bf16(a); }
 define <4 x bfloat> @test_vcreate_bf16(i64 %a) nounwind {
 ; CHECK-LABEL: test_vcreate_bf16:
 ; CHECK-NEXT:    fmov d0, x0
 ; CHECK-NEXT:    ret
 entry:
   %0 = bitcast i64 %a to <4 x bfloat>