1//===- ExpandReductions.cpp - Expand reduction intrinsics -----------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This pass implements IR expansion for reduction intrinsics, allowing targets
10// to enable the intrinsics until just before codegen.
11//
12//===----------------------------------------------------------------------===//
13
14#include "llvm/CodeGen/ExpandReductions.h"
15#include "llvm/Analysis/LoopInfo.h"
16#include "llvm/Analysis/TargetTransformInfo.h"
17#include "llvm/CodeGen/Passes.h"
18#include "llvm/IR/Dominators.h"
19#include "llvm/IR/IRBuilder.h"
20#include "llvm/IR/InstIterator.h"
21#include "llvm/IR/IntrinsicInst.h"
22#include "llvm/IR/Intrinsics.h"
23#include "llvm/InitializePasses.h"
24#include "llvm/Pass.h"
25#include "llvm/Transforms/Utils/LoopUtils.h"
26
27using namespace llvm;
28
29namespace {
30
31bool expandReductions(Function &F, const TargetTransformInfo *TTI,
32 DominatorTree *DT, LoopInfo *LI) {
33 bool Changed = false;
34 SmallVector<IntrinsicInst *, 4> Worklist;
35 for (auto &I : instructions(F)) {
36 if (auto *II = dyn_cast<IntrinsicInst>(Val: &I)) {
37 switch (II->getIntrinsicID()) {
38 default:
39 break;
40 case Intrinsic::vector_reduce_fadd:
41 case Intrinsic::vector_reduce_fmul:
42 case Intrinsic::vector_reduce_add:
43 case Intrinsic::vector_reduce_mul:
44 case Intrinsic::vector_reduce_and:
45 case Intrinsic::vector_reduce_or:
46 case Intrinsic::vector_reduce_xor:
47 case Intrinsic::vector_reduce_smax:
48 case Intrinsic::vector_reduce_smin:
49 case Intrinsic::vector_reduce_umax:
50 case Intrinsic::vector_reduce_umin:
51 case Intrinsic::vector_reduce_fmax:
52 case Intrinsic::vector_reduce_fmin:
53 case Intrinsic::vector_reduce_fmaximum:
54 case Intrinsic::vector_reduce_fminimum:
55 case Intrinsic::vector_reduce_fmaximumnum:
56 case Intrinsic::vector_reduce_fminimumnum: {
57 // Only expand if the target doesn't support this operation natively.
58 if (TTI->shouldExpandReduction(II))
59 Worklist.push_back(Elt: II);
60 break;
61 }
62 }
63 }
64 }
65
66 for (auto *II : Worklist) {
67 FastMathFlags FMF = II->getFastMathFlagsOrNone();
68 Intrinsic::ID ID = II->getIntrinsicID();
69 RecurKind RK = getMinMaxReductionRecurKind(RdxID: ID);
70 TargetTransformInfo::ReductionShuffle RS =
71 TTI->getPreferredExpandedReductionShuffle(II);
72
73 Value *Rdx = nullptr;
74 IRBuilder<> Builder(II);
75 IRBuilder<>::FastMathFlagGuard FMFGuard(Builder);
76 Builder.setFastMathFlags(FMF);
77 switch (ID) {
78 default:
79 llvm_unreachable("Unexpected intrinsic!");
80 case Intrinsic::vector_reduce_fadd:
81 case Intrinsic::vector_reduce_fmul: {
82 // FMFs must be attached to the call, otherwise it's an ordered reduction
83 // and it can't be handled by generating a shuffle sequence.
84 Value *Acc = II->getArgOperand(i: 0);
85 Value *Vec = II->getArgOperand(i: 1);
86 unsigned RdxOpcode = getArithmeticReductionInstruction(RdxID: ID);
87 if (isa<ScalableVectorType>(Val: Vec->getType())) {
88 Rdx = expandReductionViaLoop(Builder, Vec, RdxOpcode, Acc, DT, LI);
89 break;
90 }
91 if (!FMF.allowReassoc())
92 Rdx = getOrderedReduction(Builder, Acc, Src: Vec, Op: RdxOpcode, MinMaxKind: RK);
93 else {
94 if (!isPowerOf2_32(
95 Value: cast<FixedVectorType>(Val: Vec->getType())->getNumElements()))
96 continue;
97 Rdx = getShuffleReduction(Builder, Src: Vec, Op: RdxOpcode, RS, MinMaxKind: RK);
98 Rdx = Builder.CreateBinOp(Opc: (Instruction::BinaryOps)RdxOpcode, LHS: Acc, RHS: Rdx,
99 Name: "bin.rdx");
100 }
101 break;
102 }
103 case Intrinsic::vector_reduce_and:
104 case Intrinsic::vector_reduce_or: {
105 // Canonicalize logical or/and reductions:
106 // Or reduction for i1 is represented as:
107 // %val = bitcast <ReduxWidth x i1> to iReduxWidth
108 // %res = cmp ne iReduxWidth %val, 0
109 // And reduction for i1 is represented as:
110 // %val = bitcast <ReduxWidth x i1> to iReduxWidth
111 // %res = cmp eq iReduxWidth %val, 11111
112 Value *Vec = II->getArgOperand(i: 0);
113 auto *FTy = cast<FixedVectorType>(Val: Vec->getType());
114 unsigned NumElts = FTy->getNumElements();
115 if (!isPowerOf2_32(Value: NumElts))
116 continue;
117
118 if (FTy->getElementType() == Builder.getInt1Ty()) {
119 Rdx = Builder.CreateBitCast(V: Vec, DestTy: Builder.getIntNTy(N: NumElts));
120 if (ID == Intrinsic::vector_reduce_and) {
121 Rdx = Builder.CreateICmpEQ(
122 LHS: Rdx, RHS: ConstantInt::getAllOnesValue(Ty: Rdx->getType()));
123 } else {
124 assert(ID == Intrinsic::vector_reduce_or && "Expected or reduction.");
125 Rdx = Builder.CreateIsNotNull(Arg: Rdx);
126 }
127 break;
128 }
129 unsigned RdxOpcode = getArithmeticReductionInstruction(RdxID: ID);
130 Rdx = getShuffleReduction(Builder, Src: Vec, Op: RdxOpcode, RS, MinMaxKind: RK);
131 break;
132 }
133 case Intrinsic::vector_reduce_add:
134 case Intrinsic::vector_reduce_mul:
135 case Intrinsic::vector_reduce_xor:
136 case Intrinsic::vector_reduce_smax:
137 case Intrinsic::vector_reduce_smin:
138 case Intrinsic::vector_reduce_umax:
139 case Intrinsic::vector_reduce_umin: {
140 Value *Vec = II->getArgOperand(i: 0);
141 unsigned RdxOpcode = getArithmeticReductionInstruction(RdxID: ID);
142 if (isa<ScalableVectorType>(Val: Vec->getType())) {
143 Type *EltTy = Vec->getType()->getScalarType();
144 Value *Ident = getReductionIdentity(RdxID: ID, Ty: EltTy, FMF);
145 Rdx = expandReductionViaLoop(Builder, Vec, RdxOpcode, Acc: Ident, DT, LI);
146 break;
147 }
148 if (!isPowerOf2_32(
149 Value: cast<FixedVectorType>(Val: Vec->getType())->getNumElements()))
150 continue;
151 Rdx = getShuffleReduction(Builder, Src: Vec, Op: RdxOpcode, RS, MinMaxKind: RK);
152 break;
153 }
154 case Intrinsic::vector_reduce_fmax:
155 case Intrinsic::vector_reduce_fmin: {
156 // We require "nnan" to use a shuffle reduction; "nsz" is implied by the
157 // semantics of the reduction.
158 Value *Vec = II->getArgOperand(i: 0);
159 if (!isPowerOf2_32(
160 Value: cast<FixedVectorType>(Val: Vec->getType())->getNumElements()) ||
161 !FMF.noNaNs())
162 continue;
163 unsigned RdxOpcode = getArithmeticReductionInstruction(RdxID: ID);
164 Rdx = getShuffleReduction(Builder, Src: Vec, Op: RdxOpcode, RS, MinMaxKind: RK);
165 break;
166 }
167 case Intrinsic::vector_reduce_fmaximum:
168 case Intrinsic::vector_reduce_fminimum:
169 case Intrinsic::vector_reduce_fmaximumnum:
170 case Intrinsic::vector_reduce_fminimumnum: {
171 Value *Vec = II->getArgOperand(i: 0);
172 if (!isPowerOf2_32(
173 Value: cast<FixedVectorType>(Val: Vec->getType())->getNumElements()))
174 continue;
175 unsigned RdxOpcode = getArithmeticReductionInstruction(RdxID: ID);
176 Rdx = getShuffleReduction(Builder, Src: Vec, Op: RdxOpcode, RS, MinMaxKind: RK);
177 break;
178 }
179 }
180 II->replaceAllUsesWith(V: Rdx);
181 II->eraseFromParent();
182 Changed = true;
183 }
184 return Changed;
185}
186
187class ExpandReductions : public FunctionPass {
188public:
189 static char ID;
190 ExpandReductions() : FunctionPass(ID) {}
191
192 bool runOnFunction(Function &F) override {
193 const auto *TTI =&getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
194 auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>();
195 auto *LIWP = getAnalysisIfAvailable<LoopInfoWrapperPass>();
196 auto *DT = DTWP ? &DTWP->getDomTree() : nullptr;
197 auto *LI = LIWP ? &LIWP->getLoopInfo() : nullptr;
198 return expandReductions(F, TTI, DT, LI);
199 }
200
201 void getAnalysisUsage(AnalysisUsage &AU) const override {
202 AU.addRequired<TargetTransformInfoWrapperPass>();
203 AU.addPreserved<DominatorTreeWrapperPass>();
204 AU.addPreserved<LoopInfoWrapperPass>();
205 }
206};
207}
208
209char ExpandReductions::ID;
210INITIALIZE_PASS_BEGIN(ExpandReductions, "expand-reductions",
211 "Expand reduction intrinsics", false, false)
212INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
213INITIALIZE_PASS_END(ExpandReductions, "expand-reductions",
214 "Expand reduction intrinsics", false, false)
215
216FunctionPass *llvm::createExpandReductionsPass() {
217 return new ExpandReductions();
218}
219
220PreservedAnalyses ExpandReductionsPass::run(Function &F,
221 FunctionAnalysisManager &AM) {
222 const auto &TTI = AM.getResult<TargetIRAnalysis>(IR&: F);
223 auto *DT = AM.getCachedResult<DominatorTreeAnalysis>(IR&: F);
224 auto *LI = AM.getCachedResult<LoopAnalysis>(IR&: F);
225 if (!expandReductions(F, TTI: &TTI, DT, LI))
226 return PreservedAnalyses::all();
227 PreservedAnalyses PA;
228 PA.preserve<DominatorTreeAnalysis>();
229 PA.preserve<LoopAnalysis>();
230 return PA;
231}
232