[AMDGPU] Enable CodeGen for v_pk_fma_bf16 (#152578)

Author: rampitec

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -6106,6 +6106,7 @@ bool SITargetLowering::isFMAFasterThanFMulAndFAdd(const MachineFunction &MF,
   case MVT::f64:
     return true;
   case MVT::f16:
+  case MVT::bf16:
     return Subtarget->has16BitInsts() && !denormalModeIsFlushAllF64F16(MF);
   default:
     break;

llvm/test/CodeGen/AMDGPU/bf16-math.ll

@@ -370,6 +370,9 @@ define amdgpu_ps bfloat @test_clamp_bf16_folding(bfloat %src) {
 ; GCN:       ; %bb.0:
 ; GCN-NEXT:    v_exp_bf16_e64 v0, v0 clamp
 ; GCN-NEXT:    ; return to shader part epilog
+
+
+
   %exp = call bfloat @llvm.exp2.bf16(bfloat %src)
   %max = call bfloat @llvm.maxnum.bf16(bfloat %exp, bfloat 0.0)
   %clamp = call bfloat @llvm.minnum.bf16(bfloat %max, bfloat 1.0)
@@ -381,6 +384,9 @@ define amdgpu_ps float @test_clamp_v2bf16_folding(<2 x bfloat> %src0, <2 x bfloa
 ; GCN:       ; %bb.0:
 ; GCN-NEXT:    v_pk_mul_bf16 v0, v0, v1 clamp
 ; GCN-NEXT:    ; return to shader part epilog
+
+
+
   %mul = fmul <2 x bfloat> %src0, %src1
   %max = call <2 x bfloat> @llvm.maxnum.v2bf16(<2 x bfloat> %mul, <2 x bfloat> <bfloat 0.0, bfloat 0.0>)
   %clamp = call <2 x bfloat> @llvm.minnum.v2bf16(<2 x bfloat> %max, <2 x bfloat> <bfloat 1.0, bfloat 1.0>)
@@ -391,11 +397,12 @@ define amdgpu_ps float @test_clamp_v2bf16_folding(<2 x bfloat> %src0, <2 x bfloa
 define amdgpu_ps void @v_test_mul_add_v2bf16_vvv(ptr addrspace(1) %out, <2 x bfloat> %a, <2 x bfloat> %b, <2 x bfloat> %c) {
 ; GCN-LABEL: v_test_mul_add_v2bf16_vvv:
 ; GCN:       ; %bb.0:
-; GCN-NEXT:    v_pk_mul_bf16 v2, v2, v3
-; GCN-NEXT:    s_delay_alu instid0(VALU_DEP_1)
-; GCN-NEXT:    v_pk_add_bf16 v2, v2, v4
+; GCN-NEXT:    v_pk_fma_bf16 v2, v2, v3, v4
 ; GCN-NEXT:    global_store_b32 v[0:1], v2, off
 ; GCN-NEXT:    s_endpgm
+
+
+
   %mul = fmul contract <2 x bfloat> %a, %b
   %add = fadd contract <2 x bfloat> %mul, %c
   store <2 x bfloat> %add, ptr addrspace(1) %out
@@ -405,11 +412,12 @@ define amdgpu_ps void @v_test_mul_add_v2bf16_vvv(ptr addrspace(1) %out, <2 x bfl
 define amdgpu_ps void @v_test_mul_add_v2bf16_vss(ptr addrspace(1) %out, <2 x bfloat> %a, <2 x bfloat> inreg %b, <2 x bfloat> inreg %c) {
 ; GCN-LABEL: v_test_mul_add_v2bf16_vss:
 ; GCN:       ; %bb.0:
-; GCN-NEXT:    v_pk_mul_bf16 v2, v2, s0
-; GCN-NEXT:    s_delay_alu instid0(VALU_DEP_1)
-; GCN-NEXT:    v_pk_add_bf16 v2, v2, s1
+; GCN-NEXT:    v_pk_fma_bf16 v2, v2, s0, s1
 ; GCN-NEXT:    global_store_b32 v[0:1], v2, off
 ; GCN-NEXT:    s_endpgm
+
+
+
   %mul = fmul contract <2 x bfloat> %a, %b
   %add = fadd contract <2 x bfloat> %mul, %c
   store <2 x bfloat> %add, ptr addrspace(1) %out
@@ -419,11 +427,14 @@ define amdgpu_ps void @v_test_mul_add_v2bf16_vss(ptr addrspace(1) %out, <2 x bfl
 define amdgpu_ps void @v_test_mul_add_v2bf16_sss(ptr addrspace(1) %out, <2 x bfloat> inreg %a, <2 x bfloat> inreg %b, <2 x bfloat> inreg %c) {
 ; GCN-LABEL: v_test_mul_add_v2bf16_sss:
 ; GCN:       ; %bb.0:
-; GCN-NEXT:    v_pk_mul_bf16 v2, s0, s1
+; GCN-NEXT:    v_mov_b32_e32 v2, s2
 ; GCN-NEXT:    s_delay_alu instid0(VALU_DEP_1)
-; GCN-NEXT:    v_pk_add_bf16 v2, v2, s2
+; GCN-NEXT:    v_pk_fma_bf16 v2, s0, s1, v2
 ; GCN-NEXT:    global_store_b32 v[0:1], v2, off
 ; GCN-NEXT:    s_endpgm
+
+
+
   %mul = fmul contract <2 x bfloat> %a, %b
   %add = fadd contract <2 x bfloat> %mul, %c
   store <2 x bfloat> %add, ptr addrspace(1) %out
@@ -433,11 +444,12 @@ define amdgpu_ps void @v_test_mul_add_v2bf16_sss(ptr addrspace(1) %out, <2 x bfl
 define amdgpu_ps void @v_test_mul_add_v2bf16_vsc(ptr addrspace(1) %out, <2 x bfloat> %a, <2 x bfloat> inreg %b) {
 ; GCN-LABEL: v_test_mul_add_v2bf16_vsc:
 ; GCN:       ; %bb.0:
-; GCN-NEXT:    v_pk_mul_bf16 v2, v2, s0
-; GCN-NEXT:    s_delay_alu instid0(VALU_DEP_1)
-; GCN-NEXT:    v_pk_add_bf16 v2, v2, 0.5 op_sel_hi:[1,0]
+; GCN-NEXT:    v_pk_fma_bf16 v2, v2, s0, 0.5 op_sel_hi:[1,1,0]
 ; GCN-NEXT:    global_store_b32 v[0:1], v2, off
 ; GCN-NEXT:    s_endpgm
+
+
+
   %mul = fmul contract <2 x bfloat> %a, %b
   %add = fadd contract <2 x bfloat> %mul, <bfloat 0.5, bfloat 0.5>
   store <2 x bfloat> %add, ptr addrspace(1) %out
@@ -447,11 +459,14 @@ define amdgpu_ps void @v_test_mul_add_v2bf16_vsc(ptr addrspace(1) %out, <2 x bfl
 define amdgpu_ps void @v_test_mul_add_v2bf16_vll(ptr addrspace(1) %out, <2 x bfloat> %a) {
 ; GCN-LABEL: v_test_mul_add_v2bf16_vll:
 ; GCN:       ; %bb.0:
-; GCN-NEXT:    v_pk_mul_bf16 v2, 0x42c83f80, v2
-; GCN-NEXT:    s_delay_alu instid0(VALU_DEP_1)
-; GCN-NEXT:    v_pk_add_bf16 v2, 0x43484000, v2
+; GCN-NEXT:    s_mov_b32 s0, 0x43484000
+; GCN-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; GCN-NEXT:    v_pk_fma_bf16 v2, 0x42c83f80, v2, s0
 ; GCN-NEXT:    global_store_b32 v[0:1], v2, off
 ; GCN-NEXT:    s_endpgm
+
+
+
   %mul = fmul contract <2 x bfloat> %a, <bfloat 1.0, bfloat 100.0>
   %add = fadd contract <2 x bfloat> %mul, <bfloat 2.0, bfloat 200.0>
   store <2 x bfloat> %add, ptr addrspace(1) %out