1 //===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//
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 file implements utilities for the Linalg dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/Linalg/Utils/Utils.h"
14 
15 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
16 #include "mlir/Dialect/Affine/IR/AffineOps.h"
17 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
18 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
19 #include "mlir/Dialect/SCF/EDSC/Builders.h"
20 #include "mlir/Dialect/SCF/SCF.h"
21 #include "mlir/Dialect/StandardOps/IR/Ops.h"
22 #include "mlir/IR/AffineExpr.h"
23 #include "mlir/IR/AffineMap.h"
24 #include "mlir/IR/Matchers.h"
25 #include "mlir/IR/OpImplementation.h"
26 #include "mlir/Pass/Pass.h"
27 #include "mlir/Transforms/LoopUtils.h"
28 
29 using namespace mlir;
30 using namespace mlir::linalg;
31 using namespace mlir::scf;
32 
33 Optional<RegionMatcher::BinaryOpKind>
34 RegionMatcher::matchAsScalarBinaryOp(GenericOp op) {
35   auto &region = op.region();
36   if (!llvm::hasSingleElement(region))
37     return llvm::None;
38 
39   Block &block = region.front();
40   if (block.getNumArguments() != 2 ||
41       !block.getArgument(0).getType().isSignlessIntOrFloat() ||
42       !block.getArgument(1).getType().isSignlessIntOrFloat())
43     return llvm::None;
44 
45   auto &ops = block.getOperations();
46   if (!llvm::hasSingleElement(block.without_terminator()))
47     return llvm::None;
48 
49   using mlir::matchers::m_Val;
50   auto a = m_Val(block.getArgument(0));
51   auto b = m_Val(block.getArgument(1));
52 
53   auto addPattern = m_Op<linalg::YieldOp>(m_Op<AddIOp>(a, b));
54   if (addPattern.match(&ops.back()))
55     return BinaryOpKind::IAdd;
56 
57   return llvm::None;
58 }
59 
60 bool mlir::linalg::isParallelIteratorType(Attribute attr) {
61   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
62     return strAttr.getValue() == getParallelIteratorTypeName();
63   }
64   return false;
65 }
66 
67 bool mlir::linalg::isReductionIteratorType(Attribute attr) {
68   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
69     return strAttr.getValue() == getReductionIteratorTypeName();
70   }
71   return false;
72 }
73 
74 bool mlir::linalg::isWindowIteratorType(Attribute attr) {
75   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
76     return strAttr.getValue() == getWindowIteratorTypeName();
77   }
78   return false;
79 }
80 
81 /// Explicit instantiation of loop nest generator for different loop types.
82 template struct mlir::linalg::GenerateLoopNest<scf::ForOp>;
83 template struct mlir::linalg::GenerateLoopNest<scf::ParallelOp>;
84 template struct mlir::linalg::GenerateLoopNest<AffineForOp>;
85 
86 /// Given a list of subview ranges, extract individual values for lower, upper
87 /// bounds and steps and put them into the corresponding vectors.
88 static void unpackRanges(ArrayRef<Range> ranges, SmallVectorImpl<Value> &lbs,
89                          SmallVectorImpl<Value> &ubs,
90                          SmallVectorImpl<Value> &steps) {
91   for (Range range : ranges) {
92     lbs.emplace_back(range.offset);
93     ubs.emplace_back(range.size);
94     steps.emplace_back(range.stride);
95   }
96 }
97 
98 namespace mlir {
99 namespace linalg {
100 
101 /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp
102 /// is a constant then return a new value set to the smallest such constant.
103 /// Otherwise returngetSmallestBoundingIndex nullptr.
104 IntegerAttr getSmallestBoundingIndex(Value size) {
105   Optional<int64_t> boundingConst = {};
106   if (auto affineMinOp = size.getDefiningOp<AffineMinOp>()) {
107     for (auto e : affineMinOp.getAffineMap().getResults())
108       if (auto cst = e.dyn_cast<AffineConstantExpr>())
109         boundingConst = boundingConst
110                             ? std::min(boundingConst.getValue(), cst.getValue())
111                             : cst.getValue();
112   } else if (auto constIndexOp = size.getDefiningOp<ConstantOp>()) {
113     if (constIndexOp.getType().isa<IndexType>())
114       boundingConst = constIndexOp.value().cast<IntegerAttr>().getInt();
115   } else if (auto affineApplyOp = size.getDefiningOp<AffineApplyOp>()) {
116     if (auto cExpr = affineApplyOp.getAffineMap()
117                          .getResult(0)
118                          .dyn_cast<AffineConstantExpr>())
119       boundingConst = cExpr.getValue();
120   }
121   if (boundingConst && *boundingConst >= 0)
122     return Builder(size.getContext()).getIndexAttr(*boundingConst);
123   return nullptr;
124 }
125 
126 /// Specialization to build an scf "for" nest.
127 template <>
128 void GenerateLoopNest<scf::ForOp>::doit(
129     ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
130     ArrayRef<Attribute> iteratorTypes,
131     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
132     Optional<LinalgLoopDistributionOptions> distributionOptions) {
133   // Create procInfo so it dominates loops, if appropriate.
134   OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
135   Location loc = edsc::ScopedContext::getLocation();
136   SmallVector<ProcInfo, 2> procInfo;
137   if (distributionOptions.hasValue())
138     procInfo = distributionOptions->procInfo(builder, loc, loopRanges);
139 
140   SmallVector<Value, 4> lbs, ubs, steps;
141   unpackRanges(loopRanges, lbs, ubs, steps);
142   LoopNest loopNest =
143       edsc::loopNestBuilder(lbs, ubs, steps, iterArgInitValues, bodyBuilderFn);
144 
145   if (!distributionOptions.hasValue() || loopNest.loops.empty())
146     return;
147 
148   // Only supports cyclic distribution for now.
149   for (auto it : llvm::zip(loopNest.loops, procInfo,
150                            distributionOptions->distributionMethod))
151     if (std::get<2>(it) == DistributionMethod::Cyclic)
152       mapLoopToProcessorIds(std::get<0>(it), std::get<1>(it).procId,
153                             std::get<1>(it).nprocs);
154 }
155 
156 /// Specialization to build affine "for" nest.
157 template <>
158 void GenerateLoopNest<AffineForOp>::doit(
159     ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
160     ArrayRef<Attribute> iteratorTypes,
161     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
162     Optional<LinalgLoopDistributionOptions>) {
163   assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
164   SmallVector<Value, 4> lbs, ubs, steps;
165   unpackRanges(loopRanges, lbs, ubs, steps);
166 
167   // Affine loops require constant steps.
168   SmallVector<int64_t, 4> constantSteps;
169   constantSteps.reserve(steps.size());
170   for (Value v : steps) {
171     auto op = v.getDefiningOp<ConstantIndexOp>();
172     assert(op && "Affine loops require constant steps");
173     constantSteps.push_back(op.getValue());
174   }
175 
176   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
177     bodyBuilderFn(ivs, {});
178   };
179   edsc::affineLoopNestBuilder(lbs, ubs, constantSteps,
180                               bodyBuilderWithoutIterArgsFn);
181 }
182 
183 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
184 static void updateBoundsForCyclicDistribution(OpBuilder &builder, Location loc,
185                                               Value procId, Value nprocs,
186                                               Value &lb, Value &ub,
187                                               Value &step) {
188   using edsc::op::operator+;
189   using edsc::op::operator*;
190   lb = lb + (procId * step);
191   step = nprocs * step;
192 }
193 
194 /// Generates a loop nest consisting of scf.parallel and scf.for, depending
195 /// on the `iteratorTypes.` Consecutive parallel loops create a single
196 /// scf.parallel operation; each sequential loop creates a new scf.for
197 /// operation. The body of the innermost loop is populated by
198 /// `bodyBuilderFn` that accepts a range of induction variables for all
199 /// loops. `ivStorage` is used to store the partial list of induction
200 /// variables.
201 // TODO: this function can be made iterative instead. However, it
202 // will have at most as many recursive calls as nested loops, which rarely
203 // exceeds 10.
204 static void
205 generateParallelLoopNest(ValueRange lbs, ValueRange ubs, ValueRange steps,
206                          ArrayRef<Attribute> iteratorTypes,
207                          function_ref<void(ValueRange)> bodyBuilderFn,
208                          SmallVectorImpl<Value> &ivStorage,
209                          ArrayRef<DistributionMethod> distributionMethod = {}) {
210   assert(lbs.size() == ubs.size());
211   assert(lbs.size() == steps.size());
212   assert(lbs.size() == iteratorTypes.size());
213 
214   // If there are no (more) loops to be generated, generate the body and be
215   // done with it.
216   if (iteratorTypes.empty())
217     return bodyBuilderFn(ivStorage);
218 
219   // Find the outermost parallel loops and drop their types from the list.
220   unsigned nLoops = iteratorTypes.size();
221   unsigned nOuterPar =
222       nLoops - iteratorTypes.drop_while(isParallelIteratorType).size();
223 
224   // If there are no outer parallel loops, generate one sequential loop and
225   // recurse. Note that we wouldn't have dropped anything from `iteratorTypes`
226   // in this case.
227   if (nOuterPar == 0) {
228     edsc::loopNestBuilder(lbs[0], ubs[0], steps[0], [&](Value iv) {
229       ivStorage.push_back(iv);
230       generateParallelLoopNest(lbs.drop_front(), ubs.drop_front(),
231                                steps.drop_front(), iteratorTypes.drop_front(),
232                                bodyBuilderFn, ivStorage, distributionMethod);
233     });
234     return;
235   }
236   if (distributionMethod.empty()) {
237     // Generate a single parallel loop-nest operation for all outermost
238     // parallel loops and recurse.
239     edsc::OperationBuilder<scf::ParallelOp>(
240         lbs.take_front(nOuterPar), ubs.take_front(nOuterPar),
241         steps.take_front(nOuterPar),
242         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
243           edsc::ScopedContext context(nestedBuilder, nestedLoc);
244           ivStorage.append(localIvs.begin(), localIvs.end());
245           generateParallelLoopNest(
246               lbs.drop_front(nOuterPar), ubs.drop_front(nOuterPar),
247               steps.drop_front(nOuterPar), iteratorTypes.drop_front(nOuterPar),
248               bodyBuilderFn, ivStorage,
249               (distributionMethod.size() < nOuterPar)
250                   ? ArrayRef<DistributionMethod>()
251                   : distributionMethod.drop_front(nOuterPar));
252         });
253     return;
254   }
255 
256   // Process all consecutive similarly distributed loops simultaneously.
257   DistributionMethod methodToUse = distributionMethod[0];
258   unsigned numProcessed = 1;
259   for (unsigned i = 1; i < nOuterPar && i < distributionMethod.size(); ++i) {
260     if (distributionMethod[i] != methodToUse)
261       break;
262     numProcessed++;
263   }
264 
265   switch (methodToUse) {
266   case DistributionMethod::Cyclic: {
267     // Generate a single parallel loop-nest operation for all outermost
268     // parallel loops and recurse.
269     edsc::OperationBuilder<scf::ParallelOp>(
270         lbs.take_front(numProcessed), ubs.take_front(numProcessed),
271         steps.take_front(numProcessed),
272         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
273           edsc::ScopedContext context(nestedBuilder, nestedLoc);
274           ivStorage.append(localIvs.begin(), localIvs.end());
275           generateParallelLoopNest(
276               lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
277               steps.drop_front(numProcessed),
278               iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
279               (distributionMethod.size() < numProcessed)
280                   ? ArrayRef<DistributionMethod>()
281                   : distributionMethod.drop_front(numProcessed));
282         });
283     return;
284   }
285   case DistributionMethod::CyclicNumProcsGeNumIters: {
286     // Check (for the processed loops) that the iteration is in-bounds.
287     using edsc::op::slt;
288     using edsc::op::operator&&;
289     Value cond = slt(lbs[0], ubs[0]);
290     for (unsigned i = 1; i < numProcessed; ++i)
291       cond = cond && slt(lbs[i], ubs[i]);
292     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
293     edsc::conditionBuilder(cond, [&]() {
294       generateParallelLoopNest(
295           lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
296           steps.drop_front(numProcessed),
297           iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
298           distributionMethod.drop_front(numProcessed));
299     });
300     return;
301   }
302   case DistributionMethod::CyclicNumProcsEqNumIters:
303     // No check/loops needed here. Set the `%iv` to be the `%lb` and proceed
304     // with inner loop generation.
305     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
306     generateParallelLoopNest(
307         lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
308         steps.drop_front(numProcessed), iteratorTypes.drop_front(numProcessed),
309         bodyBuilderFn, ivStorage, distributionMethod.drop_front(numProcessed));
310     return;
311   }
312 }
313 
314 /// Specialization for generating a mix of parallel and sequential scf loops.
315 template <>
316 void GenerateLoopNest<scf::ParallelOp>::doit(
317     ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
318     ArrayRef<Attribute> iteratorTypes,
319     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
320     Optional<LinalgLoopDistributionOptions> distributionOptions) {
321   assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");
322   // This function may be passed more iterator types than ranges.
323   assert(iteratorTypes.size() >= loopRanges.size() &&
324          "expected iterator type for all ranges");
325   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
326   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
327   unsigned numLoops = iteratorTypes.size();
328   ivs.reserve(numLoops);
329   lbsStorage.reserve(numLoops);
330   ubsStorage.reserve(numLoops);
331   stepsStorage.reserve(numLoops);
332 
333   // Get the loop lb, ub, and step.
334   unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
335 
336   // Modify the lb, ub, and step based on the distribution options.
337   SmallVector<DistributionMethod, 0> distributionMethod;
338   if (distributionOptions) {
339     auto &options = distributionOptions.getValue();
340     OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
341     Location loc = edsc::ScopedContext::getLocation();
342     distributionMethod.assign(distributionOptions->distributionMethod.begin(),
343                               distributionOptions->distributionMethod.end());
344     SmallVector<Range, 2> parallelLoopRanges;
345     for (auto iteratorType : enumerate(iteratorTypes)) {
346       if (isParallelIteratorType(iteratorType.value()))
347         parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
348     }
349     if (distributionMethod.size() < parallelLoopRanges.size())
350       parallelLoopRanges.resize(distributionMethod.size());
351     SmallVector<ProcInfo, 2> procInfo =
352         options.procInfo(builder, loc, parallelLoopRanges);
353     unsigned index = 0;
354     for (auto iteratorType : enumerate(iteratorTypes)) {
355       if (index >= procInfo.size())
356         break;
357       if (isParallelIteratorType(iteratorType.value())) {
358         unsigned i = iteratorType.index();
359         updateBoundsForCyclicDistribution(builder, loc, procInfo[index].procId,
360                                           procInfo[index].nprocs, lbsStorage[i],
361                                           ubsStorage[i], stepsStorage[i]);
362         index++;
363       }
364     }
365   }
366   ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
367   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
368     bodyBuilderFn(ivs, {});
369   };
370   generateParallelLoopNest(lbs, ubs, steps, iteratorTypes,
371                            bodyBuilderWithoutIterArgsFn, ivs,
372                            distributionMethod);
373 
374   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
375 }
376 
377 } // namespace linalg
378 } // namespace mlir
379