[llvm] [SLP] Fix canConvertToFMA operand selection and fmul costing (PR #216425)

[Alexey Bataev via llvm-commits]

================
@@ -14382,18 +14382,44 @@ static InstructionCost canConvertToFMA(ArrayRef<Value *> VL,
   InstructionsCompatibilityAnalysis Analysis(DT, DL, TTI, TLI);
   SmallVector<BoUpSLP::ValueList> Operands = Analysis.buildOperands(S, VL);
 
-  InstructionsState OpS = getSameOpcode(Operands.front(), TLI);
-  if (!OpS.valid())
-    return InstructionCost::getInvalid();
-
-  if (OpS.isAltShuffle() || OpS.getOpcode() != Instruction::FMul)
-    return InstructionCost::getInvalid();
-  if (!CheckForContractable(Operands.front()))
+  // The fmul may sit on either side of the add/sub. Look past operand 0 only
+  // for chains that do not allow reassociation. A reassociative chain can be
+  // vectorized into a vector fmul feeding a reduction, which is usually
+  // better than the scalar fma chain this check protects.
+  bool AllowReassoc = any_of(VL, [](Value *V) {
+    auto *FPCI = dyn_cast<FPMathOperator>(V);
+    return FPCI && FPCI->getFastMathFlags().allowReassoc();
+  });
+  auto GetFMulOperandIdx = [&]() -> std::optional<unsigned> {
+    for (unsigned Idx : seq<unsigned>(0, AllowReassoc ? 1 : Operands.size())) {
+      InstructionsState CandS = getSameOpcode(Operands[Idx], TLI);
+      if (!CandS.valid() || CandS.isAltShuffle() ||
+          CandS.getOpcode() != Instruction::FMul)
+        continue;
+      if (!CheckForContractable(Operands[Idx]))
----------------
alexey-bataev wrote:

Yes, need to pass an extra argument for InstructionsState

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