[AMDGPU] Fix an assertion failure when folding into src2 of V_FMAC_F16

[foad]

Summary

D139469 "[AMDGPU] Enable OMod on more VOP3 instructions" caused an
assertion failure when trying to fold into src2 of V_FMAC_F16. It would
temporarily convert the instruction to V_FMA_F16_gfx9 and add an opsel
operand, but if the fold still failed then it would forget to remove the
opsel operand.

[arsenm]

llvm/test/CodeGen/AMDGPU/fma.f16.ll

145

Is it worth adding the same test for <2 x half>?

[foad]

llvm/test/CodeGen/AMDGPU/fma.f16.ll

145

Sure, done.

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

bool SIFoldOperands::tryAddToFoldList(SmallVectorImpl<FoldCandidate> &FoldList,
   if (!TII->isOperandLegal(*MI, OpNo, OpToFold)) {
     // Special case for v_mac_{f16, f32}_e64 if we are trying to fold into src2
     unsigned Opc = MI->getOpcode();
     unsigned NewOpc = macToMad(Opc);
     if (NewOpc != AMDGPU::INSTRUCTION_LIST_END) {
       // Check if changing this to a v_mad_{f16, f32} instruction will allow us
       // to fold the operand.
       MI->setDesc(TII->get(NewOpc));
-      if (!AMDGPU::hasNamedOperand(Opc, AMDGPU::OpName::op_sel) &&
-          AMDGPU::hasNamedOperand(NewOpc, AMDGPU::OpName::op_sel))
+      bool AddOpSel = !AMDGPU::hasNamedOperand(Opc, AMDGPU::OpName::op_sel) &&
+                      AMDGPU::hasNamedOperand(NewOpc, AMDGPU::OpName::op_sel);
+      if (AddOpSel)
         MI->addOperand(MachineOperand::CreateImm(0));
       bool FoldAsMAD = tryAddToFoldList(FoldList, MI, OpNo, OpToFold);
       if (FoldAsMAD) {
         MI->untieRegOperand(OpNo);
         return true;
       }
+      if (AddOpSel)
+        MI->removeOperand(MI->getNumExplicitOperands() - 1);
       MI->setDesc(TII->get(Opc));
     }
 
     // Special case for s_setreg_b32
     if (OpToFold->isImm()) {
       unsigned ImmOpc = 0;
       if (Opc == AMDGPU::S_SETREG_B32)
         ImmOpc = AMDGPU::S_SETREG_IMM32_B32;

llvm/test/CodeGen/AMDGPU/fma.f16.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -global-isel=0 -march=amdgcn -mcpu=gfx900 < %s | FileCheck %s -check-prefixes=GFX9,GFX9-SDAG
 ; RUN: llc -global-isel=1 -march=amdgcn -mcpu=gfx900 < %s | FileCheck %s -check-prefixes=GFX9,GFX9-GISEL
-; RUN: llc -global-isel=0 -march=amdgcn -mcpu=gfx1010 < %s | FileCheck %s -check-prefixes=GFX10
-; RUN: llc -global-isel=1 -march=amdgcn -mcpu=gfx1010 < %s | FileCheck %s -check-prefixes=GFX10
+; RUN: llc -global-isel=0 -march=amdgcn -mcpu=gfx1010 < %s | FileCheck %s -check-prefixes=GFX10,GFX10-SDAG
+; RUN: llc -global-isel=1 -march=amdgcn -mcpu=gfx1010 < %s | FileCheck %s -check-prefixes=GFX10,GFX10-GISEL
 
 declare half @llvm.fma.f16(half, half, half)
 declare half @llvm.maxnum.f16(half, half)
 
 define half @test_fma(half %x, half %y, half %z) {
 ; GFX9-LABEL: test_fma:
 ; GFX9:       ; %bb.0:
 ; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)

 ; GFX10:       ; %bb.0:
 ; GFX10-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX10-NEXT:    s_waitcnt_vscnt null, 0x0
 ; GFX10-NEXT:    v_fmamk_f16 v0, v0, 0x4200, v2
 ; GFX10-NEXT:    s_setpc_b64 s[30:31]
   %r = call half @llvm.fma.f16(half %x, half 0xH4200, half %z)
   ret half %r
 }
+
+; Regression test for a crash caused by D139469.
+define i32 @test_D139469(half %arg) {
+; GFX9-LABEL: test_D139469:
+; GFX9:       ; %bb.0: ; %bb
+; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GFX9-NEXT:    v_mul_f16_e32 v1, 0x291e, v0
+; GFX9-NEXT:    s_movk_i32 s4, 0x291e
+; GFX9-NEXT:    v_cmp_gt_f16_e32 vcc, 0, v1
+; GFX9-NEXT:    v_mov_b32_e32 v1, 0x211e
+; GFX9-NEXT:    v_fma_f16 v0, v0, s4, v1
+; GFX9-NEXT:    v_cmp_gt_f16_e64 s[4:5], 0, v0
+; GFX9-NEXT:    s_or_b64 s[4:5], vcc, s[4:5]
+; GFX9-NEXT:    v_cndmask_b32_e64 v0, 0, 1, s[4:5]
+; GFX9-NEXT:    s_setpc_b64 s[30:31]
+;
+; GFX10-SDAG-LABEL: test_D139469:
+; GFX10-SDAG:       ; %bb.0: ; %bb
+; GFX10-SDAG-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GFX10-SDAG-NEXT:    s_waitcnt_vscnt null, 0x0
+; GFX10-SDAG-NEXT:    v_mov_b32_e32 v1, 0x211e
+; GFX10-SDAG-NEXT:    v_mul_f16_e32 v2, 0x291e, v0
+; GFX10-SDAG-NEXT:    v_fmac_f16_e32 v1, 0x291e, v0
+; GFX10-SDAG-NEXT:    v_cmp_gt_f16_e32 vcc_lo, 0, v2
+; GFX10-SDAG-NEXT:    v_cmp_gt_f16_e64 s4, 0, v1
+; GFX10-SDAG-NEXT:    s_or_b32 s4, vcc_lo, s4
+; GFX10-SDAG-NEXT:    v_cndmask_b32_e64 v0, 0, 1, s4
+; GFX10-SDAG-NEXT:    s_setpc_b64 s[30:31]
+;
+; GFX10-GISEL-LABEL: test_D139469:
+; GFX10-GISEL:       ; %bb.0: ; %bb
+; GFX10-GISEL-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GFX10-GISEL-NEXT:    s_waitcnt_vscnt null, 0x0
+; GFX10-GISEL-NEXT:    s_movk_i32 s4, 0x291e
+; GFX10-GISEL-NEXT:    v_mul_f16_e32 v1, 0x291e, v0
+; GFX10-GISEL-NEXT:    v_fmaak_f16 v0, s4, v0, 0x211e
+; GFX10-GISEL-NEXT:    v_cmp_gt_f16_e32 vcc_lo, 0, v1
+; GFX10-GISEL-NEXT:    v_cmp_gt_f16_e64 s4, 0, v0
+; GFX10-GISEL-NEXT:    s_or_b32 s4, vcc_lo, s4
+; GFX10-GISEL-NEXT:    v_cndmask_b32_e64 v0, 0, 1, s4
+; GFX10-GISEL-NEXT:    s_setpc_b64 s[30:31]
+bb:
+  %i = fmul contract half %arg, 0xH291E
+  %i1 = fcmp olt half %i, 0xH0000
+  %i2 = fadd contract half %i, 0xH211E
+  %i3 = fcmp olt half %i2, 0xH0000
+  %i4 = or i1 %i1, %i3
+  %i5 = zext i1 %i4 to i32
+  ret i32 %i5
+}