[llvm] [VectorCombine] Add subvector reduction support to foldShuffleChainsToReduce (PR #199872)
Adam Scott via llvm-commits
llvm-commits at lists.llvm.org
Sun Jun 14 13:08:49 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;
----------------
as4230 wrote:
I think a lot of complexity comes from one extra case. The rewrite also handled results that draw lanes from several source vectors at once which needs more than one reduce combined with scalar ops. That path was where a good chunk of the state lived (Partial, Partials, SourcesClean, FullInter) and it buried the flow so I've dropped it and can add it in the future.
With that gone the fold reduces a single value (either a source or a covering intermediate) and the rest consolidates. The lane tracking is now a small Demand struct of two APInt masks (demanded and duplicated lanes) instead of the count vectors.
The function now reads top to bottom:
match the extract
collect the chain nodes and sources
compute demanded lanes
pick one value to reduce
cost the reduce
emit the reduce
I've pushed this so hopefully its easier to follow now.
https://github.com/llvm/llvm-project/pull/199872
More information about the llvm-commits
mailing list