[SLP]Flatten alternate associative chains into one reassociated node

Each lane peels only chain links with its own opcode, keeping the
per-lane opcode on every combine level; emission linearizes into
main-opcode chain, alt-opcode chain, and a single select shuffle.

Reviewers: hiraditya, RKSimon, bababuck

Pull Request: #215098

Author: alexey-bataev

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -14,6 +14,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/InstrTypes.h"
@@ -808,4 +809,69 @@ bool isAlternateInstruction(Instruction *I, Instruction *MainOp,
   }
   return InstructionsState(MainOp, AltOp).getMatchingMainOpOrAltOp(I) == AltOp;
 }
+
+SmallVector<SmallVector<Value *>> scanAltAssociativeOperands(
+    const InstructionsState &S, const TargetLibraryInfo &TLI,
+    ArrayRef<Value *> VL, ArrayRef<Value *> Op0, ArrayRef<Value *> Op1,
+    SmallVectorImpl<Value *> &ReassocScalars, SmallBitVector &SubLanes) {
+  assert(S.isAltShuffle() && "Expected an alternate node.");
+  const unsigned NumLanes = VL.size();
+  SmallVector<unsigned> LaneOpcodes =
+      map_to_vector(seq<unsigned>(NumLanes), [&](unsigned Lane) {
+        return isAlternateInstruction(cast<Instruction>(VL[Lane]),
+                                      S.getMainOp(), S.getAltOp(), TLI)
+                   ? S.getAltOpcode()
+                   : S.getOpcode();
+      });
+  // A lane value peels only as a single-use chain link with the lane's own
+  // opcode, keeping every combine level on the same main/alt pattern.
+  auto GetChainLink = [&](unsigned Lane, Value *V) -> Instruction * {
+    auto *I = dyn_cast<Instruction>(V);
+    if (!I || !I->hasOneUse() || I->getOpcode() != LaneOpcodes[Lane] ||
+        !isReassocChainLink(I))
+      return nullptr;
+    return I;
+  };
+  SmallVector<SmallVector<Value *>> Columns;
+  Columns.emplace_back(Op0.begin(), Op0.end());
+  Columns.emplace_back(Op1.begin(), Op1.end());
+  // The chain link of a commutative lane may sit in the second column;
+  // normalize so every lane's link leads.
+  for (unsigned Lane : seq<unsigned>(NumLanes)) {
+    if (GetChainLink(Lane, Columns[0][Lane]))
+      continue;
+    Instruction *Link = GetChainLink(Lane, Columns[1][Lane]);
+    if (!Link || !Link->isCommutative())
+      return {};
+    std::swap(Columns[0][Lane], Columns[1][Lane]);
+  }
+  // Peel the leading column while every lane stays a matching chain link.
+  while (all_of(seq<unsigned>(NumLanes), [&](unsigned Lane) {
+    return GetChainLink(Lane, Columns[0][Lane]) != nullptr;
+  })) {
+    SmallVector<Value *> NewColumn(NumLanes);
+    for (unsigned Lane : seq<unsigned>(NumLanes)) {
+      Instruction *Link = GetChainLink(Lane, Columns[0][Lane]);
+      ReassocScalars.push_back(Link);
+      // The chain of a commutative lane may continue in the second operand;
+      // keep the chain link as the running value.
+      unsigned RunningOp = Link->isCommutative() &&
+                                   !GetChainLink(Lane, Link->getOperand(0)) &&
+                                   GetChainLink(Lane, Link->getOperand(1))
+                               ? 1
+                               : 0;
+      NewColumn[Lane] = Link->getOperand(1 - RunningOp);
+      Columns[0][Lane] = Link->getOperand(RunningOp);
+    }
+    Columns.insert(std::next(Columns.begin()), std::move(NewColumn));
+  }
+  assert(!ReassocScalars.empty() &&
+         "Normalization guarantees at least one peeled level.");
+  SubLanes.resize(NumLanes);
+  for (unsigned Lane : seq<unsigned>(NumLanes))
+    if (LaneOpcodes[Lane] == Instruction::Sub ||
+        LaneOpcodes[Lane] == Instruction::FSub)
+      SubLanes.set(Lane);
+  return Columns;
+}
 } // namespace llvm::slpvectorizer

llvm/lib/Transforms/Vectorize/SLPVectorizer/SLPCompatibilityAnalysis.cpp

@@ -19,6 +19,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/IVDescriptors.h"
 #include "llvm/IR/Instruction.h"
@@ -306,6 +307,18 @@ convertTo(Instruction *I, const InstructionsState &S);
 /// the given \p MainOp and \p AltOp instructions.
 bool isAlternateInstruction(Instruction *I, Instruction *MainOp,
                             Instruction *AltOp, const TargetLibraryInfo &TLI);
+
+/// Peel the per-lane associative chains of an alternate node into operand
+/// columns. Lanes peel in lockstep and only chain links with the lane's own
+/// opcode, so every combine level keeps the root's main/alt opcode pattern
+/// and a subtract lane never becomes an add of a negated leaf. Only the
+/// leading (running) column peels: peeling a subtracted subtract would flip
+/// signs. \p SubLanes records the subtract lanes for the realignment sign
+/// query. Returns the flattened columns, empty when no level peels.
+SmallVector<SmallVector<Value *>> scanAltAssociativeOperands(
+    const InstructionsState &S, const TargetLibraryInfo &TLI,
+    ArrayRef<Value *> VL, ArrayRef<Value *> Op0, ArrayRef<Value *> Op1,
+    SmallVectorImpl<Value *> &ReassocScalars, SmallBitVector &SubLanes);
 } // namespace llvm::slpvectorizer
 
 #endif // LLVM_LIB_TRANSFORMS_VECTORIZE_SLPVECTORIZER_SLPCOMPATIBILITYANALYSIS_H

llvm/lib/Transforms/Vectorize/SLPVectorizer/SLPCompatibilityAnalysis.h

@@ -4,17 +4,7 @@
 define void @test() {
 ; CHECK-LABEL: @test(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = lshr i64 0, 0
-; CHECK-NEXT:    [[TMP1:%.*]] = sub i64 0, 1
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr i64 [[TMP1]], 0
-; CHECK-NEXT:    [[UMIN120:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP0]], i64 [[TMP2]])
-; CHECK-NEXT:    [[TMP3:%.*]] = sub i64 0, 0
-; CHECK-NEXT:    [[TMP4:%.*]] = lshr i64 [[TMP3]], 0
-; CHECK-NEXT:    [[UMIN122:%.*]] = call i64 @llvm.umin.i64(i64 [[UMIN120]], i64 [[TMP4]])
-; CHECK-NEXT:    [[TMP5:%.*]] = add i64 0, 1
-; CHECK-NEXT:    [[TMP6:%.*]] = lshr i64 [[TMP5]], 0
-; CHECK-NEXT:    [[UMIN123:%.*]] = call i64 @llvm.umin.i64(i64 [[UMIN122]], i64 [[TMP6]])
-; CHECK-NEXT:    [[UMIN124:%.*]] = call i64 @llvm.umin.i64(i64 [[UMIN123]], i64 0)
+; CHECK-NEXT:    [[UMIN124:%.*]] = call i64 @llvm.umin.i64(i64 0, i64 0)
 ; CHECK-NEXT:    ret void
 ;
 entry:

llvm/test/Transforms/SLPVectorizer/X86/BinOpSameOpcodeHelper.ll

@@ -89,13 +89,10 @@ define void @test_supernode_addsub_alt(ptr %Aarray, ptr %Barray, ptr %Carray, pt
 ; ENABLED-NEXT:    [[TMP0:%.*]] = load <2 x double>, ptr [[AARRAY:%.*]], align 8
 ; ENABLED-NEXT:    [[TMP1:%.*]] = load <2 x double>, ptr [[BARRAY:%.*]], align 8
 ; ENABLED-NEXT:    [[TMP2:%.*]] = load <2 x double>, ptr [[CARRAY:%.*]], align 8
-; ENABLED-NEXT:    [[TMP3:%.*]] = shufflevector <2 x double> [[TMP0]], <2 x double> [[TMP2]], <2 x i32> <i32 0, i32 3>
-; ENABLED-NEXT:    [[TMP4:%.*]] = fsub fast <2 x double> [[TMP3]], [[TMP1]]
-; ENABLED-NEXT:    [[TMP5:%.*]] = fadd fast <2 x double> [[TMP3]], [[TMP1]]
-; ENABLED-NEXT:    [[TMP6:%.*]] = shufflevector <2 x double> [[TMP4]], <2 x double> [[TMP5]], <2 x i32> <i32 0, i32 3>
-; ENABLED-NEXT:    [[TMP7:%.*]] = shufflevector <2 x double> [[TMP2]], <2 x double> [[TMP0]], <2 x i32> <i32 0, i32 3>
-; ENABLED-NEXT:    [[TMP8:%.*]] = fsub fast <2 x double> [[TMP6]], [[TMP7]]
-; ENABLED-NEXT:    [[TMP9:%.*]] = fadd fast <2 x double> [[TMP6]], [[TMP7]]
+; ENABLED-NEXT:    [[TMP4:%.*]] = fsub reassoc nsz arcp contract afn <2 x double> [[TMP0]], [[TMP1]]
+; ENABLED-NEXT:    [[TMP8:%.*]] = fsub reassoc nsz arcp contract afn <2 x double> [[TMP4]], [[TMP2]]
+; ENABLED-NEXT:    [[TMP5:%.*]] = fadd reassoc nsz arcp contract afn <2 x double> [[TMP0]], [[TMP1]]
+; ENABLED-NEXT:    [[TMP9:%.*]] = fadd reassoc nsz arcp contract afn <2 x double> [[TMP5]], [[TMP2]]
 ; ENABLED-NEXT:    [[TMP10:%.*]] = shufflevector <2 x double> [[TMP8]], <2 x double> [[TMP9]], <2 x i32> <i32 0, i32 3>
 ; ENABLED-NEXT:    store <2 x double> [[TMP10]], ptr [[SARRAY:%.*]], align 8
 ; ENABLED-NEXT:    ret void
@@ -133,16 +130,12 @@ define void @test_addsub_alt_commuted(ptr %Aarray, ptr %Barray, ptr %Carray, ptr
 ; ENABLED-NEXT:    [[TMP1:%.*]] = load <2 x double>, ptr [[BARRAY:%.*]], align 8
 ; ENABLED-NEXT:    [[TMP2:%.*]] = load <2 x double>, ptr [[CARRAY:%.*]], align 8
 ; ENABLED-NEXT:    [[TMP3:%.*]] = load <2 x double>, ptr [[DARRAY:%.*]], align 8
-; ENABLED-NEXT:    [[TMP4:%.*]] = shufflevector <2 x double> [[TMP0]], <2 x double> [[TMP2]], <2 x i32> <i32 0, i32 3>
-; ENABLED-NEXT:    [[TMP5:%.*]] = fsub fast <2 x double> [[TMP4]], [[TMP1]]
-; ENABLED-NEXT:    [[TMP12:%.*]] = fadd fast <2 x double> [[TMP4]], [[TMP1]]
-; ENABLED-NEXT:    [[TMP7:%.*]] = shufflevector <2 x double> [[TMP5]], <2 x double> [[TMP12]], <2 x i32> <i32 0, i32 3>
-; ENABLED-NEXT:    [[TMP8:%.*]] = fsub fast <2 x double> [[TMP7]], [[TMP3]]
-; ENABLED-NEXT:    [[TMP13:%.*]] = fadd fast <2 x double> [[TMP7]], [[TMP3]]
-; ENABLED-NEXT:    [[TMP14:%.*]] = shufflevector <2 x double> [[TMP8]], <2 x double> [[TMP13]], <2 x i32> <i32 0, i32 3>
-; ENABLED-NEXT:    [[TMP11:%.*]] = shufflevector <2 x double> [[TMP2]], <2 x double> [[TMP0]], <2 x i32> <i32 0, i32 3>
-; ENABLED-NEXT:    [[TMP6:%.*]] = fsub fast <2 x double> [[TMP14]], [[TMP11]]
-; ENABLED-NEXT:    [[TMP9:%.*]] = fadd fast <2 x double> [[TMP14]], [[TMP11]]
+; ENABLED-NEXT:    [[TMP4:%.*]] = fsub reassoc nsz arcp contract afn <2 x double> [[TMP0]], [[TMP1]]
+; ENABLED-NEXT:    [[TMP5:%.*]] = fsub reassoc nsz arcp contract afn <2 x double> [[TMP4]], [[TMP3]]
+; ENABLED-NEXT:    [[TMP6:%.*]] = fsub reassoc nsz arcp contract afn <2 x double> [[TMP5]], [[TMP2]]
+; ENABLED-NEXT:    [[TMP7:%.*]] = fadd reassoc nsz arcp contract afn <2 x double> [[TMP0]], [[TMP1]]
+; ENABLED-NEXT:    [[TMP8:%.*]] = fadd reassoc nsz arcp contract afn <2 x double> [[TMP7]], [[TMP3]]
+; ENABLED-NEXT:    [[TMP9:%.*]] = fadd reassoc nsz arcp contract afn <2 x double> [[TMP8]], [[TMP2]]
 ; ENABLED-NEXT:    [[TMP10:%.*]] = shufflevector <2 x double> [[TMP6]], <2 x double> [[TMP9]], <2 x i32> <i32 0, i32 3>
 ; ENABLED-NEXT:    store <2 x double> [[TMP10]], ptr [[SARRAY:%.*]], align 8
 ; ENABLED-NEXT:    ret void

llvm/test/Transforms/SLPVectorizer/X86/supernode.ll

@@ -5,24 +5,13 @@
 define void @foo(ptr %c, ptr %d) {
 ; X86-LABEL: @foo(
 ; X86-NEXT:  entry:
-; X86-NEXT:    [[ARRAYIDX1:%.*]] = getelementptr inbounds i8, ptr [[C:%.*]], i64 4
-; X86-NEXT:    [[ARRAYIDX4:%.*]] = getelementptr inbounds i8, ptr [[C]], i64 1
-; X86-NEXT:    [[ARRAYIDX12:%.*]] = getelementptr inbounds i8, ptr [[C]], i64 2
+; X86-NEXT:    [[ARRAYIDX4:%.*]] = getelementptr inbounds i8, ptr [[C:%.*]], i64 1
 ; X86-NEXT:    [[ADD_PTR53:%.*]] = getelementptr inbounds float, ptr [[D:%.*]], i64 -4
-; X86-NEXT:    [[TMP0:%.*]] = load i8, ptr [[ARRAYIDX4]], align 1
-; X86-NEXT:    [[TMP1:%.*]] = load i8, ptr [[ARRAYIDX1]], align 1
-; X86-NEXT:    [[CONV5:%.*]] = zext i8 [[TMP0]] to i32
-; X86-NEXT:    [[CONV2:%.*]] = zext i8 [[TMP1]] to i32
-; X86-NEXT:    [[SHL6:%.*]] = shl nuw nsw i32 [[CONV5]], 2
-; X86-NEXT:    [[AND:%.*]] = and i32 [[CONV2]], 3
-; X86-NEXT:    [[TMP2:%.*]] = load <2 x i8>, ptr [[ARRAYIDX12]], align 1
-; X86-NEXT:    [[TMP3:%.*]] = zext <2 x i8> [[TMP2]] to <2 x i16>
-; X86-NEXT:    [[TMP4:%.*]] = shl <2 x i16> [[TMP3]], splat (i16 2)
-; X86-NEXT:    [[TMP5:%.*]] = insertelement <4 x i32> poison, i32 [[SHL6]], i64 0
-; X86-NEXT:    [[TMP6:%.*]] = zext <2 x i16> [[TMP4]] to <2 x i32>
-; X86-NEXT:    [[TMP7:%.*]] = shufflevector <2 x i32> [[TMP6]], <2 x i32> poison, <4 x i32> <i32 0, i32 1, i32 poison, i32 poison>
-; X86-NEXT:    [[TMP8:%.*]] = shufflevector <4 x i32> [[TMP5]], <4 x i32> [[TMP7]], <4 x i32> <i32 0, i32 4, i32 5, i32 poison>
-; X86-NEXT:    [[TMP9:%.*]] = insertelement <4 x i32> [[TMP8]], i32 [[AND]], i64 3
+; X86-NEXT:    [[TMP0:%.*]] = load <4 x i8>, ptr [[ARRAYIDX4]], align 1
+; X86-NEXT:    [[TMP1:%.*]] = zext <4 x i8> [[TMP0]] to <4 x i32>
+; X86-NEXT:    [[TMP2:%.*]] = shl nuw nsw <4 x i32> [[TMP1]], <i32 2, i32 2, i32 2, i32 3>
+; X86-NEXT:    [[TMP3:%.*]] = and <4 x i32> [[TMP1]], <i32 2, i32 2, i32 2, i32 3>
+; X86-NEXT:    [[TMP9:%.*]] = shufflevector <4 x i32> [[TMP2]], <4 x i32> [[TMP3]], <4 x i32> <i32 0, i32 1, i32 2, i32 7>
 ; X86-NEXT:    [[TMP10:%.*]] = add nsw <4 x i32> undef, [[TMP9]]
 ; X86-NEXT:    [[TMP11:%.*]] = sitofp <4 x i32> [[TMP10]] to <4 x float>
 ; X86-NEXT:    [[TMP12:%.*]] = fdiv <4 x float> [[TMP11]], undef