[AMDGPU] Allow folding to FMAMK with SGPR and immediate operand on GFX10+

[jayfoad]

Allow foldImmediate to create instructions like:

v_fmamk_f32 v0, s0, 0x42000000, v0

This instruction has two "scalar values": s0 and 0x42000000. On GFX10+
this is allowed. This fold was originally implemented before the
compiler supported GFX10, when all ASICs were limited to one scalar
value.

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

Author: Jay Foad (jayfoad)

Changes

Allow foldImmediate to create instructions like:

v_fmamk_f32 v0, s0, 0x42000000, v0

This instruction has two "scalar values": s0 and 0x42000000. On GFX10+
this is allowed. This fold was originally implemented before the
compiler supported GFX10, when all ASICs were limited to one scalar
value.

---

Full diff: https://github.com/llvm/llvm-project/pull/72258.diff

2 Files Affected:

• (modified) llvm/lib/Target/AMDGPU/SIInstrInfo.cpp (+4-2)
• (modified) llvm/test/CodeGen/AMDGPU/fmul-to-ldexp.ll (+27-6)

diff --git a/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp b/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
index 027b695c3bb1a74..5f78dfff1e98852 100644
--- a/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
@@ -3445,8 +3445,10 @@ bool SIInstrInfo::FoldImmediate(MachineInstr &UseMI, MachineInstr &DefMI,
         (Src1->isReg() && Src1->getReg() == Reg)) {
       MachineOperand *RegSrc =
           Src1->isReg() && Src1->getReg() == Reg ? Src0 : Src1;
-      if (!RegSrc->isReg() ||
-          RI.isSGPRClass(MRI->getRegClass(RegSrc->getReg())))
+      if (!RegSrc->isReg())
+        return false;
+      if (RI.isSGPRClass(MRI->getRegClass(RegSrc->getReg())) &&
+          ST.getConstantBusLimit(Opc) < 2)
         return false;
 
       if (!Src2->isReg() || RI.isSGPRClass(MRI->getRegClass(Src2->getReg())))
diff --git a/llvm/test/CodeGen/AMDGPU/fmul-to-ldexp.ll b/llvm/test/CodeGen/AMDGPU/fmul-to-ldexp.ll
index 5216f565c862097..0ff5ea6ee7b1f12 100644
--- a/llvm/test/CodeGen/AMDGPU/fmul-to-ldexp.ll
+++ b/llvm/test/CodeGen/AMDGPU/fmul-to-ldexp.ll
@@ -715,12 +715,33 @@ define amdgpu_ps i32 @s_mul_fma_32_f32(float inreg %x, float inreg %y) {
 ; GFX9-NEXT:    v_readfirstlane_b32 s0, v0
 ; GFX9-NEXT:    ; return to shader part epilog
 ;
-; GFX1011-LABEL: s_mul_fma_32_f32:
-; GFX1011:       ; %bb.0:
-; GFX1011-NEXT:    v_mov_b32_e32 v0, s1
-; GFX1011-NEXT:    v_fmac_f32_e64 v0, 0x42000000, s0
-; GFX1011-NEXT:    v_readfirstlane_b32 s0, v0
-; GFX1011-NEXT:    ; return to shader part epilog
+; GFX10-SDAG-LABEL: s_mul_fma_32_f32:
+; GFX10-SDAG:       ; %bb.0:
+; GFX10-SDAG-NEXT:    v_mov_b32_e32 v0, s1
+; GFX10-SDAG-NEXT:    v_fmamk_f32 v0, s0, 0x42000000, v0
+; GFX10-SDAG-NEXT:    v_readfirstlane_b32 s0, v0
+; GFX10-SDAG-NEXT:    ; return to shader part epilog
+;
+; GFX10-GISEL-LABEL: s_mul_fma_32_f32:
+; GFX10-GISEL:       ; %bb.0:
+; GFX10-GISEL-NEXT:    v_mov_b32_e32 v0, s1
+; GFX10-GISEL-NEXT:    v_fmac_f32_e64 v0, 0x42000000, s0
+; GFX10-GISEL-NEXT:    v_readfirstlane_b32 s0, v0
+; GFX10-GISEL-NEXT:    ; return to shader part epilog
+;
+; GFX11-SDAG-LABEL: s_mul_fma_32_f32:
+; GFX11-SDAG:       ; %bb.0:
+; GFX11-SDAG-NEXT:    v_mov_b32_e32 v0, s1
+; GFX11-SDAG-NEXT:    v_fmamk_f32 v0, s0, 0x42000000, v0
+; GFX11-SDAG-NEXT:    v_readfirstlane_b32 s0, v0
+; GFX11-SDAG-NEXT:    ; return to shader part epilog
+;
+; GFX11-GISEL-LABEL: s_mul_fma_32_f32:
+; GFX11-GISEL:       ; %bb.0:
+; GFX11-GISEL-NEXT:    v_mov_b32_e32 v0, s1
+; GFX11-GISEL-NEXT:    v_fmac_f32_e64 v0, 0x42000000, s0
+; GFX11-GISEL-NEXT:    v_readfirstlane_b32 s0, v0
+; GFX11-GISEL-NEXT:    ; return to shader part epilog
   %mul = fmul contract float %x, 32.0
   %fma = fadd contract float %mul, %y
   %cast = bitcast float %fma to i32

[jayfoad]

We are only successfully folding to fmamk for sdag. To make it work for gisel we need something like #72128.

[rovka]

LGTM, thanks!