[llvm] [VectorCombine] Add subvector reduction support to foldShuffleChainsToReduce (PR #199872)
Simon Pilgrim via llvm-commits
llvm-commits at lists.llvm.org
Fri Jun 12 07:03:01 PDT 2026
================
@@ -3999,265 +3953,365 @@ 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)) {
+ if (!getReductionForBinop(BO->getOpcode()))
+ return false;
+ CommonBinOp = BO->getOpcode();
+ } else if (auto *MMI = dyn_cast<MinMaxIntrinsic>(VecOpEE)) {
+ CommonCallOp = MMI->getIntrinsicID();
+ } else {
+ return false;
}
- InstWorklist.push(VecOpEE);
-
- bool IsPartialReduction = false;
- bool HasLaneDuplication = 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 floating-point reductions, track FMF intersection across all binops.
+ FastMathFlags CommonFMF;
+ bool IsFloatReduction = 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;
+ 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 (!CommonCallOp)
- CommonCallOp = II->getIntrinsicID();
- if (II->getIntrinsicID() != *CommonCallOp)
- 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;
+ SmallPtrSet<Value *, 16> KeptChain;
+ SmallVector<Value *, 2> Sources;
+ DenseMap<Value *, unsigned> SrcSizes;
+
+ auto AddSource = [&](Value *V) -> bool {
+ auto *VT = dyn_cast<FixedVectorType>(V->getType());
+ if (!VT)
+ return false;
+ if (SrcSizes.insert({V, VT->getNumElements()}).second)
+ Sources.push_back(V);
+ return true;
+ };
- 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.
-
- if (!ShouldBeCallOrBinInst)
- return false;
+ struct StackEntry {
+ Value *V;
+ unsigned ChildIdx;
+ };
+ SmallVector<StackEntry, 16> Stack;
- if (!IsFirstCallOrBinInst && any_of(PrevVecV, equal_to(nullptr)))
- return false;
+ auto Enqueue = [&](Value *V) -> bool {
+ if (Visited.size() >= MaxChainNodes)
+ return false;
+ if (!Visited.insert(V).second)
+ return true;
+ if (!IsChainNode(V))
+ return AddSource(V);
+ Stack.push_back({V, 0});
+ return true;
+ };
- if (BinOp != (IsFirstCallOrBinInst ? VecOpEE : PrevVecV[0]))
- return false;
- IsFirstCallOrBinInst = false;
+ // VecOpEE is a chain node so push directly.
+ Visited.insert(VecOpEE);
+ Stack.push_back({VecOpEE, 0});
- if (!CommonBinOp)
- CommonBinOp = BinOp->getOpcode();
+ while (!Stack.empty()) {
+ auto &Top = Stack.back();
+ Value *V = Top.V;
- if (BinOp->getOpcode() != *CommonBinOp)
+ // 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;
-
- switch (*CommonBinOp) {
- case BinaryOperator::Add:
- case BinaryOperator::Mul:
- case BinaryOperator::Or:
- case BinaryOperator::And:
- case BinaryOperator::Xor:
- case BinaryOperator::FAdd:
- case BinaryOperator::FMul: {
- auto *Op0 = BinOp->getOperand(0);
- auto *Op1 = BinOp->getOperand(1);
- PrevVecV[0] = Op0;
- PrevVecV[1] = Op1;
- break;
+ } else {
+ // 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.
+ bool Keep = isa<ShuffleVectorInst>(V);
+ if (!Keep) {
+ auto *U = cast<User>(V);
+ Keep = KeptChain.contains(U->getOperand(0)) ||
+ KeptChain.contains(U->getOperand(1));
}
- default:
+ if (Keep) {
+ KeptChain.insert(V);
+ ChainPostorder.push_back(V);
+ } else if (!AddSource(V)) {
return false;
}
+ Stack.pop_back();
+ }
+ }
- // For FP reductions, require reassoc on every binop and collect FMF.
- if (*CommonBinOp == Instruction::FAdd ||
- *CommonBinOp == Instruction::FMul) {
- if (!BinOp->hasAllowReassoc())
- return false;
- if (!IsFloatReduction) {
- CommonFMF = BinOp->getFastMathFlags();
- IsFloatReduction = true;
- } else {
- CommonFMF &= BinOp->getFastMathFlags();
- }
- }
-
- ShouldBeCallOrBinInst ^= 1;
+ // No chain nodes (e.g. extract(binop(x, y), 0)). Rebuilding would just
+ // recreate the original extract and refold forever when extraction is free.
+ if (ChainPostorder.empty())
+ return false;
- OrigCost +=
- TTI.getArithmeticInstrCost(*CommonBinOp, BinOp->getType(), CostKind);
+ bool IsIdempotent =
+ CommonCallOp || (CommonBinOp && Instruction::isIdempotent(*CommonBinOp));
- 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;
+ // For FP reductions, require reassoc on every binop and collect FMF.
+ for (Value *V : ChainPostorder) {
+ auto *BinOp = dyn_cast<BinaryOperator>(V);
+ if (!BinOp || !BinOp->getType()->isFPOrFPVectorTy())
+ continue;
+ if (!BinOp->hasAllowReassoc())
+ return false;
+ if (!IsFloatReduction) {
+ CommonFMF = BinOp->getFastMathFlags();
+ IsFloatReduction = true;
+ } else {
+ CommonFMF &= BinOp->getFastMathFlags();
+ }
+ }
- if (SVInst != PrevVecV[1])
- return false;
+ // Top-down demanded elements. How many times does each lane of each chain
+ // value feed the extracted lane 0? Reverse postorder visits every use before
+ // its value. A binop forwards its demand to both operands and a shuffle
+ // follows its mask back to the source lane.
+ DenseMap<Value *, SmallVector<unsigned, 16>> Demand;
----------------
RKSimon wrote:
Can this be done simpler with APInt masks instead? Or BitVector if you have to.
https://github.com/llvm/llvm-project/pull/199872
More information about the llvm-commits
mailing list