[llvm] [AMDGPU] Add OMOD folding for TRANS bfloat16 instructions (PR #208595)

[Jay Foad via llvm-commits]

================
@@ -2363,6 +2358,117 @@ SIFoldOperandsImpl::isOMod(const MachineInstr &MI) const {
 
     return std::pair(nullptr, SIOutMods::NONE);
   }
+  case AMDGPU::V_FMA_MIXLO_BF16:
+  case AMDGPU::V_FMA_MIX_BF16_t16: {
+    if (MFI->getMode().FP32Denormals.Output != DenormalMode::PreserveSign ||
+        MI.mayRaiseFPException())
+      return std::pair(nullptr, SIOutMods::NONE);
+
+    const MachineOperand *Src0 = TII->getNamedOperand(MI, AMDGPU::OpName::src0);
+    const MachineOperand *Src1 = TII->getNamedOperand(MI, AMDGPU::OpName::src1);
+    const MachineOperand *Src2 = TII->getNamedOperand(MI, AMDGPU::OpName::src2);
+
+    if (!Src2 || !Src2->isImm() || Src2->getImm() != 0)
+      return std::pair(nullptr, SIOutMods::NONE);
+
+    // op_sel bits are in src*_modifiers but don't affect folding
+    const MachineOperand *Src0Mods =
+        TII->getNamedOperand(MI, AMDGPU::OpName::src0_modifiers);
+    const MachineOperand *Src1Mods =
+        TII->getNamedOperand(MI, AMDGPU::OpName::src1_modifiers);
+    if ((Src0Mods &&
+         (Src0Mods->getImm() & (SISrcMods::NEG | SISrcMods::ABS))) ||
+        (Src1Mods &&
+         (Src1Mods->getImm() & (SISrcMods::NEG | SISrcMods::ABS))) ||
+        TII->hasModifiersSet(MI, AMDGPU::OpName::omod) ||
+        TII->hasModifiersSet(MI, AMDGPU::OpName::clamp))
+      return std::pair(nullptr, SIOutMods::NONE);
+
+    if (!Src0->isReg() || !Src1->isReg())
+      return std::pair(nullptr, SIOutMods::NONE);
+
+    const MachineOperand *RegOp = nullptr;
+    int64_t ImmValue = 0;
+
+    const MachineInstr *Src1Def = MRI->getVRegDef(Src1->getReg());
+    if (Src1Def && Src1Def->isMoveImmediate() &&
+        Src1Def->getOperand(1).isImm()) {
+      ImmValue = Src1Def->getOperand(1).getImm();
+      RegOp = Src0;
+    } else {
+      const MachineInstr *Src0Def = MRI->getVRegDef(Src0->getReg());
+      if (Src0Def && Src0Def->isMoveImmediate() &&
+          Src0Def->getOperand(1).isImm()) {
+        ImmValue = Src0Def->getOperand(1).getImm();
+        RegOp = Src1;
+      } else {
+        return std::pair(nullptr, SIOutMods::NONE);
+      }
+    }
+
+    int OMod = getOModValue(AMDGPU::V_MUL_F32_e64, ImmValue);
+    if (OMod == SIOutMods::NONE)
+      return std::pair(nullptr, SIOutMods::NONE);
+
+    return std::pair(RegOp, OMod);
+  }
+  case AMDGPU::V_CVT_PK_BF16_F32_e64: {
+    // Fold cvt_pk_bf16_f32(fma_mix_f32_bf16(x, 1.0, x)) -> x with MUL2 omod
+    if (MFI->getMode().FP32Denormals.Output != DenormalMode::PreserveSign ||
----------------
jayfoad wrote:

> > I don't see V_CVT_PK_BF16_F32 in any of the tests you added.
> 
> %mul2 = fmul nsz bfloat %cos, 2.0 ==> something like %fma = v_fma_mix_f32_bf16 %cos, 1.0. %cos %mul2 = V_CVT_PK_BF16_F32 %fma

I see, it only shows up in the test _without_ your SIFoldOperands patch. Can you precommit the test to show the current codegen?

> > I'd also prefer this part in a separate PR because I think it's the first time isOMod is handling something with vector type, and we need to be careful about what that means: does it apply to all of the vector elements or just the first one? You seem to be assuming "just the first one".
> 
> We are handling three cases of omod:
> 
> 1. %mul2 = fmul nsz bfloat %cos, 2.0
> 2. %mul4 = fmul nsz bfloat %cos, 4.0
> 3. %div2 = fmul nsz bfloat %cos, 0.5
> 
> Do you mean we need to separate the PR for case 1? Thanks.

Yes.

https://github.com/llvm/llvm-project/pull/208595