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
27using namespace llvm;
28
29namespace {
30
31unsigned 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
60RecurKind 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
79bool 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
176class ExpandReductions : public FunctionPass {
177public:
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
195char ExpandReductions::ID;
196INITIALIZE_PASS_BEGIN(ExpandReductions, "expand-reductions",
197                      "Expand reduction intrinsics", false, false)
198INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
199INITIALIZE_PASS_END(ExpandReductions, "expand-reductions",
200                    "Expand reduction intrinsics", false, false)
201
202FunctionPass *llvm::createExpandReductionsPass() {
203  return new ExpandReductions();
204}
205
206PreservedAnalyses 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