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 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 
31 unsigned getOpcode(Intrinsic::ID ID) {
32   switch (ID) {
33   case Intrinsic::vector_reduce_fadd:
34     return Instruction::FAdd;
35   case Intrinsic::vector_reduce_fmul:
36     return Instruction::FMul;
37   case Intrinsic::vector_reduce_add:
38     return Instruction::Add;
39   case Intrinsic::vector_reduce_mul:
40     return Instruction::Mul;
41   case Intrinsic::vector_reduce_and:
42     return Instruction::And;
43   case Intrinsic::vector_reduce_or:
44     return Instruction::Or;
45   case Intrinsic::vector_reduce_xor:
46     return Instruction::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     return Instruction::ICmp;
52   case Intrinsic::vector_reduce_fmax:
53   case Intrinsic::vector_reduce_fmin:
54     return Instruction::FCmp;
55   default:
56     llvm_unreachable("Unexpected ID");
57   }
58 }
59 
60 RecurKind getRK(Intrinsic::ID ID) {
61   switch (ID) {
62   case Intrinsic::vector_reduce_smax:
63     return RecurKind::SMax;
64   case Intrinsic::vector_reduce_smin:
65     return RecurKind::SMin;
66   case Intrinsic::vector_reduce_umax:
67     return RecurKind::UMax;
68   case Intrinsic::vector_reduce_umin:
69     return RecurKind::UMin;
70   case Intrinsic::vector_reduce_fmax:
71     return RecurKind::FMax;
72   case Intrinsic::vector_reduce_fmin:
73     return RecurKind::FMin;
74   default:
75     return RecurKind::None;
76   }
77 }
78 
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::vector_reduce_fadd:
87       case Intrinsic::vector_reduce_fmul:
88       case Intrinsic::vector_reduce_add:
89       case Intrinsic::vector_reduce_mul:
90       case Intrinsic::vector_reduce_and:
91       case Intrinsic::vector_reduce_or:
92       case Intrinsic::vector_reduce_xor:
93       case Intrinsic::vector_reduce_smax:
94       case Intrinsic::vector_reduce_smin:
95       case Intrinsic::vector_reduce_umax:
96       case Intrinsic::vector_reduce_umin:
97       case Intrinsic::vector_reduce_fmax:
98       case Intrinsic::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     RecurKind RK = getRK(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::vector_reduce_fadd:
120     case Intrinsic::vector_reduce_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), RK);
127       else {
128         if (!isPowerOf2_32(
129                 cast<FixedVectorType>(Vec->getType())->getNumElements()))
130           continue;
131 
132         Rdx = getShuffleReduction(Builder, Vec, getOpcode(ID), RK);
133         Rdx = Builder.CreateBinOp((Instruction::BinaryOps)getOpcode(ID),
134                                   Acc, Rdx, "bin.rdx");
135       }
136       break;
137     }
138     case Intrinsic::vector_reduce_add:
139     case Intrinsic::vector_reduce_mul:
140     case Intrinsic::vector_reduce_and:
141     case Intrinsic::vector_reduce_or:
142     case Intrinsic::vector_reduce_xor:
143     case Intrinsic::vector_reduce_smax:
144     case Intrinsic::vector_reduce_smin:
145     case Intrinsic::vector_reduce_umax:
146     case Intrinsic::vector_reduce_umin: {
147       Value *Vec = II->getArgOperand(0);
148       if (!isPowerOf2_32(
149               cast<FixedVectorType>(Vec->getType())->getNumElements()))
150         continue;
151 
152       Rdx = getShuffleReduction(Builder, Vec, getOpcode(ID), RK);
153       break;
154     }
155     case Intrinsic::vector_reduce_fmax:
156     case Intrinsic::vector_reduce_fmin: {
157       // We require "nnan" to use a shuffle reduction; "nsz" is implied by the
158       // semantics of the reduction.
159       Value *Vec = II->getArgOperand(0);
160       if (!isPowerOf2_32(
161               cast<FixedVectorType>(Vec->getType())->getNumElements()) ||
162           !FMF.noNaNs())
163         continue;
164 
165       Rdx = getShuffleReduction(Builder, Vec, getOpcode(ID), RK);
166       break;
167     }
168     }
169     II->replaceAllUsesWith(Rdx);
170     II->eraseFromParent();
171     Changed = true;
172   }
173   return Changed;
174 }
175 
176 class ExpandReductions : public FunctionPass {
177 public:
178   static char ID;
179   ExpandReductions() : FunctionPass(ID) {
180     initializeExpandReductionsPass(*PassRegistry::getPassRegistry());
181   }
182 
183   bool runOnFunction(Function &F) override {
184     const auto *TTI =&getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
185     return expandReductions(F, TTI);
186   }
187 
188   void getAnalysisUsage(AnalysisUsage &AU) const override {
189     AU.addRequired<TargetTransformInfoWrapperPass>();
190     AU.setPreservesCFG();
191   }
192 };
193 }
194 
195 char ExpandReductions::ID;
196 INITIALIZE_PASS_BEGIN(ExpandReductions, "expand-reductions",
197                       "Expand reduction intrinsics", false, false)
198 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
199 INITIALIZE_PASS_END(ExpandReductions, "expand-reductions",
200                     "Expand reduction intrinsics", false, false)
201 
202 FunctionPass *llvm::createExpandReductionsPass() {
203   return new ExpandReductions();
204 }
205 
206 PreservedAnalyses ExpandReductionsPass::run(Function &F,
207                                             FunctionAnalysisManager &AM) {
208   const auto &TTI = AM.getResult<TargetIRAnalysis>(F);
209   if (!expandReductions(F, &TTI))
210     return PreservedAnalyses::all();
211   PreservedAnalyses PA;
212   PA.preserveSet<CFGAnalyses>();
213   return PA;
214 }
215