Skip to content

Commit b170f17

Browse files
shiltianshiltian
authored
[AMDGPU] Add support for safe bfloat16 fdiv on targets with bf16 trans instructions (#154373)
Recent changes introduced custom lowering for bf16 fdiv on targets that support bf16 trans instructions, but only covered the unsafe version. This PR extends that support to the safe variant. For the safe version, the op is lowered by converting to float, performing the div in float, and converting the result back to bf16. This matches the behavior on targets that don't support bf16 trans instructions. Fixes SWDEV-550381.
1 parent b35b629 commit b170f17

File tree

2 files changed

+71
-5
lines changed

2 files changed

+71
-5
lines changed

‎llvm/lib/Target/AMDGPU/SIISelLowering.cpp

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11540,9 +11540,22 @@ SDValue SITargetLowering::LowerFDIV16(SDValue Op, SelectionDAG &DAG) const {
1154011540
return FastLowered;
1154111541

1154211542
SDLoc SL(Op);
11543+
EVT VT = Op.getValueType();
1154311544
SDValue LHS = Op.getOperand(0);
1154411545
SDValue RHS = Op.getOperand(1);
1154511546

11547+
SDValue LHSExt = DAG.getNode(ISD::FP_EXTEND, SL, MVT::f32, LHS);
11548+
SDValue RHSExt = DAG.getNode(ISD::FP_EXTEND, SL, MVT::f32, RHS);
11549+
11550+
if (VT == MVT::bf16) {
11551+
SDValue ExtDiv =
11552+
DAG.getNode(ISD::FDIV, SL, MVT::f32, LHSExt, RHSExt, Op->getFlags());
11553+
return DAG.getNode(ISD::FP_ROUND, SL, MVT::bf16, ExtDiv,
11554+
DAG.getTargetConstant(0, SL, MVT::i32));
11555+
}
11556+
11557+
assert(VT == MVT::f16);
11558+
1154611559
// a32.u = opx(V_CVT_F32_F16, a.u); // CVT to F32
1154711560
// b32.u = opx(V_CVT_F32_F16, b.u); // CVT to F32
1154811561
// r32.u = opx(V_RCP_F32, b32.u); // rcp = 1 / d
@@ -11559,9 +11572,6 @@ SDValue SITargetLowering::LowerFDIV16(SDValue Op, SelectionDAG &DAG) const {
1155911572
// We will use ISD::FMA on targets that don't support ISD::FMAD.
1156011573
unsigned FMADOpCode =
1156111574
isOperationLegal(ISD::FMAD, MVT::f32) ? ISD::FMAD : ISD::FMA;
11562-
11563-
SDValue LHSExt = DAG.getNode(ISD::FP_EXTEND, SL, MVT::f32, LHS);
11564-
SDValue RHSExt = DAG.getNode(ISD::FP_EXTEND, SL, MVT::f32, RHS);
1156511575
SDValue NegRHSExt = DAG.getNode(ISD::FNEG, SL, MVT::f32, RHSExt);
1156611576
SDValue Rcp =
1156711577
DAG.getNode(AMDGPUISD::RCP, SL, MVT::f32, RHSExt, Op->getFlags());

‎llvm/test/CodeGen/AMDGPU/fdiv.bf16.ll

Lines changed: 58 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,68 @@
22
; RUN: llc -mtriple=amdgcn -mcpu=gfx1250 -mattr=+real-true16 -denormal-fp-math-f32=preserve-sign < %s | FileCheck -check-prefixes=GFX1250-TRUE16 %s
33
; RUN: llc -mtriple=amdgcn -mcpu=gfx1250 -mattr=-real-true16 -denormal-fp-math-f32=preserve-sign < %s | FileCheck -check-prefixes=GFX1250-FAKE16 %s
44

5-
/* TODO: Support safe bf16 fdiv lowering.
65
define bfloat @v_fdiv_bf16(bfloat %x, bfloat %y) {
6+
; GFX1250-TRUE16-LABEL: v_fdiv_bf16:
7+
; GFX1250-TRUE16: ; %bb.0:
8+
; GFX1250-TRUE16-NEXT: s_wait_loadcnt_dscnt 0x0
9+
; GFX1250-TRUE16-NEXT: s_wait_kmcnt 0x0
10+
; GFX1250-TRUE16-NEXT: v_mov_b16_e32 v2.l, 0
11+
; GFX1250-TRUE16-NEXT: v_mov_b16_e32 v2.h, v1.l
12+
; GFX1250-TRUE16-NEXT: v_mov_b16_e32 v1.h, v0.l
13+
; GFX1250-TRUE16-NEXT: s_delay_alu instid0(VALU_DEP_3) | instskip(NEXT) | instid1(VALU_DEP_1)
14+
; GFX1250-TRUE16-NEXT: v_mov_b16_e32 v1.l, v2.l
15+
; GFX1250-TRUE16-NEXT: v_div_scale_f32 v0, null, v2, v2, v1
16+
; GFX1250-TRUE16-NEXT: v_div_scale_f32 v4, vcc_lo, v1, v2, v1
17+
; GFX1250-TRUE16-NEXT: s_delay_alu instid0(VALU_DEP_2) | instskip(SKIP_2) | instid1(TRANS32_DEP_1)
18+
; GFX1250-TRUE16-NEXT: v_rcp_f32_e32 v3, v0
19+
; GFX1250-TRUE16-NEXT: s_denorm_mode 15
20+
; GFX1250-TRUE16-NEXT: v_nop
21+
; GFX1250-TRUE16-NEXT: v_fma_f32 v5, -v0, v3, 1.0
22+
; GFX1250-TRUE16-NEXT: s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_1)
23+
; GFX1250-TRUE16-NEXT: v_fmac_f32_e32 v3, v5, v3
24+
; GFX1250-TRUE16-NEXT: v_mul_f32_e32 v5, v4, v3
25+
; GFX1250-TRUE16-NEXT: s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_1)
26+
; GFX1250-TRUE16-NEXT: v_fma_f32 v6, -v0, v5, v4
27+
; GFX1250-TRUE16-NEXT: v_fmac_f32_e32 v5, v6, v3
28+
; GFX1250-TRUE16-NEXT: s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_1) | instid1(VALU_DEP_1)
29+
; GFX1250-TRUE16-NEXT: v_fma_f32 v0, -v0, v5, v4
30+
; GFX1250-TRUE16-NEXT: s_denorm_mode 12
31+
; GFX1250-TRUE16-NEXT: v_div_fmas_f32 v0, v0, v3, v5
32+
; GFX1250-TRUE16-NEXT: s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_1)
33+
; GFX1250-TRUE16-NEXT: v_div_fixup_f32 v0, v0, v2, v1
34+
; GFX1250-TRUE16-NEXT: v_cvt_pk_bf16_f32 v0, v0, s0
35+
; GFX1250-TRUE16-NEXT: s_set_pc_i64 s[30:31]
36+
;
37+
; GFX1250-FAKE16-LABEL: v_fdiv_bf16:
38+
; GFX1250-FAKE16: ; %bb.0:
39+
; GFX1250-FAKE16-NEXT: s_wait_loadcnt_dscnt 0x0
40+
; GFX1250-FAKE16-NEXT: s_wait_kmcnt 0x0
41+
; GFX1250-FAKE16-NEXT: v_dual_lshlrev_b32 v1, 16, v1 :: v_dual_lshlrev_b32 v0, 16, v0
42+
; GFX1250-FAKE16-NEXT: s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_1) | instid1(VALU_DEP_2)
43+
; GFX1250-FAKE16-NEXT: v_div_scale_f32 v2, null, v1, v1, v0
44+
; GFX1250-FAKE16-NEXT: v_div_scale_f32 v4, vcc_lo, v0, v1, v0
45+
; GFX1250-FAKE16-NEXT: v_rcp_f32_e32 v3, v2
46+
; GFX1250-FAKE16-NEXT: s_denorm_mode 15
47+
; GFX1250-FAKE16-NEXT: v_nop
48+
; GFX1250-FAKE16-NEXT: s_delay_alu instid0(TRANS32_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_1)
49+
; GFX1250-FAKE16-NEXT: v_fma_f32 v5, -v2, v3, 1.0
50+
; GFX1250-FAKE16-NEXT: v_fmac_f32_e32 v3, v5, v3
51+
; GFX1250-FAKE16-NEXT: s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_1)
52+
; GFX1250-FAKE16-NEXT: v_mul_f32_e32 v5, v4, v3
53+
; GFX1250-FAKE16-NEXT: v_fma_f32 v6, -v2, v5, v4
54+
; GFX1250-FAKE16-NEXT: s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_1)
55+
; GFX1250-FAKE16-NEXT: v_fmac_f32_e32 v5, v6, v3
56+
; GFX1250-FAKE16-NEXT: v_fma_f32 v2, -v2, v5, v4
57+
; GFX1250-FAKE16-NEXT: s_denorm_mode 12
58+
; GFX1250-FAKE16-NEXT: s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_1)
59+
; GFX1250-FAKE16-NEXT: v_div_fmas_f32 v2, v2, v3, v5
60+
; GFX1250-FAKE16-NEXT: v_div_fixup_f32 v0, v2, v1, v0
61+
; GFX1250-FAKE16-NEXT: s_delay_alu instid0(VALU_DEP_1)
62+
; GFX1250-FAKE16-NEXT: v_cvt_pk_bf16_f32 v0, v0, s0
63+
; GFX1250-FAKE16-NEXT: s_set_pc_i64 s[30:31]
764
%fdiv = fdiv bfloat %x, %y
865
ret bfloat %fdiv
966
}
10-
*/
1167

1268
define bfloat @v_rcp_bf16(bfloat %x) {
1369
; GFX1250-TRUE16-LABEL: v_rcp_bf16:

0 commit comments

Comments
 (0)