[RISCV] Implement support for bf16 select when zfbfmin is enabled

These test cases previously caused an error. RISCVInstrInfo::copyPhysReg also needed a tweak in order to account for copying bf16 values in FPR16 registers.

Differential Revision: https://reviews.llvm.org/D156883

Author: asb

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -376,6 +376,9 @@ RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
     setOperationAction(ISD::FP_EXTEND, MVT::f32, Custom);
     setOperationAction(ISD::FP_EXTEND, MVT::f64, Custom);
     setOperationAction(ISD::ConstantFP, MVT::bf16, Expand);
+    setOperationAction(ISD::SELECT_CC, MVT::bf16, Expand);
+    setOperationAction(ISD::SELECT, MVT::bf16, Promote);
+    setOperationAction(ISD::SETCC, MVT::bf16, Promote);
   }
 
   if (Subtarget.hasStdExtZfhOrZfhminOrZhinxOrZhinxmin()) {

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

@@ -330,9 +330,10 @@ void RISCVInstrInfo::copyPhysReg(MachineBasicBlock &MBB,
     if (STI.hasStdExtZfh()) {
       Opc = RISCV::FSGNJ_H;
     } else {
-      assert(STI.hasStdExtF() && STI.hasStdExtZfhmin() &&
+      assert(STI.hasStdExtF() &&
+             (STI.hasStdExtZfhmin() || STI.hasStdExtZfbfmin()) &&
              "Unexpected extensions");
-      // Zfhmin subset doesn't have FSGNJ_H, replaces FSGNJ_H with FSGNJ_S.
+      // Zfhmin/Zfbfmin doesn't have FSGNJ_H, replace FSGNJ_H with FSGNJ_S.
       DstReg = TRI->getMatchingSuperReg(DstReg, RISCV::sub_16,
                                         &RISCV::FPR32RegClass);
       SrcReg = TRI->getMatchingSuperReg(SrcReg, RISCV::sub_16,

llvm/test/CodeGen/RISCV/bfloat-select-fcmp.ll

@@ -0,0 +1,327 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -mtriple=riscv32 -mattr=+experimental-zfbfmin -verify-machineinstrs \
+; RUN:   -target-abi ilp32f < %s | FileCheck %s
+; RUN: llc -mtriple=riscv64 -mattr=+experimental-zfbfmin -verify-machineinstrs \
+; RUN:   -target-abi lp64f < %s | FileCheck %s
+
+define bfloat @select_fcmp_false(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_false:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fmv.s fa0, fa1
+; CHECK-NEXT:    ret
+  %1 = fcmp false bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_oeq(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_oeq:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    feq.s a0, fa5, fa4
+; CHECK-NEXT:    bnez a0, .LBB1_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB1_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp oeq bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_ogt(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_ogt:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    flt.s a0, fa4, fa5
+; CHECK-NEXT:    bnez a0, .LBB2_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB2_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp ogt bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_oge(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_oge:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fle.s a0, fa4, fa5
+; CHECK-NEXT:    bnez a0, .LBB3_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB3_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp oge bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_olt(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_olt:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    flt.s a0, fa5, fa4
+; CHECK-NEXT:    bnez a0, .LBB4_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB4_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp olt bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_ole(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_ole:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    fle.s a0, fa5, fa4
+; CHECK-NEXT:    bnez a0, .LBB5_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB5_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp ole bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_one(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_one:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    flt.s a0, fa5, fa4
+; CHECK-NEXT:    flt.s a1, fa4, fa5
+; CHECK-NEXT:    or a0, a1, a0
+; CHECK-NEXT:    bnez a0, .LBB6_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB6_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp one bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_ord(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_ord:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa1
+; CHECK-NEXT:    feq.s a0, fa5, fa5
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa0
+; CHECK-NEXT:    feq.s a1, fa4, fa4
+; CHECK-NEXT:    and a0, a1, a0
+; CHECK-NEXT:    bnez a0, .LBB7_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa4, fa5
+; CHECK-NEXT:  .LBB7_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa4
+; CHECK-NEXT:    ret
+  %1 = fcmp ord bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_ueq(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_ueq:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    flt.s a0, fa5, fa4
+; CHECK-NEXT:    flt.s a1, fa4, fa5
+; CHECK-NEXT:    or a0, a1, a0
+; CHECK-NEXT:    beqz a0, .LBB8_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB8_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp ueq bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_ugt(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_ugt:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    fle.s a0, fa5, fa4
+; CHECK-NEXT:    beqz a0, .LBB9_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB9_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp ugt bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_uge(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_uge:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    flt.s a0, fa5, fa4
+; CHECK-NEXT:    beqz a0, .LBB10_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB10_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp uge bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_ult(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_ult:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fle.s a0, fa4, fa5
+; CHECK-NEXT:    beqz a0, .LBB11_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB11_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp ult bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_ule(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_ule:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    flt.s a0, fa4, fa5
+; CHECK-NEXT:    beqz a0, .LBB12_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB12_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp ule bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_une(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_une:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa0
+; CHECK-NEXT:    feq.s a0, fa5, fa4
+; CHECK-NEXT:    beqz a0, .LBB13_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa5, fa4
+; CHECK-NEXT:  .LBB13_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa5
+; CHECK-NEXT:    ret
+  %1 = fcmp une bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_uno(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_uno:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa1
+; CHECK-NEXT:    feq.s a0, fa5, fa5
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa0
+; CHECK-NEXT:    feq.s a1, fa4, fa4
+; CHECK-NEXT:    and a0, a1, a0
+; CHECK-NEXT:    beqz a0, .LBB14_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    fmv.s fa4, fa5
+; CHECK-NEXT:  .LBB14_2:
+; CHECK-NEXT:    fcvt.bf16.s fa0, fa4
+; CHECK-NEXT:    ret
+  %1 = fcmp uno bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+define bfloat @select_fcmp_true(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_true:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    ret
+  %1 = fcmp true bfloat %a, %b
+  %2 = select i1 %1, bfloat %a, bfloat %b
+  ret bfloat %2
+}
+
+; Ensure that ISel succeeds for a select+fcmp that has an i32 result type.
+define i32 @i32_select_fcmp_oeq(bfloat %a, bfloat %b, i32 %c, i32 %d) nounwind {
+; CHECK-LABEL: i32_select_fcmp_oeq:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa0
+; CHECK-NEXT:    feq.s a2, fa4, fa5
+; CHECK-NEXT:    bnez a2, .LBB16_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    mv a0, a1
+; CHECK-NEXT:  .LBB16_2:
+; CHECK-NEXT:    ret
+  %1 = fcmp oeq bfloat %a, %b
+  %2 = select i1 %1, i32 %c, i32 %d
+  ret i32 %2
+}
+
+define i32 @select_fcmp_oeq_1_2(bfloat %a, bfloat %b) {
+; CHECK-LABEL: select_fcmp_oeq_1_2:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa0
+; CHECK-NEXT:    feq.s a0, fa4, fa5
+; CHECK-NEXT:    li a1, 2
+; CHECK-NEXT:    sub a0, a1, a0
+; CHECK-NEXT:    ret
+  %1 = fcmp fast oeq bfloat %a, %b
+  %2 = select i1 %1, i32 1, i32 2
+  ret i32 %2
+}
+
+define signext i32 @select_fcmp_uge_negone_zero(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_uge_negone_zero:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa0
+; CHECK-NEXT:    fle.s a0, fa4, fa5
+; CHECK-NEXT:    addi a0, a0, -1
+; CHECK-NEXT:    ret
+  %1 = fcmp ugt bfloat %a, %b
+  %2 = select i1 %1, i32 -1, i32 0
+  ret i32 %2
+}
+
+define signext i32 @select_fcmp_uge_1_2(bfloat %a, bfloat %b) nounwind {
+; CHECK-LABEL: select_fcmp_uge_1_2:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fcvt.s.bf16 fa5, fa1
+; CHECK-NEXT:    fcvt.s.bf16 fa4, fa0
+; CHECK-NEXT:    fle.s a0, fa4, fa5
+; CHECK-NEXT:    addi a0, a0, 1
+; CHECK-NEXT:    ret
+  %1 = fcmp ugt bfloat %a, %b
+  %2 = select i1 %1, i32 1, i32 2
+  ret i32 %2
+}