1 //===--- ExpandReductions.cpp - Expand experimental 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 experimental intrinsics until just before codegen.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #include "llvm/CodeGen/ExpandReductions.h"
15 #include "llvm/Analysis/TargetTransformInfo.h"
16 #include "llvm/CodeGen/Passes.h"
17 #include "llvm/IR/Function.h"
18 #include "llvm/IR/IRBuilder.h"
19 #include "llvm/IR/InstIterator.h"
20 #include "llvm/IR/IntrinsicInst.h"
21 #include "llvm/IR/Intrinsics.h"
22 #include "llvm/IR/Module.h"
23 #include "llvm/InitializePasses.h"
24 #include "llvm/Pass.h"
25 #include "llvm/Transforms/Utils/LoopUtils.h"
26
27 using namespace llvm;
28
29 namespace {
30
getOpcode(Intrinsic::ID ID)31 unsigned getOpcode(Intrinsic::ID ID) {
32 switch (ID) {
33 case Intrinsic::experimental_vector_reduce_v2_fadd:
34 return Instruction::FAdd;
35 case Intrinsic::experimental_vector_reduce_v2_fmul:
36 return Instruction::FMul;
37 case Intrinsic::experimental_vector_reduce_add:
38 return Instruction::Add;
39 case Intrinsic::experimental_vector_reduce_mul:
40 return Instruction::Mul;
41 case Intrinsic::experimental_vector_reduce_and:
42 return Instruction::And;
43 case Intrinsic::experimental_vector_reduce_or:
44 return Instruction::Or;
45 case Intrinsic::experimental_vector_reduce_xor:
46 return Instruction::Xor;
47 case Intrinsic::experimental_vector_reduce_smax:
48 case Intrinsic::experimental_vector_reduce_smin:
49 case Intrinsic::experimental_vector_reduce_umax:
50 case Intrinsic::experimental_vector_reduce_umin:
51 return Instruction::ICmp;
52 case Intrinsic::experimental_vector_reduce_fmax:
53 case Intrinsic::experimental_vector_reduce_fmin:
54 return Instruction::FCmp;
55 default:
56 llvm_unreachable("Unexpected ID");
57 }
58 }
59
getMRK(Intrinsic::ID ID)60 RecurrenceDescriptor::MinMaxRecurrenceKind getMRK(Intrinsic::ID ID) {
61 switch (ID) {
62 case Intrinsic::experimental_vector_reduce_smax:
63 return RecurrenceDescriptor::MRK_SIntMax;
64 case Intrinsic::experimental_vector_reduce_smin:
65 return RecurrenceDescriptor::MRK_SIntMin;
66 case Intrinsic::experimental_vector_reduce_umax:
67 return RecurrenceDescriptor::MRK_UIntMax;
68 case Intrinsic::experimental_vector_reduce_umin:
69 return RecurrenceDescriptor::MRK_UIntMin;
70 case Intrinsic::experimental_vector_reduce_fmax:
71 return RecurrenceDescriptor::MRK_FloatMax;
72 case Intrinsic::experimental_vector_reduce_fmin:
73 return RecurrenceDescriptor::MRK_FloatMin;
74 default:
75 return RecurrenceDescriptor::MRK_Invalid;
76 }
77 }
78
expandReductions(Function & F,const TargetTransformInfo * TTI)79 bool expandReductions(Function &F, const TargetTransformInfo *TTI) {
80 bool Changed = false;
81 SmallVector<IntrinsicInst *, 4> Worklist;
82 for (auto &I : instructions(F)) {
83 if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
84 switch (II->getIntrinsicID()) {
85 default: break;
86 case Intrinsic::experimental_vector_reduce_v2_fadd:
87 case Intrinsic::experimental_vector_reduce_v2_fmul:
88 case Intrinsic::experimental_vector_reduce_add:
89 case Intrinsic::experimental_vector_reduce_mul:
90 case Intrinsic::experimental_vector_reduce_and:
91 case Intrinsic::experimental_vector_reduce_or:
92 case Intrinsic::experimental_vector_reduce_xor:
93 case Intrinsic::experimental_vector_reduce_smax:
94 case Intrinsic::experimental_vector_reduce_smin:
95 case Intrinsic::experimental_vector_reduce_umax:
96 case Intrinsic::experimental_vector_reduce_umin:
97 case Intrinsic::experimental_vector_reduce_fmax:
98 case Intrinsic::experimental_vector_reduce_fmin:
99 if (TTI->shouldExpandReduction(II))
100 Worklist.push_back(II);
101
102 break;
103 }
104 }
105 }
106
107 for (auto *II : Worklist) {
108 FastMathFlags FMF =
109 isa<FPMathOperator>(II) ? II->getFastMathFlags() : FastMathFlags{};
110 Intrinsic::ID ID = II->getIntrinsicID();
111 RecurrenceDescriptor::MinMaxRecurrenceKind MRK = getMRK(ID);
112
113 Value *Rdx = nullptr;
114 IRBuilder<> Builder(II);
115 IRBuilder<>::FastMathFlagGuard FMFGuard(Builder);
116 Builder.setFastMathFlags(FMF);
117 switch (ID) {
118 default: llvm_unreachable("Unexpected intrinsic!");
119 case Intrinsic::experimental_vector_reduce_v2_fadd:
120 case Intrinsic::experimental_vector_reduce_v2_fmul: {
121 // FMFs must be attached to the call, otherwise it's an ordered reduction
122 // and it can't be handled by generating a shuffle sequence.
123 Value *Acc = II->getArgOperand(0);
124 Value *Vec = II->getArgOperand(1);
125 if (!FMF.allowReassoc())
126 Rdx = getOrderedReduction(Builder, Acc, Vec, getOpcode(ID), MRK);
127 else {
128 if (!isPowerOf2_32(Vec->getType()->getVectorNumElements()))
129 continue;
130
131 Rdx = getShuffleReduction(Builder, Vec, getOpcode(ID), MRK);
132 Rdx = Builder.CreateBinOp((Instruction::BinaryOps)getOpcode(ID),
133 Acc, Rdx, "bin.rdx");
134 }
135 break;
136 }
137 case Intrinsic::experimental_vector_reduce_add:
138 case Intrinsic::experimental_vector_reduce_mul:
139 case Intrinsic::experimental_vector_reduce_and:
140 case Intrinsic::experimental_vector_reduce_or:
141 case Intrinsic::experimental_vector_reduce_xor:
142 case Intrinsic::experimental_vector_reduce_smax:
143 case Intrinsic::experimental_vector_reduce_smin:
144 case Intrinsic::experimental_vector_reduce_umax:
145 case Intrinsic::experimental_vector_reduce_umin:
146 case Intrinsic::experimental_vector_reduce_fmax:
147 case Intrinsic::experimental_vector_reduce_fmin: {
148 Value *Vec = II->getArgOperand(0);
149 if (!isPowerOf2_32(Vec->getType()->getVectorNumElements()))
150 continue;
151
152 Rdx = getShuffleReduction(Builder, Vec, getOpcode(ID), MRK);
153 break;
154 }
155 }
156 II->replaceAllUsesWith(Rdx);
157 II->eraseFromParent();
158 Changed = true;
159 }
160 return Changed;
161 }
162
163 class ExpandReductions : public FunctionPass {
164 public:
165 static char ID;
ExpandReductions()166 ExpandReductions() : FunctionPass(ID) {
167 initializeExpandReductionsPass(*PassRegistry::getPassRegistry());
168 }
169
runOnFunction(Function & F)170 bool runOnFunction(Function &F) override {
171 const auto *TTI =&getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
172 return expandReductions(F, TTI);
173 }
174
getAnalysisUsage(AnalysisUsage & AU) const175 void getAnalysisUsage(AnalysisUsage &AU) const override {
176 AU.addRequired<TargetTransformInfoWrapperPass>();
177 AU.setPreservesCFG();
178 }
179 };
180 }
181
182 char ExpandReductions::ID;
183 INITIALIZE_PASS_BEGIN(ExpandReductions, "expand-reductions",
184 "Expand reduction intrinsics", false, false)
INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)185 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
186 INITIALIZE_PASS_END(ExpandReductions, "expand-reductions",
187 "Expand reduction intrinsics", false, false)
188
189 FunctionPass *llvm::createExpandReductionsPass() {
190 return new ExpandReductions();
191 }
192
run(Function & F,FunctionAnalysisManager & AM)193 PreservedAnalyses ExpandReductionsPass::run(Function &F,
194 FunctionAnalysisManager &AM) {
195 const auto &TTI = AM.getResult<TargetIRAnalysis>(F);
196 if (!expandReductions(F, &TTI))
197 return PreservedAnalyses::all();
198 PreservedAnalyses PA;
199 PA.preserveSet<CFGAnalyses>();
200 return PA;
201 }
202