[AMDGPU] Improve AMDGPU sqrt and inverse sqrt handling for bf16 by aeft · Pull Request #180291 · llvm/llvm-project

aeft · 2026-02-06T22:03:21Z

Following the approach in lowerFSQRTF16, this adds a custom lowering for bf16 FSQRT that extends to f32, calls amdgcn_sqrt (v_sqrt_f32), and rounds back to bf16. Denormal inputs are preserved via scaling, since bf16 denormals remain denormals after extending to f32.

Fixes: #172045

aeft · 2026-02-06T22:03:41Z

Unlike f16, this introduces a behavior change for denormals: bf16 and f32 share the same exponent format (8-bit, bias 127), so bf16 denormals remain denormals after extending to f32. v_sqrt_f32 always (may?) flushes f32 denormal inputs, so the old code had a scaling workaround (multiply by 2^32) to handle this in denorm-preserve mode. This PR removes that workaround. (For f16, this issue doesn't exist: f16 denormals become normal f32 values after extension, so lowerFSQRTF16 is lossless.)

Is dropping denormal preservation for bf16 sqrt acceptable?

llvmbot · 2026-02-06T22:03:56Z

@llvm/pr-subscribers-llvm-analysis

@llvm/pr-subscribers-backend-amdgpu

Author: Alex Wang (aeft)

Changes

Following the approach in lowerFSQRTF16, this adds a custom lowering for bf16 FSQRT that extends to f32, calls amdgcn_sqrt (v_sqrt_f32), and rounds back to bf16, reducing the sqrt expansion from ~20 instructions to just v_sqrt_f32 plus bf16 conversion/rounding.

Fixes: #172045

Patch is 33.34 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/180291.diff

3 Files Affected:

(modified) llvm/lib/Target/AMDGPU/SIISelLowering.cpp (+27-9)
(modified) llvm/lib/Target/AMDGPU/SIISelLowering.h (+1)
(modified) llvm/test/CodeGen/AMDGPU/bf16.ll (+34-374)

diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index cdf6fb97d0b3b..9cb39937a2881 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -221,15 +221,14 @@ SITargetLowering::SITargetLowering(const TargetMachine &TM,
 
   if (isTypeLegal(MVT::bf16)) {
     for (unsigned Opc :
-         {ISD::FADD,     ISD::FSUB,       ISD::FMUL,    ISD::FDIV,
-          ISD::FREM,     ISD::FMA,        ISD::FMINNUM, ISD::FMAXNUM,
-          ISD::FMINIMUM, ISD::FMAXIMUM,   ISD::FSQRT,   ISD::FCBRT,
-          ISD::FSIN,     ISD::FCOS,       ISD::FPOW,    ISD::FPOWI,
-          ISD::FLDEXP,   ISD::FFREXP,     ISD::FLOG,    ISD::FLOG2,
-          ISD::FLOG10,   ISD::FEXP,       ISD::FEXP2,   ISD::FEXP10,
-          ISD::FCEIL,    ISD::FTRUNC,     ISD::FRINT,   ISD::FNEARBYINT,
-          ISD::FROUND,   ISD::FROUNDEVEN, ISD::FFLOOR,  ISD::FCANONICALIZE,
-          ISD::SETCC}) {
+         {ISD::FADD,       ISD::FSUB,     ISD::FMUL,          ISD::FDIV,
+          ISD::FREM,       ISD::FMA,      ISD::FMINNUM,       ISD::FMAXNUM,
+          ISD::FMINIMUM,   ISD::FMAXIMUM, ISD::FCBRT,         ISD::FSIN,
+          ISD::FCOS,       ISD::FPOW,     ISD::FPOWI,         ISD::FLDEXP,
+          ISD::FFREXP,     ISD::FLOG,     ISD::FLOG2,         ISD::FLOG10,
+          ISD::FEXP,       ISD::FEXP2,    ISD::FEXP10,        ISD::FCEIL,
+          ISD::FTRUNC,     ISD::FRINT,    ISD::FNEARBYINT,    ISD::FROUND,
+          ISD::FROUNDEVEN, ISD::FFLOOR,   ISD::FCANONICALIZE, ISD::SETCC}) {
       setOperationAction(Opc, MVT::bf16, Promote);
     }
 
@@ -246,6 +245,8 @@ SITargetLowering::SITargetLowering(const TargetMachine &TM,
     // sources.
     setOperationAction(ISD::FP_TO_SINT, MVT::i32, Custom);
     setOperationAction(ISD::FP_TO_UINT, MVT::i32, Custom);
+
+    setOperationAction(ISD::FSQRT, MVT::bf16, Custom);
   }
 
   setTruncStoreAction(MVT::v2i32, MVT::v2i16, Expand);
@@ -7046,6 +7047,8 @@ SDValue SITargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
       return lowerFSQRTF32(Op, DAG);
     if (VT == MVT::f64)
       return lowerFSQRTF64(Op, DAG);
+    if (VT == MVT::bf16)
+      return lowerFSQRTBF16(Op, DAG);
     return SDValue();
   }
   case ISD::FSIN:
@@ -13184,6 +13187,21 @@ SDValue SITargetLowering::lowerFSQRTF64(SDValue Op, SelectionDAG &DAG) const {
                      Flags);
 }
 
+SDValue SITargetLowering::lowerFSQRTBF16(SDValue Op, SelectionDAG &DAG) const {
+  SDLoc SL(Op);
+  assert(!Subtarget->hasBF16TransInsts());
+  SDNodeFlags Flags = Op->getFlags();
+  SDValue Ext =
+      DAG.getNode(ISD::FP_EXTEND, SL, MVT::f32, Op.getOperand(0), Flags);
+
+  SDValue SqrtID = DAG.getTargetConstant(Intrinsic::amdgcn_sqrt, SL, MVT::i32);
+  SDValue Sqrt =
+      DAG.getNode(ISD::INTRINSIC_WO_CHAIN, SL, MVT::f32, SqrtID, Ext, Flags);
+
+  return DAG.getNode(ISD::FP_ROUND, SL, MVT::bf16, Sqrt,
+                     DAG.getTargetConstant(0, SL, MVT::i32), Flags);
+}
+
 SDValue SITargetLowering::LowerTrig(SDValue Op, SelectionDAG &DAG) const {
   SDLoc DL(Op);
   EVT VT = Op.getValueType();
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.h b/llvm/lib/Target/AMDGPU/SIISelLowering.h
index 5a4f249c166aa..1d9f73827c81b 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.h
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.h
@@ -129,6 +129,7 @@ class SITargetLowering final : public AMDGPUTargetLowering {
   SDValue lowerFSQRTF16(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerFSQRTF32(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerFSQRTF64(SDValue Op, SelectionDAG &DAG) const;
+  SDValue lowerFSQRTBF16(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerATOMIC_CMP_SWAP(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerBRCOND(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerRETURNADDR(SDValue Op, SelectionDAG &DAG) const;
diff --git a/llvm/test/CodeGen/AMDGPU/bf16.ll b/llvm/test/CodeGen/AMDGPU/bf16.ll
index c404225da6f05..34423da59eee9 100644
--- a/llvm/test/CodeGen/AMDGPU/bf16.ll
+++ b/llvm/test/CodeGen/AMDGPU/bf16.ll
@@ -28388,24 +28388,7 @@ define bfloat @v_sqrt_bf16(bfloat %a) {
 ; GFX8:       ; %bb.0:
 ; GFX8-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX8-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
-; GFX8-NEXT:    s_mov_b32 s4, 0xf800000
-; GFX8-NEXT:    v_mul_f32_e32 v1, 0x4f800000, v0
-; GFX8-NEXT:    v_cmp_gt_f32_e32 vcc, s4, v0
-; GFX8-NEXT:    v_cndmask_b32_e32 v0, v0, v1, vcc
-; GFX8-NEXT:    v_sqrt_f32_e32 v1, v0
-; GFX8-NEXT:    v_add_u32_e64 v2, s[4:5], -1, v1
-; GFX8-NEXT:    v_fma_f32 v3, -v2, v1, v0
-; GFX8-NEXT:    v_cmp_ge_f32_e64 s[4:5], 0, v3
-; GFX8-NEXT:    v_cndmask_b32_e64 v2, v1, v2, s[4:5]
-; GFX8-NEXT:    v_add_u32_e64 v3, s[4:5], 1, v1
-; GFX8-NEXT:    v_fma_f32 v1, -v3, v1, v0
-; GFX8-NEXT:    v_cmp_lt_f32_e64 s[4:5], 0, v1
-; GFX8-NEXT:    v_cndmask_b32_e64 v1, v2, v3, s[4:5]
-; GFX8-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX8-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc
-; GFX8-NEXT:    v_mov_b32_e32 v2, 0x260
-; GFX8-NEXT:    v_cmp_class_f32_e32 vcc, v0, v2
-; GFX8-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc
+; GFX8-NEXT:    v_sqrt_f32_e32 v0, v0
 ; GFX8-NEXT:    v_bfe_u32 v1, v0, 16, 1
 ; GFX8-NEXT:    v_add_u32_e32 v1, vcc, v1, v0
 ; GFX8-NEXT:    v_add_u32_e32 v1, vcc, 0x7fff, v1
@@ -28419,28 +28402,11 @@ define bfloat @v_sqrt_bf16(bfloat %a) {
 ; GFX900:       ; %bb.0:
 ; GFX900-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX900-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
-; GFX900-NEXT:    s_mov_b32 s4, 0xf800000
-; GFX900-NEXT:    v_mul_f32_e32 v1, 0x4f800000, v0
-; GFX900-NEXT:    v_cmp_gt_f32_e32 vcc, s4, v0
-; GFX900-NEXT:    v_cndmask_b32_e32 v0, v0, v1, vcc
-; GFX900-NEXT:    v_sqrt_f32_e32 v1, v0
-; GFX900-NEXT:    v_add_u32_e32 v2, -1, v1
-; GFX900-NEXT:    v_fma_f32 v3, -v2, v1, v0
-; GFX900-NEXT:    v_cmp_ge_f32_e64 s[4:5], 0, v3
-; GFX900-NEXT:    v_add_u32_e32 v3, 1, v1
-; GFX900-NEXT:    v_cndmask_b32_e64 v2, v1, v2, s[4:5]
-; GFX900-NEXT:    v_fma_f32 v1, -v3, v1, v0
-; GFX900-NEXT:    v_cmp_lt_f32_e64 s[4:5], 0, v1
-; GFX900-NEXT:    v_cndmask_b32_e64 v1, v2, v3, s[4:5]
-; GFX900-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX900-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc
-; GFX900-NEXT:    v_mov_b32_e32 v2, 0x260
-; GFX900-NEXT:    v_cmp_class_f32_e32 vcc, v0, v2
-; GFX900-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc
-; GFX900-NEXT:    v_bfe_u32 v1, v0, 16, 1
+; GFX900-NEXT:    v_sqrt_f32_e32 v0, v0
 ; GFX900-NEXT:    s_movk_i32 s4, 0x7fff
-; GFX900-NEXT:    v_add3_u32 v1, v1, v0, s4
+; GFX900-NEXT:    v_bfe_u32 v1, v0, 16, 1
 ; GFX900-NEXT:    v_or_b32_e32 v2, 0x400000, v0
+; GFX900-NEXT:    v_add3_u32 v1, v1, v0, s4
 ; GFX900-NEXT:    v_cmp_u_f32_e32 vcc, v0, v0
 ; GFX900-NEXT:    v_cndmask_b32_e32 v0, v1, v2, vcc
 ; GFX900-NEXT:    v_lshrrev_b32_e32 v0, 16, v0
@@ -28450,29 +28416,8 @@ define bfloat @v_sqrt_bf16(bfloat %a) {
 ; GFX950:       ; %bb.0:
 ; GFX950-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX950-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
-; GFX950-NEXT:    s_mov_b32 s0, 0xf800000
-; GFX950-NEXT:    v_mul_f32_e32 v1, 0x4f800000, v0
-; GFX950-NEXT:    v_cmp_gt_f32_e32 vcc, s0, v0
-; GFX950-NEXT:    s_nop 1
-; GFX950-NEXT:    v_cndmask_b32_e32 v0, v0, v1, vcc
-; GFX950-NEXT:    v_sqrt_f32_e32 v1, v0
-; GFX950-NEXT:    s_nop 0
-; GFX950-NEXT:    v_add_u32_e32 v2, -1, v1
-; GFX950-NEXT:    v_fma_f32 v3, -v2, v1, v0
-; GFX950-NEXT:    v_cmp_ge_f32_e64 s[0:1], 0, v3
-; GFX950-NEXT:    v_add_u32_e32 v3, 1, v1
+; GFX950-NEXT:    v_sqrt_f32_e32 v0, v0
 ; GFX950-NEXT:    s_nop 0
-; GFX950-NEXT:    v_cndmask_b32_e64 v2, v1, v2, s[0:1]
-; GFX950-NEXT:    v_fma_f32 v1, -v3, v1, v0
-; GFX950-NEXT:    v_cmp_lt_f32_e64 s[0:1], 0, v1
-; GFX950-NEXT:    s_nop 1
-; GFX950-NEXT:    v_cndmask_b32_e64 v1, v2, v3, s[0:1]
-; GFX950-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX950-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc
-; GFX950-NEXT:    v_mov_b32_e32 v2, 0x260
-; GFX950-NEXT:    v_cmp_class_f32_e32 vcc, v0, v2
-; GFX950-NEXT:    s_nop 1
-; GFX950-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc
 ; GFX950-NEXT:    v_cvt_pk_bf16_f32 v0, v0, s0
 ; GFX950-NEXT:    s_setpc_b64 s[30:31]
 ;
@@ -28480,22 +28425,7 @@ define bfloat @v_sqrt_bf16(bfloat %a) {
 ; GFX10:       ; %bb.0:
 ; GFX10-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX10-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
-; GFX10-NEXT:    v_mul_f32_e32 v1, 0x4f800000, v0
-; GFX10-NEXT:    v_cmp_gt_f32_e32 vcc_lo, 0xf800000, v0
-; GFX10-NEXT:    v_cndmask_b32_e32 v0, v0, v1, vcc_lo
-; GFX10-NEXT:    v_sqrt_f32_e32 v1, v0
-; GFX10-NEXT:    v_add_nc_u32_e32 v2, -1, v1
-; GFX10-NEXT:    v_add_nc_u32_e32 v3, 1, v1
-; GFX10-NEXT:    v_fma_f32 v4, -v2, v1, v0
-; GFX10-NEXT:    v_fma_f32 v5, -v3, v1, v0
-; GFX10-NEXT:    v_cmp_ge_f32_e64 s4, 0, v4
-; GFX10-NEXT:    v_cndmask_b32_e64 v1, v1, v2, s4
-; GFX10-NEXT:    v_cmp_lt_f32_e64 s4, 0, v5
-; GFX10-NEXT:    v_cndmask_b32_e64 v1, v1, v3, s4
-; GFX10-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX10-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc_lo
-; GFX10-NEXT:    v_cmp_class_f32_e64 vcc_lo, v0, 0x260
-; GFX10-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc_lo
+; GFX10-NEXT:    v_sqrt_f32_e32 v0, v0
 ; GFX10-NEXT:    v_bfe_u32 v1, v0, 16, 1
 ; GFX10-NEXT:    v_or_b32_e32 v2, 0x400000, v0
 ; GFX10-NEXT:    v_cmp_u_f32_e32 vcc_lo, v0, v0
@@ -28509,30 +28439,9 @@ define bfloat @v_sqrt_bf16(bfloat %a) {
 ; GFX11TRUE16-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX11TRUE16-NEXT:    v_mov_b16_e32 v1.l, 0
 ; GFX11TRUE16-NEXT:    v_mov_b16_e32 v1.h, v0.l
-; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_1) | instid1(VALU_DEP_1)
-; GFX11TRUE16-NEXT:    v_cmp_gt_f32_e32 vcc_lo, 0xf800000, v1
-; GFX11TRUE16-NEXT:    v_mul_f32_e32 v0, 0x4f800000, v1
-; GFX11TRUE16-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc_lo
-; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_3) | instid1(VALU_DEP_2)
-; GFX11TRUE16-NEXT:    v_sqrt_f32_e32 v1, v0
+; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_4) | instid1(VALU_DEP_3)
+; GFX11TRUE16-NEXT:    v_sqrt_f32_e32 v0, v1
 ; GFX11TRUE16-NEXT:    s_waitcnt_depctr depctr_va_vdst(0)
-; GFX11TRUE16-NEXT:    v_add_nc_u32_e32 v2, -1, v1
-; GFX11TRUE16-NEXT:    v_add_nc_u32_e32 v3, 1, v1
-; GFX11TRUE16-NEXT:    v_fma_f32 v4, -v2, v1, v0
-; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_2) | instskip(NEXT) | instid1(VALU_DEP_2)
-; GFX11TRUE16-NEXT:    v_fma_f32 v5, -v3, v1, v0
-; GFX11TRUE16-NEXT:    v_cmp_ge_f32_e64 s0, 0, v4
-; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_3)
-; GFX11TRUE16-NEXT:    v_cndmask_b32_e64 v1, v1, v2, s0
-; GFX11TRUE16-NEXT:    v_cmp_lt_f32_e64 s0, 0, v5
-; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_1)
-; GFX11TRUE16-NEXT:    v_cndmask_b32_e64 v1, v1, v3, s0
-; GFX11TRUE16-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_1) | instid1(VALU_DEP_2)
-; GFX11TRUE16-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc_lo
-; GFX11TRUE16-NEXT:    v_cmp_class_f32_e64 vcc_lo, v0, 0x260
-; GFX11TRUE16-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc_lo
-; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_2) | instid1(VALU_DEP_3)
 ; GFX11TRUE16-NEXT:    v_bfe_u32 v1, v0, 16, 1
 ; GFX11TRUE16-NEXT:    v_or_b32_e32 v2, 0x400000, v0
 ; GFX11TRUE16-NEXT:    v_cmp_u_f32_e32 vcc_lo, v0, v0
@@ -28546,30 +28455,9 @@ define bfloat @v_sqrt_bf16(bfloat %a) {
 ; GFX11FAKE16:       ; %bb.0:
 ; GFX11FAKE16-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX11FAKE16-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
-; GFX11FAKE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_1) | instid1(VALU_DEP_2)
-; GFX11FAKE16-NEXT:    v_mul_f32_e32 v1, 0x4f800000, v0
-; GFX11FAKE16-NEXT:    v_cmp_gt_f32_e32 vcc_lo, 0xf800000, v0
-; GFX11FAKE16-NEXT:    v_cndmask_b32_e32 v0, v0, v1, vcc_lo
-; GFX11FAKE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_3) | instid1(VALU_DEP_2)
-; GFX11FAKE16-NEXT:    v_sqrt_f32_e32 v1, v0
+; GFX11FAKE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_4) | instid1(VALU_DEP_3)
+; GFX11FAKE16-NEXT:    v_sqrt_f32_e32 v0, v0
 ; GFX11FAKE16-NEXT:    s_waitcnt_depctr depctr_va_vdst(0)
-; GFX11FAKE16-NEXT:    v_add_nc_u32_e32 v2, -1, v1
-; GFX11FAKE16-NEXT:    v_add_nc_u32_e32 v3, 1, v1
-; GFX11FAKE16-NEXT:    v_fma_f32 v4, -v2, v1, v0
-; GFX11FAKE16-NEXT:    s_delay_alu instid0(VALU_DEP_2) | instskip(NEXT) | instid1(VALU_DEP_2)
-; GFX11FAKE16-NEXT:    v_fma_f32 v5, -v3, v1, v0
-; GFX11FAKE16-NEXT:    v_cmp_ge_f32_e64 s0, 0, v4
-; GFX11FAKE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_3)
-; GFX11FAKE16-NEXT:    v_cndmask_b32_e64 v1, v1, v2, s0
-; GFX11FAKE16-NEXT:    v_cmp_lt_f32_e64 s0, 0, v5
-; GFX11FAKE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(NEXT) | instid1(VALU_DEP_1)
-; GFX11FAKE16-NEXT:    v_cndmask_b32_e64 v1, v1, v3, s0
-; GFX11FAKE16-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX11FAKE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_1) | instid1(VALU_DEP_2)
-; GFX11FAKE16-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc_lo
-; GFX11FAKE16-NEXT:    v_cmp_class_f32_e64 vcc_lo, v0, 0x260
-; GFX11FAKE16-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc_lo
-; GFX11FAKE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_2) | instid1(VALU_DEP_3)
 ; GFX11FAKE16-NEXT:    v_bfe_u32 v1, v0, 16, 1
 ; GFX11FAKE16-NEXT:    v_or_b32_e32 v2, 0x400000, v0
 ; GFX11FAKE16-NEXT:    v_cmp_u_f32_e32 vcc_lo, v0, v0
@@ -28675,28 +28563,10 @@ define bfloat @v_rsq_bf16(bfloat %x) {
 ; GFX8:       ; %bb.0:
 ; GFX8-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX8-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
-; GFX8-NEXT:    s_mov_b32 s4, 0xf800000
-; GFX8-NEXT:    v_mul_f32_e32 v1, 0x4f800000, v0
-; GFX8-NEXT:    v_cmp_gt_f32_e32 vcc, s4, v0
-; GFX8-NEXT:    v_cndmask_b32_e32 v0, v0, v1, vcc
-; GFX8-NEXT:    v_sqrt_f32_e32 v1, v0
-; GFX8-NEXT:    v_add_u32_e64 v2, s[4:5], -1, v1
-; GFX8-NEXT:    v_fma_f32 v3, -v2, v1, v0
-; GFX8-NEXT:    v_cmp_ge_f32_e64 s[4:5], 0, v3
-; GFX8-NEXT:    v_cndmask_b32_e64 v2, v1, v2, s[4:5]
-; GFX8-NEXT:    v_add_u32_e64 v3, s[4:5], 1, v1
-; GFX8-NEXT:    v_fma_f32 v1, -v3, v1, v0
-; GFX8-NEXT:    v_cmp_lt_f32_e64 s[4:5], 0, v1
-; GFX8-NEXT:    v_cndmask_b32_e64 v1, v2, v3, s[4:5]
-; GFX8-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX8-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc
-; GFX8-NEXT:    v_mov_b32_e32 v2, 0x260
-; GFX8-NEXT:    v_cmp_class_f32_e32 vcc, v0, v2
-; GFX8-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc
+; GFX8-NEXT:    v_sqrt_f32_e32 v0, v0
 ; GFX8-NEXT:    v_bfe_u32 v1, v0, 16, 1
 ; GFX8-NEXT:    v_add_u32_e32 v1, vcc, v1, v0
-; GFX8-NEXT:    s_movk_i32 s4, 0x7fff
-; GFX8-NEXT:    v_add_u32_e32 v1, vcc, s4, v1
+; GFX8-NEXT:    v_add_u32_e32 v1, vcc, 0x7fff, v1
 ; GFX8-NEXT:    v_or_b32_e32 v2, 0x400000, v0
 ; GFX8-NEXT:    v_cmp_u_f32_e32 vcc, v0, v0
 ; GFX8-NEXT:    v_cndmask_b32_e32 v0, v1, v2, vcc
@@ -28725,28 +28595,11 @@ define bfloat @v_rsq_bf16(bfloat %x) {
 ; GFX900:       ; %bb.0:
 ; GFX900-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX900-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
-; GFX900-NEXT:    s_mov_b32 s4, 0xf800000
-; GFX900-NEXT:    v_mul_f32_e32 v1, 0x4f800000, v0
-; GFX900-NEXT:    v_cmp_gt_f32_e32 vcc, s4, v0
-; GFX900-NEXT:    v_cndmask_b32_e32 v0, v0, v1, vcc
-; GFX900-NEXT:    v_sqrt_f32_e32 v1, v0
+; GFX900-NEXT:    v_sqrt_f32_e32 v0, v0
 ; GFX900-NEXT:    s_movk_i32 s6, 0x7fff
-; GFX900-NEXT:    v_add_u32_e32 v2, -1, v1
-; GFX900-NEXT:    v_fma_f32 v3, -v2, v1, v0
-; GFX900-NEXT:    v_cmp_ge_f32_e64 s[4:5], 0, v3
-; GFX900-NEXT:    v_add_u32_e32 v3, 1, v1
-; GFX900-NEXT:    v_cndmask_b32_e64 v2, v1, v2, s[4:5]
-; GFX900-NEXT:    v_fma_f32 v1, -v3, v1, v0
-; GFX900-NEXT:    v_cmp_lt_f32_e64 s[4:5], 0, v1
-; GFX900-NEXT:    v_cndmask_b32_e64 v1, v2, v3, s[4:5]
-; GFX900-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX900-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc
-; GFX900-NEXT:    v_mov_b32_e32 v2, 0x260
-; GFX900-NEXT:    v_cmp_class_f32_e32 vcc, v0, v2
-; GFX900-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc
 ; GFX900-NEXT:    v_bfe_u32 v1, v0, 16, 1
-; GFX900-NEXT:    v_add3_u32 v1, v1, v0, s6
 ; GFX900-NEXT:    v_or_b32_e32 v2, 0x400000, v0
+; GFX900-NEXT:    v_add3_u32 v1, v1, v0, s6
 ; GFX900-NEXT:    v_cmp_u_f32_e32 vcc, v0, v0
 ; GFX900-NEXT:    v_cndmask_b32_e32 v0, v1, v2, vcc
 ; GFX900-NEXT:    v_and_b32_e32 v0, 0xffff0000, v0
@@ -28773,37 +28626,15 @@ define bfloat @v_rsq_bf16(bfloat %x) {
 ; GFX950:       ; %bb.0:
 ; GFX950-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX950-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
-; GFX950-NEXT:    s_mov_b32 s0, 0xf800000
-; GFX950-NEXT:    v_mul_f32_e32 v1, 0x4f800000, v0
-; GFX950-NEXT:    v_cmp_gt_f32_e32 vcc, s0, v0
-; GFX950-NEXT:    s_nop 1
-; GFX950-NEXT:    v_cndmask_b32_e32 v0, v0, v1, vcc
-; GFX950-NEXT:    v_sqrt_f32_e32 v1, v0
-; GFX950-NEXT:    s_nop 0
-; GFX950-NEXT:    v_add_u32_e32 v2, -1, v1
-; GFX950-NEXT:    v_fma_f32 v3, -v2, v1, v0
-; GFX950-NEXT:    v_cmp_ge_f32_e64 s[0:1], 0, v3
-; GFX950-NEXT:    v_add_u32_e32 v3, 1, v1
+; GFX950-NEXT:    v_sqrt_f32_e32 v0, v0
 ; GFX950-NEXT:    s_nop 0
-; GFX950-NEXT:    v_cndmask_b32_e64 v2, v1, v2, s[0:1]
-; GFX950-NEXT:    v_fma_f32 v1, -v3, v1, v0
-; GFX950-NEXT:    v_cmp_lt_f32_e64 s[0:1], 0, v1
-; GFX950-NEXT:    s_nop 1
-; GFX950-NEXT:    v_cndmask_b32_e64 v1, v2, v3, s[0:1]
-; GFX950-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX950-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc
-; GFX950-NEXT:    v_mov_b32_e32 v2, 0x260
-; GFX950-NEXT:    v_cmp_class_f32_e32 vcc, v0, v2
-; GFX950-NEXT:    s_nop 1
-; GFX950-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc
 ; GFX950-NEXT:    v_cvt_pk_bf16_f32 v0, v0, s0
 ; GFX950-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
 ; GFX950-NEXT:    v_div_scale_f32 v1, s[0:1], v0, v0, 1.0
 ; GFX950-NEXT:    v_rcp_f32_e32 v2, v1
-; GFX950-NEXT:    s_nop 0
-; GFX950-NEXT:    v_fma_f32 v3, -v1, v2, 1.0
-; GFX950-NEXT:    v_fmac_f32_e32 v2, v3, v2
 ; GFX950-NEXT:    v_div_scale_f32 v3, vcc, 1.0, v0, 1.0
+; GFX950-NEXT:    v_fma_f32 v4, -v1, v2, 1.0
+; GFX950-NEXT:    v_fmac_f32_e32 v2, v4, v2
 ; GFX950-NEXT:    v_mul_f32_e32 v4, v3, v2
 ; GFX950-NEXT:    v_fma_f32 v5, -v1, v4, v3
 ; GFX950-NEXT:    v_fmac_f32_e32 v4, v5, v2
@@ -28817,22 +28648,7 @@ define bfloat @v_rsq_bf16(bfloat %x) {
 ; GFX10:       ; %bb.0:
 ; GFX10-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX10-NEXT:    v_lshlrev_b32_e32 v0, 16, v0
-; GFX10-NEXT:    v_mul_f32_e32 v1, 0x4f800000, v0
-; GFX10-NEXT:    v_cmp_gt_f32_e32 vcc_lo, 0xf800000, v0
-; GFX10-NEXT:    v_cndmask_b32_e32 v0, v0, v1, vcc_lo
-; GFX10-NEXT:    v_sqrt_f32_e32 v1, v0
-; GFX10-NEXT:    v_add_nc_u32_e32 v2, -1, v1
-; GFX10-NEXT:    v_add_nc_u32_e32 v3, 1, v1
-; GFX10-NEXT:    v_fma_f32 v4, -v2, v1, v0
-; GFX10-NEXT:    v_fma_f32 v5, -v3, v1, v0
-; GFX10-NEXT:    v_cmp_ge_f32_e64 s4, 0, v4
-; GFX10-NEXT:    v_cndmask_b32_e64 v1, v1, v2, s4
-; GFX10-NEXT:    v_cmp_lt_f32_e64 s4, 0, v5
-; GFX10-NEXT:    v_cndmask_b32_e64 v1, v1, v3, s4
-; GFX10-NEXT:    v_mul_f32_e32 v2, 0x37800000, v1
-; GFX10-NEXT:    v_cndmask_b32_e32 v1, v1, v2, vcc_lo
-; GFX10-NEXT:    v_cmp_class_f32_e64 vcc_lo, v0, 0x260
-; GFX10-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc_lo
+; GFX10-NEXT:    v_sqrt_f32_e32 v0, v0
 ; GFX10-NEXT:    v_bfe_u32 v1, v0, 16, 1
 ; GFX10-NEXT:    v_or_b32_e32 v2, 0x400000, v0
 ; GFX10-NEXT:    v_cmp_u_f32_e32 vcc_lo, v0, v0
@@ -28863,30 +28679,9 @@ define bfloat @v_rsq_bf16(bfloat %x) {
 ; GFX11TRUE16-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX11TRUE16-NEXT:    v_mov_b16_e32 v1.l, 0
 ; GFX11TRUE16-NEXT:    v_mov_b16_e32 v1.h, v0.l
-; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_1) | instid1(VALU_DEP_1)
-; GFX11TRUE16-NEXT:    v_cmp_gt_f32_e32 vcc_lo, 0xf800000, v1
-; GFX11TRUE16-NEXT:    v_mul_f32_e32 v0, 0x4f800000, v1
-; GFX11TRUE16-NEXT:    v_cndmask_b32_e32 v0, v1, v0, vcc_lo
-; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_3) | instid1(VALU_DEP_2)
-; GFX11TRUE16-NEXT:    v_sqrt_f32_e32 v1, v0
+; GFX11TRUE16-NEXT:    s_delay_alu instid0(VALU_DEP_1) | instskip(SKIP_4) | instid1(VALU_DEP_3)
+; GFX11TRUE16-NEXT:    v_sqrt_f32_e32 v0, v1
 ; GFX11TRUE16-NEXT:    s_waitcnt_depctr depctr_va_vdst(0)
-; GFX11TRUE16-NEXT:    v_ad...
[truncated]

github-actions · 2026-02-06T22:30:49Z

🐧 Linux x64 Test Results

189863 tests passed
5064 tests skipped

✅ The build succeeded and all tests passed.

github-actions · 2026-02-06T22:30:49Z

🪟 Windows x64 Test Results

130705 tests passed
2925 tests skipped

✅ The build succeeded and all tests passed.

arsenm · 2026-02-07T09:22:33Z

so the old code had a scaling workaround (multiply by 2^32)

I wouldn't call this a workaround, it's just how you need to implement it.

to handle this in denorm-preserve mode. This PR removes that workaround. (For f16, this issue doesn't exist: f16 denormals become normal f32 values after extension, so lowerFSQRTF16 is lossless.)

Is dropping denormal preservation for bf16 sqrt acceptable?

This should yield the correct result by default. If this degrades the result away from the correctly rounded result, it may require approx funds. What numeric testing have you tried this on?

aeft · 2026-02-08T01:49:37Z

What numeric testing have you tried this on?

https://godbolt.org/z/MEP9rMEx8

I wrote a test over all positive normal bf16 values. For each value, it computes the correctly-rounded f32 sqrt, then checks whether perturbing the f32 result by ±1 ULP (matching v_sqrt_f32's error spec) changes the bf16 rounding. No mismatches were found.

This shows that the correction in lowerFSQRTF32 is unnecessary for bf16: v_sqrt_f32's 1 ULP accuracy is already sufficient.

aeft · 2026-02-08T02:21:51Z

Hi, @arsenm. I added the scale-up/down logic to handle the denormal value. This should yield the correct result as the current behavior.

I have two questions and appreciate your feedback:

I use needsDenormHandlingF32 in the new lowerFSQRTBF16. This is consistent with the current behavior, since the old Promote path also checks needsDenormHandlingF32 in lowerFSQRTF32 after promoting to f32. Is this okay?
I don't have access to AMDGPU hardware. If on-hardware verification is needed (e.g., compare results for all bf16 values), could anyone help with that? Or do you have any recommendations? I’m happy to provide a small reproducer or scripts to run the check.

arsenm

Have you done any numeric testing on this?

arsenm · 2026-02-09T16:41:09Z

@@ -13184,6 +13187,45 @@ SDValue SITargetLowering::lowerFSQRTF64(SDValue Op, SelectionDAG &DAG) const {
                     Flags);
 }

+SDValue SITargetLowering::lowerFSQRTBF16(SDValue Op, SelectionDAG &DAG) const {


This is mostly duplicating the f32 path. Can you merge these into one function?

aeft · 2026-02-09T17:30:24Z

What numeric testing have you tried this on?

https://godbolt.org/z/MEP9rMEx8

I wrote a test over all positive normal bf16 values. For each value, it computes the correctly-rounded f32 sqrt, then checks whether perturbing the f32 result by ±1 ULP (matching v_sqrt_f32's error spec) changes the bf16 rounding. No mismatches were found.

This shows that the correction in lowerFSQRTF32 is unnecessary for bf16: v_sqrt_f32's 1 ULP accuracy is already sufficient.

Hi, @arsenm. Is this the numeric testing you were looking for?

github-actions · 2026-02-09T18:42:38Z

✅ With the latest revision this PR passed the C/C++ code formatter.

arsenm · 2026-02-09T21:00:05Z

Hi, @arsenm. Is this the numeric testing you were looking for?

I was more hoping for some preexisting spec test, that will be run with this implementation routinely. This example probably indicates it's correct, but it would be nice to have something to go on in the test suite

aeft · 2026-02-09T21:43:35Z

Hi @arsenm

I searched for existing spec tests that would routinely cover this. The closest I found are:

libc/test/src/math/smoke/sqrtbf16_test.cpp : smoke test for libc's sqrtbf16(), only tests special values
(NaN, Inf, 0, small integers)
libc/test/src/math/exhaustive/sqrtf_test.cpp : exhaustive f32 test using MPFR

However, these test the CPU-side libc implementation, not the AMDGPU codegen path. I don't see a preexisting runtime accuracy test for AMDGPU math operations. The GPU integration test infrastructure (llvm-gpu-loader) exists but is only used for libc functionality tests (malloc, printf, etc.), not math accuracy.

Could you please point me to what you had in mind? Would it be helpful to add an exhaustive runtime test for AMDGPU bf16 sqrt (like sqrtf_test.cpp)? I'd need some guidance on the right test infrastructure for that. Thanks!

aeft · 2026-02-22T04:54:26Z

Gentle ping @arsenm. Hi, just checking in to see if you've had a chance to look at the latest updates and my questions regarding the spec tests. Thanks!

arsenm · 2026-04-10T20:16:36Z

I don't like the NeedsCorrection argument. Can you reframe this as the original value type?

arsenm · 2026-04-10T20:19:48Z

Could you please point me to what you had in mind? Would it be helpful to add an exhaustive runtime test for AMDGPU bf16 sqrt (like sqrtf_test.cpp)? I'd need some guidance on the right test infrastructure for that. Thanks!

We have a little bit of underdeveloped and underused tests in https://github.com/llvm/llvm-test-suite/tree/main/External/HIP, and there's also https://github.com/llvm/offload-test-suite which I haven't tried to use yet. llvm/llvm-test-suite#370 adds a vaguely similar test, but I would hope we would have exhaustive coverage for all the 16-bit cases

aeft · 2026-04-14T06:23:28Z

Thanks for your suggestion! I will take a look at the test suite.

[AMDGPU] Improve AMDGPU sqrt and inverse sqrt handling for bf16

1d2e950

llvmbot added the backend:AMDGPU label Feb 6, 2026

fixup! [AMDGPU] Improve AMDGPU sqrt and inverse sqrt handling for bf16

2edb5b8

llvmbot added the llvm:analysis Includes value tracking, cost tables and constant folding label Feb 6, 2026

fixup! [AMDGPU] Improve AMDGPU sqrt and inverse sqrt handling for bf16

7726b97

arsenm reviewed Feb 9, 2026

View reviewed changes

fixup! [AMDGPU] Improve AMDGPU sqrt and inverse sqrt handling for bf16

f0bb7dd

fixup! [AMDGPU] Improve AMDGPU sqrt and inverse sqrt handling for bf16

d91d072

aeft requested a review from arsenm February 12, 2026 02:38

arsenm reviewed Apr 10, 2026

View reviewed changes

arsenm added the floating-point Floating-point math label Apr 10, 2026

Uh oh!

Conversation

aeft commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

aeft commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🐧 Linux x64 Test Results

Uh oh!

github-actions Bot commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🪟 Windows x64 Test Results

Uh oh!

arsenm commented Feb 7, 2026

Uh oh!

aeft commented Feb 8, 2026

Uh oh!

aeft commented Feb 8, 2026

Uh oh!

arsenm left a comment

Choose a reason for hiding this comment

Uh oh!

arsenm Feb 9, 2026

Choose a reason for hiding this comment

Uh oh!

aeft Feb 9, 2026

Choose a reason for hiding this comment

Uh oh!

aeft commented Feb 9, 2026

Uh oh!

github-actions Bot commented Feb 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

arsenm commented Feb 9, 2026

Uh oh!

aeft commented Feb 9, 2026

Uh oh!

aeft commented Feb 22, 2026

Uh oh!

arsenm Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

arsenm commented Apr 10, 2026

Uh oh!

aeft commented Apr 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aeft commented Feb 6, 2026 •

edited

Loading

aeft commented Feb 6, 2026 •

edited

Loading

llvmbot commented Feb 6, 2026 •

edited

Loading

github-actions Bot commented Feb 6, 2026 •

edited

Loading

github-actions Bot commented Feb 6, 2026 •

edited

Loading

github-actions Bot commented Feb 9, 2026 •

edited

Loading