[llvm] [AMDGPU] Allow folding FMA with SGPR and immediate operand on GFX10+ (PR #72258)

[Jay Foad via llvm-commits]

https://github.com/jayfoad created https://github.com/llvm/llvm-project/pull/72258

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.


>From e0834a697f7ea7fc94295a54161f6b9f82e803d6 Mon Sep 17 00:00:00 2001
From: Jay Foad <jay.foad at amd.com>
Date: Tue, 14 Nov 2023 13:36:34 +0000
Subject: [PATCH] [AMDGPU] Allow folding FMA with SGPR and immediate operand on
 GFX10+

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.
---
 llvm/lib/Target/AMDGPU/SIInstrInfo.cpp    |  6 +++--
 llvm/test/CodeGen/AMDGPU/fmul-to-ldexp.ll | 33 ++++++++++++++++++-----
 2 files changed, 31 insertions(+), 8 deletions(-)

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