[AArch64] Add some basic handling for bf16 constants.

[dmgreen]

Summary

This adds some basic handling for bf16 constants, attempting to treat them a lot like fp16 constants where it can. Zero immediates get lowered to FMOVH0, others either get lowered to FMOVWHr(MOVi32imm) or use FMOVHi if they can. Without fp16 they get expanded. This may not always be optimal, but fixes a gap in our lowering. See https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AArch64/f16-imm.ll for the equivalent fp16 test.

[stuij]

Yea nice. Looks pretty straightforward to me.

llvm/lib/CodeGen/TargetLoweringBase.cpp

#include "llvm/IR/VPIntrinsics.def"
 
   // Most targets also ignore the @llvm.readcyclecounter intrinsic.
   setOperationAction(ISD::READCYCLECOUNTER, MVT::i64, Expand);
 
   // ConstantFP nodes default to expand.  Targets can either change this to
   // Legal, in which case all fp constants are legal, or use isFPImmLegal()
   // to optimize expansions for certain constants.
   setOperationAction(ISD::ConstantFP,
-                     {MVT::f16, MVT::f32, MVT::f64, MVT::f80, MVT::f128},
+                     {MVT::bf16, MVT::f16, MVT::f32, MVT::f64, MVT::f80, MVT::f128},
                      Expand);
 
   // These library functions default to expand.
   setOperationAction({ISD::FCBRT, ISD::FLOG, ISD::FLOG2, ISD::FLOG10, ISD::FEXP,
                       ISD::FEXP2, ISD::FFLOOR, ISD::FNEARBYINT, ISD::FCEIL,
                       ISD::FRINT, ISD::FTRUNC, ISD::LROUND, ISD::LLROUND,
                       ISD::LRINT, ISD::LLRINT},
                      {MVT::f32, MVT::f64, MVT::f128}, Expand);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasNEON()) {
     // conversion happens in two steps: v4i32 -> v4f32 -> v4f16
     for (auto Op : {ISD::SINT_TO_FP, ISD::UINT_TO_FP, ISD::STRICT_SINT_TO_FP,
                     ISD::STRICT_UINT_TO_FP})
       for (auto VT : {MVT::v2i32, MVT::v2i64, MVT::v4i32})
         setOperationAction(Op, VT, Custom);
 
     if (Subtarget->hasFullFP16()) {
       setOperationAction(ISD::ConstantFP, MVT::f16, Legal);
+      setOperationAction(ISD::ConstantFP, MVT::bf16, Legal);
 
       setOperationAction(ISD::SINT_TO_FP, MVT::v8i8, Custom);
       setOperationAction(ISD::UINT_TO_FP, MVT::v8i8, Custom);
       setOperationAction(ISD::SINT_TO_FP, MVT::v16i8, Custom);
       setOperationAction(ISD::UINT_TO_FP, MVT::v16i8, Custom);
       setOperationAction(ISD::SINT_TO_FP, MVT::v4i16, Custom);
       setOperationAction(ISD::UINT_TO_FP, MVT::v4i16, Custom);
       setOperationAction(ISD::SINT_TO_FP, MVT::v8i16, Custom);
bool AArch64TargetLowering::isFPImmLegal(const APFloat &Imm, EVT VT,
   // We can materialize #0.0 as fmov $Rd, XZR for 64-bit, 32-bit cases, and
   // 16-bit case when target has full fp16 support.
   // FIXME: We should be able to handle f128 as well with a clever lowering.
   const APInt ImmInt = Imm.bitcastToAPInt();
   if (VT == MVT::f64)
     IsLegal = AArch64_AM::getFP64Imm(ImmInt) != -1 || Imm.isPosZero();
   else if (VT == MVT::f32)
     IsLegal = AArch64_AM::getFP32Imm(ImmInt) != -1 || Imm.isPosZero();
-  else if (VT == MVT::f16)
+  else if (VT == MVT::f16 || VT == MVT::bf16)
     IsLegal =
         (Subtarget->hasFullFP16() && AArch64_AM::getFP16Imm(ImmInt) != -1) ||
         Imm.isPosZero();
 
   // If we can not materialize in immediate field for fmov, check if the
   // value can be encoded as the immediate operand of a logical instruction.
   // The immediate value will be created with either MOVZ, MOVN, or ORR.
   // TODO: fmov h0, w0 is also legal, however we don't have an isel pattern to

llvm/lib/Target/AArch64/AArch64InstrFormats.td

 def fpimm16 : Operand<f16>,
               FPImmLeaf<f16, [{
       return AArch64_AM::getFP16Imm(Imm) != -1;
     }], fpimm16XForm> {
   let ParserMatchClass = FPImmOperand;
   let PrintMethod = "printFPImmOperand";
 }
 
+def fpimmbf16 : Operand<bf16>,
+                FPImmLeaf<bf16, [{
+      return AArch64_AM::getFP16Imm(Imm) != -1;
+    }], fpimm16XForm>;
+
 def fpimm32 : Operand<f32>,
               FPImmLeaf<f32, [{
       return AArch64_AM::getFP32Imm(Imm) != -1;
     }], fpimm32XForm> {
   let ParserMatchClass = FPImmOperand;
   let PrintMethod = "printFPImmOperand";
 }
 

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 let isReMaterializable = 1, isCodeGenOnly = 1, isAsCheapAsAMove = 1 in {
 def FMOVH0 : Pseudo<(outs FPR16:$Rd), (ins), [(set f16:$Rd, (fpimm0))]>,
     Sched<[WriteF]>;
 def FMOVS0 : Pseudo<(outs FPR32:$Rd), (ins), [(set f32:$Rd, (fpimm0))]>,
     Sched<[WriteF]>;
 def FMOVD0 : Pseudo<(outs FPR64:$Rd), (ins), [(set f64:$Rd, (fpimm0))]>,
     Sched<[WriteF]>;
 }
+
 // Similarly add aliases
 def : InstAlias<"fmov $Rd, #0.0", (FMOVWHr FPR16:$Rd, WZR), 0>,
     Requires<[HasFullFP16]>;
 def : InstAlias<"fmov $Rd, #0.0", (FMOVWSr FPR32:$Rd, WZR), 0>;
 def : InstAlias<"fmov $Rd, #0.0", (FMOVXDr FPR64:$Rd, XZR), 0>;
 
-// Pattern for FP16 immediates
+def : Pat<(bf16 fpimm0),
+          (FMOVH0)>;
+
+// Pattern for FP16 and BF16 immediates
 let Predicates = [HasFullFP16] in {
   def : Pat<(f16 fpimm:$in),
             (FMOVWHr (MOVi32imm (bitcast_fpimm_to_i32 f16:$in)))>;
+
+  def : Pat<(bf16 fpimm:$in),
+            (FMOVWHr (MOVi32imm (bitcast_fpimm_to_i32 bf16:$in)))>;
 }
 
 //===----------------------------------------------------------------------===//
 // Floating point conversion instruction.
 //===----------------------------------------------------------------------===//
 
 defm FCVT : FPConversion<"fcvt">;
 

 //===----------------------------------------------------------------------===//
 // Floating point immediate move.
 //===----------------------------------------------------------------------===//
 
 let isReMaterializable = 1, isAsCheapAsAMove = 1 in {
 defm FMOV : FPMoveImmediate<"fmov">;
 }
 
+let Predicates = [HasFullFP16] in {
+  def : Pat<(bf16 fpimmbf16:$in),
+            (FMOVHi (fpimm16XForm bf16:$in))>;
+}
+
 //===----------------------------------------------------------------------===//
 // Advanced SIMD two vector instructions.
 //===----------------------------------------------------------------------===//
 
 defm UABDL   : SIMDLongThreeVectorBHSabdl<1, 0b0111, "uabdl",
                                           AArch64uabd>;
 // Match UABDL in log2-shuffle patterns.
 def : Pat<(abs (v8i16 (sub (zext (v8i8 V64:$opA)),

llvm/test/CodeGen/AArch64/bf16-imm.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=aarch64 -mattr=+fullfp16 | FileCheck %s --check-prefixes=CHECK,CHECK-FP16
+; RUN: llc < %s -mtriple=aarch64 -mattr=-fullfp16 | FileCheck %s --check-prefixes=CHECK,CHECK-NOFP16
+
+define bfloat @Const0() {
+; CHECK-LABEL: Const0:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    movi d0, #0000000000000000
+; CHECK-NEXT:    ret
+entry:
+  ret bfloat 0xR0000
+}
+
+define bfloat @Const1() {
+; CHECK-FP16-LABEL: Const1:
+; CHECK-FP16:       // %bb.0: // %entry
+; CHECK-FP16-NEXT:    fmov h0, #1.00000000
+; CHECK-FP16-NEXT:    ret
+;
+; CHECK-NOFP16-LABEL: Const1:
+; CHECK-NOFP16:       // %bb.0: // %entry
+; CHECK-NOFP16-NEXT:    adrp x8, .LCPI1_0
+; CHECK-NOFP16-NEXT:    ldr h0, [x8, :lo12:.LCPI1_0]
+; CHECK-NOFP16-NEXT:    ret
+entry:
+  ret bfloat 0xR3C00
+}
+
+define bfloat @Const2() {
+; CHECK-FP16-LABEL: Const2:
+; CHECK-FP16:       // %bb.0: // %entry
+; CHECK-FP16-NEXT:    fmov h0, #0.12500000
+; CHECK-FP16-NEXT:    ret
+;
+; CHECK-NOFP16-LABEL: Const2:
+; CHECK-NOFP16:       // %bb.0: // %entry
+; CHECK-NOFP16-NEXT:    adrp x8, .LCPI2_0
+; CHECK-NOFP16-NEXT:    ldr h0, [x8, :lo12:.LCPI2_0]
+; CHECK-NOFP16-NEXT:    ret
+entry:
+  ret bfloat 0xR3000
+}
+
+define bfloat @Const3() {
+; CHECK-FP16-LABEL: Const3:
+; CHECK-FP16:       // %bb.0: // %entry
+; CHECK-FP16-NEXT:    fmov h0, #30.00000000
+; CHECK-FP16-NEXT:    ret
+;
+; CHECK-NOFP16-LABEL: Const3:
+; CHECK-NOFP16:       // %bb.0: // %entry
+; CHECK-NOFP16-NEXT:    adrp x8, .LCPI3_0
+; CHECK-NOFP16-NEXT:    ldr h0, [x8, :lo12:.LCPI3_0]
+; CHECK-NOFP16-NEXT:    ret
+entry:
+  ret bfloat 0xR4F80
+}
+
+define bfloat @Const4() {
+; CHECK-FP16-LABEL: Const4:
+; CHECK-FP16:       // %bb.0: // %entry
+; CHECK-FP16-NEXT:    fmov h0, #31.00000000
+; CHECK-FP16-NEXT:    ret
+;
+; CHECK-NOFP16-LABEL: Const4:
+; CHECK-NOFP16:       // %bb.0: // %entry
+; CHECK-NOFP16-NEXT:    adrp x8, .LCPI4_0
+; CHECK-NOFP16-NEXT:    ldr h0, [x8, :lo12:.LCPI4_0]
+; CHECK-NOFP16-NEXT:    ret
+entry:
+  ret bfloat 0xR4FC0
+}
+
+define bfloat @Const5() {
+; CHECK-FP16-LABEL: Const5:
+; CHECK-FP16:       // %bb.0: // %entry
+; CHECK-FP16-NEXT:    mov w8, #12272 // =0x2ff0
+; CHECK-FP16-NEXT:    fmov h0, w8
+; CHECK-FP16-NEXT:    ret
+;
+; CHECK-NOFP16-LABEL: Const5:
+; CHECK-NOFP16:       // %bb.0: // %entry
+; CHECK-NOFP16-NEXT:    adrp x8, .LCPI5_0
+; CHECK-NOFP16-NEXT:    ldr h0, [x8, :lo12:.LCPI5_0]
+; CHECK-NOFP16-NEXT:    ret
+entry:
+  ret bfloat 0xR2FF0
+}
+
+define bfloat @Const6() {
+; CHECK-FP16-LABEL: Const6:
+; CHECK-FP16:       // %bb.0: // %entry
+; CHECK-FP16-NEXT:    mov w8, #20417 // =0x4fc1
+; CHECK-FP16-NEXT:    fmov h0, w8
+; CHECK-FP16-NEXT:    ret
+;
+; CHECK-NOFP16-LABEL: Const6:
+; CHECK-NOFP16:       // %bb.0: // %entry
+; CHECK-NOFP16-NEXT:    adrp x8, .LCPI6_0
+; CHECK-NOFP16-NEXT:    ldr h0, [x8, :lo12:.LCPI6_0]
+; CHECK-NOFP16-NEXT:    ret
+entry:
+  ret bfloat 0xR4FC1
+}
+
+define bfloat @Const7() {
+; CHECK-FP16-LABEL: Const7:
+; CHECK-FP16:       // %bb.0: // %entry
+; CHECK-FP16-NEXT:    mov w8, #20480 // =0x5000
+; CHECK-FP16-NEXT:    fmov h0, w8
+; CHECK-FP16-NEXT:    ret
+;
+; CHECK-NOFP16-LABEL: Const7:
+; CHECK-NOFP16:       // %bb.0: // %entry
+; CHECK-NOFP16-NEXT:    adrp x8, .LCPI7_0
+; CHECK-NOFP16-NEXT:    ldr h0, [x8, :lo12:.LCPI7_0]
+; CHECK-NOFP16-NEXT:    ret
+entry:
+  ret bfloat 0xR5000
+}
+