[llvm] [AMDGPU][True16] Fix MadMix selection (PR #205431)

[Petar Avramovic via llvm-commits]

================
@@ -7299,21 +7311,134 @@ AMDGPUInstructionSelector::selectVOP3PMadMixMods(MachineOperand &Root) const {
   Register Src;
   unsigned Mods;
   bool Matched;
-  bool NeedsWiden;
-  std::tie(Src, Mods) = selectVOP3PMadMixModsImpl(Root, Matched, NeedsWiden);
+  std::tie(Src, Mods) = selectVOP3PMadMixModsImpl(Root, Matched);
+
+  if (madMixSrcNeedsWiden(Src))
+    std::tie(Src, Mods) = selectVOP3ModsImpl(Root.getReg());
 
-  MachineRegisterInfo *RegInfo = MRI;
   return {{
-      [=](MachineInstrBuilder &MIB) {
-        Register Reg = Src;
-        if (NeedsWiden)
-          Reg = createVOP3PSrc32FromLo16(Src, MIB.getInstr(), *RegInfo);
-        MIB.addReg(Reg);
-      },
+      [=](MachineInstrBuilder &MIB) { MIB.addReg(Src); },
       [=](MachineInstrBuilder &MIB) { MIB.addImm(Mods); } // src_mods
   }};
 }
 
+// With real true16 instructions a 16-bit value lives in a VGPR_16, which cannot
+// be used as a source of the 32-bit mix instructions. Widening it needs a
+// REG_SEQUENCE
+bool AMDGPUInstructionSelector::selectVOP3PMadMixF32(MachineInstr &I) const {
+  if (!Subtarget->useRealTrue16Insts() || !Subtarget->hasFmaMixInsts())
+    return false;
+
+  Register Dst = I.getOperand(0).getReg();
+  if (MRI->getType(Dst) != LLT::scalar(32) ||
+      RBI.getRegBank(Dst, *MRI, TRI)->getID() != AMDGPU::VGPRRegBankID)
+    return false;
+
+  struct MixSrc {
+    Register Reg;
+    int64_t Imm = 0;
+    unsigned Mods = SISrcMods::NONE;
+    bool IsImm = false;
+    bool NeedsWiden = false;
+  } Srcs[3];
+
+  const auto MatchReg = [&](MixSrc &S, MachineOperand &Op) {
+    bool Matched;
+    std::tie(S.Reg, S.Mods) = selectVOP3PMadMixModsImpl(Op, Matched);
+    S.NeedsWiden = madMixSrcNeedsWiden(S.Reg);
+  };
+  const auto SetImm = [](MixSrc &S, int64_t Imm, unsigned Mods) {
+    S.IsImm = true;
+    S.Imm = Imm;
+    S.Mods = Mods;
+  };
+
+  switch (I.getOpcode()) {
+  case TargetOpcode::G_FMA:
+    MatchReg(Srcs[0], I.getOperand(1));
+    MatchReg(Srcs[1], I.getOperand(2));
+    MatchReg(Srcs[2], I.getOperand(3));
+    break;
+  case TargetOpcode::G_FADD:
+    // (fadd x, y) -> (fma x, 1.0, y)
+    MatchReg(Srcs[0], I.getOperand(1));
+    SetImm(Srcs[1], 0x3C00 /*half 1.0*/, SISrcMods::OP_SEL_1);
+    MatchReg(Srcs[2], I.getOperand(2));
+    break;
+  case TargetOpcode::G_FMUL:
+    // (fmul x, y) -> (fma x, y, -0.0)
+    MatchReg(Srcs[0], I.getOperand(1));
+    MatchReg(Srcs[1], I.getOperand(2));
+    SetImm(Srcs[2], 0, SISrcMods::NEG);
+    break;
+  case TargetOpcode::G_FSUB:
+    // (fsub x, y) -> (fma y, -1.0, x)
+    MatchReg(Srcs[0], I.getOperand(2));
+    SetImm(Srcs[1], 0xBC00 /*half -1.0*/, SISrcMods::OP_SEL_1);
+    MatchReg(Srcs[2], I.getOperand(1));
+    break;
+  default:
+    return false;
+  }
+
+  if (none_of(Srcs, [](const MixSrc &S) { return S.NeedsWiden; }))
+    return false;
+
+  // Constrain the sources that are about to be widened before creating
+  // anything, so that a failure here leaves the function untouched.
+  for (const MixSrc &S : Srcs) {
----------------
petar-avramovic wrote:

Think it should be possible to merge all three
  `for (const MixSrc &S : Srcs) {`
loops

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