[llvm] [SLP] Merge redundant input operands of binary-operator nodes (PR #208670)

[Balakrishna Bandlapalli via llvm-commits]

================
@@ -13145,9 +13256,163 @@ BoUpSLP::getScalarsVectorizationLegality(ArrayRef<Value *> VL, unsigned Depth,
   return ScalarsVectorizationLegality(S, /*IsLegal=*/true);
 }
 
+void BoUpSLP::updateReuseShuffleForMergedOps(
+    const EdgeInfo &UserTreeIdx, SmallVectorImpl<int> &ReuseShuffleIndices,
+    SmallDenseMap<Value *, unsigned, 16> &UniquePositions,
+    ArrayRef<Value *> UniqueValues, bool MergedOp) const {
+  if (!MergedOp)
+    return;
+  // The operand node is built from the merged scalar list, but the reuse
+  // shuffle mask must map the user node's operand lanes onto the (unique)
+  // merged scalars so that codegen reads the correct lane.
+  ArrayRef<Value *> InVL = UserTreeIdx.UserTE->getOperand(UserTreeIdx.EdgeIdx);
+  ReuseShuffleIndices.clear();
+  for (Value *V : InVL) {
+    auto Res = UniquePositions.try_emplace(V, UniqueValues.size());
+    ReuseShuffleIndices.emplace_back(Res.first->second);
+    assert(!Res.second && "Merged operand value not found in unique values.");
+  }
+}
+
+bool BoUpSLP::isInputVectorizable(ArrayRef<Value *> ParentVL,
+                                  ArrayRef<Value *> InVL,
+                                  SmallVectorImpl<Value *> &MergedVL) {
+  SmallVector<Value *> VL(InVL.begin(), InVL.end());
+  assert((allConstant(VL) || allSameType(VL)) && "Invalid types!");
+
+  SmallVector<int> ReuseShuffleIndices;
+  SmallVector<Value *> UniqueValues;
+  auto TryToFindDuplicates = [&]() {
+    SmallDenseMap<Value *, unsigned, 16> UniquePositions(VL.size());
+    for (Value *V : VL) {
+      if (isConstant(V)) {
+        ReuseShuffleIndices.emplace_back(
+            isa<UndefValue>(V) ? PoisonMaskElem : UniqueValues.size());
+        UniqueValues.emplace_back(V);
+        continue;
+      }
+      auto Res = UniquePositions.try_emplace(V, UniqueValues.size());
+      ReuseShuffleIndices.emplace_back(Res.first->second);
+      if (Res.second)
+        UniqueValues.emplace_back(V);
+    }
+    size_t NumUniqueScalarValues = UniqueValues.size();
+    if (NumUniqueScalarValues == VL.size()) {
+      ReuseShuffleIndices.clear();
+    } else {
+      if (NumUniqueScalarValues <= 1 ||
+          (UniquePositions.size() == 1 &&
+           all_of(UniqueValues,
+                  [](Value *V) {
+                    return isa<UndefValue>(V) || !isConstant(V);
+                  })) ||
+          !has_single_bit(NumUniqueScalarValues))
+        return false;
+      VL = UniqueValues;
+    }
+    return true;
+  };
+  InstructionsState S = getSameOpcode(VL, *TLI);
+  if (!S)
+    return false;
+  if (!isa<BinaryOperator>(S.getMainOp()))
+    return false;
+  // Check that every instruction appears once in this bundle.
+  if (!TryToFindDuplicates())
+    return false;
+
+  // Bail out if the unique operations in the input operand vector are not
+  // (at least) half of the parent VL - otherwise merging is not profitable.
+  if (PowerOf2Ceil(ParentVL.size()) / 2 < PowerOf2Ceil(VL.size()))
+    return false;
+  // The merged operands must be vectorizable on their own.
+  OrdersType CurrentOrder;
+  SmallVector<Value *> PointerOps;
+  StridedPtrInfo SPtrInfo;
+  SmallVector<int> ExpandShuffleMask;
+  TreeEntry::EntryState State = getScalarsVectorizationState(
+      S, VL, /*IsScatterVectorizeUserTE=*/false, CurrentOrder, PointerOps,
+      SPtrInfo, ExpandShuffleMask);
+  if (State != TreeEntry::Vectorize)
+    return false;
+
+  MergedVL.append(VL.begin(), VL.end());
+  return true;
+}
+
+bool BoUpSLP::canMergeInputOperands(ArrayRef<Value *> InVL,
+                                    ArrayRef<Value *> InLeft,
+                                    ArrayRef<Value *> InRight,
+                                    SmallVectorImpl<Value *> &MergedVL) {
+  // We expect the minimum size of the merged VL to be at least 4; smaller
+  // sizes are better vectorized with splats if not merged.
+  if (PowerOf2Ceil(InVL.size()) < 4)
+    return false;
+  if (!isInputVectorizable(InVL, InLeft, MergedVL) ||
+      !isInputVectorizable(InVL, InRight, MergedVL))
+    return false;
+  if (MergedVL.size() != InLeft.size() || MergedVL.size() != InRight.size())
+    return false;
+  // Both operand nodes are built from the single MergedVL bundle and are
+  // distinguished only by their reuse-shuffle masks selecting lanes out of one
+  // merged vector. For codegen to reconstruct each operand correctly:
+  //  - every merged scalar must occupy a distinct lane (otherwise the two
+  //    operand nodes can no longer be told apart by their reuse masks and they
+  //    collapse to the same vector), and
+  //  - the merged bundle must vectorize as a single uniform node. A
+  //    mixed-opcode bundle would become an alt-shuffle node whose vectorized
+  //    value is a shuffle of two different source vectors; the merged-operand
+  //    codegen only unwraps a single source operand (getShuffleInput) and
+  //    would silently drop the other half.
+  SmallPtrSet<Value *, 8> SeenVals;
+  for (Value *V : MergedVL)
+    if (!SeenVals.insert(V).second)
+      return false;
+  InstructionsState MergedS = getSameOpcode(MergedVL, *TLI);
+  if (!MergedS || MergedS.isAltShuffle())
+    return false;
+  OrdersType MergedOrder;
+  SmallVector<Value *> MergedPointerOps;
+  StridedPtrInfo MergedSPtrInfo;
+  SmallVector<int> MergedExpandShuffleMask;
+  if (getScalarsVectorizationState(MergedS, MergedVL,
+                                   /*IsScatterVectorizeUserTE=*/false,
+                                   MergedOrder, MergedPointerOps,
+                                   MergedSPtrInfo,
+                                   MergedExpandShuffleMask) != TreeEntry::Vectorize)
+    return false;
+  return true;
+}
+
+bool BoUpSLP::handleMergedOperands(TreeEntry *TE, ArrayRef<Value *> VL,
+                                   ArrayRef<Value *> Left,
+                                   ArrayRef<Value *> Right, unsigned Depth) {
+  if (!SLPMergeInputOps)
+    return false;
+  SmallVector<Value *> MergedVL;
+  if (!canMergeInputOperands(VL, Left, Right, MergedVL))
+    return false;
+  LLVM_DEBUG(dbgs() << "SLP: Merging input operands in " << F->getName()
+                    << "\n");
+
+  TE->setOperand(0, Left);
+  TE->setOperand(1, Right);
----------------
amd-bbandlap wrote:

Left and Right are preserved so that ReuseShuffleIndicies of the nodes are updated correctly as MergedVL is passed as main VL to buildTreeRec  

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