[NVPTX] Expand select_cc on bfloat16 type

[qingyunqu]

Summary

Expand select_cc on bfloat16 and bfloat16v2 type.

[qingyunqu]

@bkramer Could you please help to merge this PR? I don't have the merge permission. Thanks!

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   setOperationAction(ISD::BUILD_VECTOR, MVT::v2bf16, Custom);
   setOperationAction(ISD::EXTRACT_VECTOR_ELT, MVT::v2bf16, Custom);
   setOperationAction(ISD::INSERT_VECTOR_ELT, MVT::v2bf16, Expand);
   setOperationAction(ISD::VECTOR_SHUFFLE, MVT::v2bf16, Expand);
 
   setBF16OperationAction(ISD::SETCC, MVT::bf16, Legal, Promote);
   setBF16OperationAction(ISD::SETCC, MVT::v2bf16, Legal, Expand);
   // Operations not directly supported by NVPTX.
-  for (MVT VT : {MVT::f16, MVT::v2f16, MVT::f32, MVT::f64, MVT::i1, MVT::i8,
-                 MVT::i16, MVT::i32, MVT::i64}) {
+  for (MVT VT : {MVT::bf16, MVT::f16, MVT::v2bf16, MVT::v2f16, MVT::f32,
+                 MVT::f64, MVT::i1, MVT::i8, MVT::i16, MVT::i32, MVT::i64}) {
     setOperationAction(ISD::SELECT_CC, VT, Expand);
     setOperationAction(ISD::BR_CC, VT, Expand);
   }
 
   // Some SIGN_EXTEND_INREG can be done using cvt instruction.
   // For others we will expand to a SHL/SRA pair.
   setOperationAction(ISD::SIGN_EXTEND_INREG, MVT::i64, Legal);
   setOperationAction(ISD::SIGN_EXTEND_INREG, MVT::i32, Legal);

llvm/test/CodeGen/NVPTX/bf16-instructions.ll

 ; SM80:       cvt.rn.bf16.f32 [[R:%rs[0-9]+]], [[FR]];
 
 ; CHECK:      st.param.b16    [func_retval0+0], [[R]];
 ; CHECK-NEXT: ret;
 define bfloat @test_fadd_imm_1(bfloat %a) #0 {
   %r = fadd bfloat %a, 1.0
   ret bfloat %r
 }
+
+; CHECK-LABEL: test_select_cc_bf16_f64(
+; CHECK-DAG:      ld.param.f64    [[A:%fd[0-9]+]], [test_select_cc_bf16_f64_param_0];
+; CHECK-DAG:      ld.param.f64    [[B:%fd[0-9]+]], [test_select_cc_bf16_f64_param_1];
+; CHECK:          setp.lt.f64     [[P:%p[0-9]+]], [[A]], [[B]];
+; CHECK-DAG:      ld.param.b16    [[C:%rs[0-9]+]], [test_select_cc_bf16_f64_param_2];
+; CHECK-DAG:      ld.param.b16    [[D:%rs[0-9]+]], [test_select_cc_bf16_f64_param_3];
+; CHECK:          selp.b16        [[R:%rs[0-9]+]], [[C]], [[D]], [[P]];
+; CHECK-NEXT:     st.param.b16    [func_retval0+0], [[R]];
+; CHECK-NEXT:     ret;
+define bfloat @test_select_cc_bf16_f64(double %a, double %b, bfloat %c, bfloat %d) #0 {
+  %cc = fcmp olt double %a, %b
+  %r = select i1 %cc, bfloat %c, bfloat %d
+  ret bfloat %r
+}