[llvm] [VectorCombine] Add subvector reduction support to foldShuffleChainsToReduce (PR #199872)
Adam Scott via llvm-commits
llvm-commits at lists.llvm.org
Mon Jun 1 19:53:13 PDT 2026
================
@@ -3973,226 +3950,317 @@ bool VectorCombine::foldShuffleChainsToReduce(Instruction &I) {
if (!FVT)
return false;
- int64_t VecSize = FVT->getNumElements();
- if (VecSize < 2)
+ if (FVT->getNumElements() < 2)
return false;
- // Number of levels would be ~log2(n), considering we always partition
- // by half for this fold pattern.
- unsigned int NumLevels = Log2_64_Ceil(VecSize), VisitedCnt = 0;
- int64_t ShuffleMaskHalf = 1, ExpectedParityMask = 0;
+ std::optional<Instruction::BinaryOps> CommonBinOp;
+ std::optional<Intrinsic::ID> CommonCallOp;
- // This is how we generalise for all element sizes.
- // At each step, if vector size is odd, we need non-poison
- // values to cover the dominant half so we don't miss out on any element.
- //
- // This mask will help us retrieve this as we go from bottom to top:
- //
- // Mask Set -> N = N * 2 - 1
- // Mask Unset -> N = N * 2
- for (int Cur = VecSize, Mask = NumLevels - 1; Cur > 1;
- Cur = (Cur + 1) / 2, --Mask) {
- if (Cur & 1)
- ExpectedParityMask |= (1ll << Mask);
+ if (auto *BO = dyn_cast<BinaryOperator>(VecOpEE)) {
+ switch (BO->getOpcode()) {
+ case Instruction::Add:
+ case Instruction::Mul:
+ case Instruction::Or:
+ case Instruction::And:
+ case Instruction::Xor:
+ CommonBinOp = BO->getOpcode();
+ break;
+ default:
+ return false;
+ }
+ } else if (auto *MMI = dyn_cast<MinMaxIntrinsic>(VecOpEE)) {
+ CommonCallOp = MMI->getIntrinsicID();
+ } else {
+ return false;
}
- InstWorklist.push(VecOpEE);
-
- bool IsPartialReduction = false;
-
- while (!InstWorklist.empty()) {
- Value *CI = InstWorklist.front();
- InstWorklist.pop();
-
- if (auto *II = dyn_cast<IntrinsicInst>(CI)) {
- if (!ShouldBeCallOrBinInst)
- return false;
-
- if (!IsFirstCallOrBinInst && any_of(PrevVecV, equal_to(nullptr)))
- return false;
-
- // For the first found call/bin op, the vector has to come from the
- // extract element op.
- if (II != (IsFirstCallOrBinInst ? VecOpEE : PrevVecV[0]))
- return false;
- IsFirstCallOrBinInst = false;
-
- if (!CommonCallOp)
- CommonCallOp = II->getIntrinsicID();
- if (II->getIntrinsicID() != *CommonCallOp)
- return false;
-
- switch (II->getIntrinsicID()) {
- case Intrinsic::umin:
- case Intrinsic::umax:
- case Intrinsic::smin:
- case Intrinsic::smax: {
- auto *Op0 = II->getOperand(0);
- auto *Op1 = II->getOperand(1);
- PrevVecV[0] = Op0;
- PrevVecV[1] = Op1;
- break;
- }
- default:
- return false;
- }
- ShouldBeCallOrBinInst ^= 1;
-
- IntrinsicCostAttributes ICA(
- *CommonCallOp, II->getType(),
- {PrevVecV[0]->getType(), PrevVecV[1]->getType()});
- OrigCost += TTI.getIntrinsicInstrCost(ICA, CostKind);
-
- // We may need a swap here since it can be (a, b) or (b, a)
- // and accordingly change as we go up.
- if (!isa<ShuffleVectorInst>(PrevVecV[1]))
- std::swap(PrevVecV[0], PrevVecV[1]);
- InstWorklist.push(PrevVecV[1]);
- InstWorklist.push(PrevVecV[0]);
- } else if (auto *BinOp = dyn_cast<BinaryOperator>(CI)) {
- // Similar logic for bin ops.
+ auto IsChainNode = [&](Value *V) {
+ if (auto *BO = dyn_cast<BinaryOperator>(V))
+ return CommonBinOp && BO->getOpcode() == *CommonBinOp;
+ if (auto *MMI = dyn_cast<MinMaxIntrinsic>(V))
+ return CommonCallOp && MMI->getIntrinsicID() == *CommonCallOp;
+ if (auto *SVI = dyn_cast<ShuffleVectorInst>(V))
+ return isa<PoisonValue>(SVI->getOperand(1));
+ return false;
+ };
- if (!ShouldBeCallOrBinInst)
- return false;
+ // Walk VecOpEE bottom-up. Chain nodes are single-source shuffles and
+ // matching-opcode binops/intrinsics. Anything else is a leaf source.
+ constexpr unsigned MaxChainNodes = 32;
+ SmallVector<Value *, 16> ChainPostorder;
+ SmallPtrSet<Value *, 16> Visited;
+ SmallVector<Value *, 2> Sources;
+ DenseMap<Value *, unsigned> SrcSizes;
- if (!IsFirstCallOrBinInst && any_of(PrevVecV, equal_to(nullptr)))
- return false;
+ struct StackEntry {
+ Value *V;
+ unsigned ChildIdx;
+ };
+ SmallVector<StackEntry, 16> Stack;
- if (BinOp != (IsFirstCallOrBinInst ? VecOpEE : PrevVecV[0]))
+ auto Enqueue = [&](Value *V) -> bool {
+ if (Visited.size() >= MaxChainNodes)
+ return false;
+ if (!Visited.insert(V).second)
+ return true;
+ if (!IsChainNode(V)) {
+ auto *VT = dyn_cast<FixedVectorType>(V->getType());
+ if (!VT)
return false;
- IsFirstCallOrBinInst = false;
+ if (SrcSizes.insert({V, VT->getNumElements()}).second)
+ Sources.push_back(V);
+ return true;
+ }
+ Stack.push_back({V, 0});
+ return true;
+ };
- if (!CommonBinOp)
- CommonBinOp = BinOp->getOpcode();
+ // VecOpEE is a chain node so push directly.
+ Visited.insert(VecOpEE);
+ Stack.push_back({VecOpEE, 0});
- if (BinOp->getOpcode() != *CommonBinOp)
- return false;
+ while (!Stack.empty()) {
+ auto &Top = Stack.back();
+ Value *V = Top.V;
- switch (*CommonBinOp) {
- case BinaryOperator::Add:
- case BinaryOperator::Mul:
- case BinaryOperator::Or:
- case BinaryOperator::And:
- case BinaryOperator::Xor: {
- auto *Op0 = BinOp->getOperand(0);
- auto *Op1 = BinOp->getOperand(1);
- PrevVecV[0] = Op0;
- PrevVecV[1] = Op1;
- break;
- }
- default:
+ // Chain shuffles always have poison as op1, so only op0 matters.
+ unsigned NumOps = isa<ShuffleVectorInst>(V) ? 1 : 2;
+ if (Top.ChildIdx < NumOps) {
+ Value *Child = cast<User>(V)->getOperand(Top.ChildIdx++);
+ if (!Enqueue(Child))
return false;
- }
- ShouldBeCallOrBinInst ^= 1;
+ } else {
+ ChainPostorder.push_back(V);
+ Stack.pop_back();
+ }
+ }
- OrigCost +=
- TTI.getArithmeticInstrCost(*CommonBinOp, BinOp->getType(), CostKind);
-
- if (!isa<ShuffleVectorInst>(PrevVecV[1]))
- std::swap(PrevVecV[0], PrevVecV[1]);
- InstWorklist.push(PrevVecV[1]);
- InstWorklist.push(PrevVecV[0]);
- } else if (auto *SVInst = dyn_cast<ShuffleVectorInst>(CI)) {
- // We shouldn't have any null values in the previous vectors,
- // is so, there was a mismatch in pattern.
- if (ShouldBeCallOrBinInst || any_of(PrevVecV, equal_to(nullptr)))
- return false;
+ // Demote any binop/intrinsic whose operands are not themselves chain nodes.
+ // Walking past would re-derive the value from its operands, leaving the
+ // original alive if anything downstream still uses it.
+ SmallPtrSet<Value *, 16> KeptChain;
+ SmallVector<Value *, 16> NewChainPostorder;
+ NewChainPostorder.reserve(ChainPostorder.size());
+ for (Value *V : ChainPostorder) {
+ if (isa<ShuffleVectorInst>(V)) {
+ KeptChain.insert(V);
+ NewChainPostorder.push_back(V);
+ continue;
+ }
+ auto *U = cast<User>(V);
+ if (KeptChain.contains(U->getOperand(0)) ||
+ KeptChain.contains(U->getOperand(1))) {
+ KeptChain.insert(V);
+ NewChainPostorder.push_back(V);
+ continue;
+ }
+ auto *VT = cast<FixedVectorType>(V->getType());
+ if (SrcSizes.insert({V, VT->getNumElements()}).second)
+ Sources.push_back(V);
+ }
+ ChainPostorder = std::move(NewChainPostorder);
- if (SVInst != PrevVecV[1])
- return false;
+ bool IsIdempotent =
+ CommonCallOp || (CommonBinOp && Instruction::isIdempotent(*CommonBinOp));
- ArrayRef<int> CurMask;
- if (!match(SVInst, m_Shuffle(m_Specific(PrevVecV[0]), m_Poison(),
- m_Mask(CurMask))))
- return false;
+ // Each output lane has a bitmask of contributing source lanes and a
+ // poison flag. Shuffles permute these records. Binops union them.
+ struct LaneInfo {
+ SmallDenseMap<Value *, APInt, 2> SrcBits;
----------------
as4230 wrote:
I renamed LaneInfo::SrcBits to SrcElts and Partial::Bits to Elts for consistency. They store per lane masks so the convention matches DemandedElts rather than KnownBits anyways.
https://github.com/llvm/llvm-project/pull/199872
More information about the llvm-commits
mailing list