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/FoldUtils.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 static Value emitOrFoldComposedAffineApply(OpBuilder &b, Location loc,
61                                            AffineMap map,
62                                            ValueRange operandsRef,
63                                            OperationFolder *folder) {
64   SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
65   fullyComposeAffineMapAndOperands(&map, &operands);
66   canonicalizeMapAndOperands(&map, &operands);
67   return folder ? folder->create<AffineApplyOp>(b, loc, map, operands)
68                 : b.create<AffineApplyOp>(loc, map, operands);
69 }
70 
71 SmallVector<Value, 4> mlir::linalg::applyMapToValues(OpBuilder &b, Location loc,
72                                                      AffineMap map,
73                                                      ValueRange values,
74                                                      OperationFolder *folder) {
75   SmallVector<Value, 4> res;
76   res.reserve(map.getNumResults());
77   unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols();
78   // For each `expr` in `map`, applies the `expr` to the values extracted from
79   // ranges. If the resulting application can be folded into a Value, the
80   // folding occurs eagerly. Otherwise, an affine.apply operation is emitted.
81   for (auto expr : map.getResults()) {
82     AffineMap map = AffineMap::get(numDims, numSym, expr);
83     res.push_back(emitOrFoldComposedAffineApply(b, loc, map, values, folder));
84   }
85   return res;
86 }
87 
88 /// Returns all the operands of `linalgOp` that are not views.
89 /// Asserts that these operands are value types to allow transformations like
90 /// tiling to just use the values when cloning `linalgOp`.
91 SmallVector<Value, 4>
92 mlir::linalg::getAssumedNonViewOperands(LinalgOp linalgOp) {
93   auto *op = linalgOp.getOperation();
94   unsigned numViews = linalgOp.getNumInputsAndOutputs();
95   unsigned nOperands = op->getNumOperands() - numViews;
96   SmallVector<Value, 4> res;
97   res.reserve(nOperands);
98   for (unsigned i = 0; i < nOperands; ++i) {
99     res.push_back(op->getOperand(numViews + i));
100     auto t = res.back().getType();
101     (void)t;
102     assert((t.isSignlessIntOrIndexOrFloat() || t.isa<VectorType>()) &&
103            "expected scalar or vector type");
104   }
105   return res;
106 }
107 
108 bool mlir::linalg::isParallelIteratorType(Attribute attr) {
109   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
110     return strAttr.getValue() == getParallelIteratorTypeName();
111   }
112   return false;
113 }
114 
115 bool mlir::linalg::isReductionIteratorType(Attribute attr) {
116   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
117     return strAttr.getValue() == getReductionIteratorTypeName();
118   }
119   return false;
120 }
121 
122 bool mlir::linalg::isWindowIteratorType(Attribute attr) {
123   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
124     return strAttr.getValue() == getWindowIteratorTypeName();
125   }
126   return false;
127 }
128 
129 /// Explicit instantiation of loop nest generator for different loop types.
130 template struct mlir::linalg::GenerateLoopNest<scf::ForOp>;
131 template struct mlir::linalg::GenerateLoopNest<scf::ParallelOp>;
132 template struct mlir::linalg::GenerateLoopNest<AffineForOp>;
133 
134 /// Given a list of subview ranges, extract individual values for lower, upper
135 /// bounds and steps and put them into the corresponding vectors.
136 static void unpackRanges(ArrayRef<SubViewOp::Range> ranges,
137                          SmallVectorImpl<Value> &lbs,
138                          SmallVectorImpl<Value> &ubs,
139                          SmallVectorImpl<Value> &steps) {
140   for (SubViewOp::Range range : ranges) {
141     lbs.emplace_back(range.offset);
142     ubs.emplace_back(range.size);
143     steps.emplace_back(range.stride);
144   }
145 }
146 
147 namespace mlir {
148 namespace linalg {
149 
150 /// Return the linearized list of all view dimensions in a linalgOp.
151 SmallVector<Value, 8> getViewSizes(OpBuilder &builder, LinalgOp linalgOp) {
152   auto loc = linalgOp.getLoc();
153   SmallVector<Value, 8> res;
154   SmallVector<unsigned, 4> ranks;
155   for (auto v : linalgOp.getInputsAndOutputBuffers()) {
156     MemRefType t = v.getType().template cast<MemRefType>();
157     ranks.push_back(t.getRank());
158     for (unsigned i = 0; i < t.getRank(); ++i)
159       res.push_back(builder.create<DimOp>(loc, v, i));
160   }
161 
162   auto attr = linalgOp.template getAttrOfType<IntegerAttr>("symbol_source");
163   if (attr) {
164     // Find the correct position for inserting values for symbols.
165     unsigned numSymb = ranks[attr.getInt()], symbolsPos = 0;
166     for (unsigned idx = 0; idx < attr.getInt(); idx++)
167       symbolsPos += ranks[idx];
168 
169     // Append the end of the value list that corresponds to the
170     // values mapping to symbols. Since inside concatinated map symbols are
171     // repeated we have to repeat the sizes as well.
172 
173     // Reserve is mandatory to avoid a potential undefined behavior with
174     // pushing back to smallvector from itself.
175     res.reserve(res.size() + ranks.size() * numSymb);
176     for (unsigned idx = 0, s = ranks.size(); idx < s; ++idx)
177       for (unsigned idx2 = 0; idx2 < numSymb; ++idx2)
178         res.push_back(res[symbolsPos + idx2]);
179   }
180   return res;
181 }
182 
183 Optional<SmallVector<Value, 4>>
184 getLoopRanges(OpBuilder &builder, LinalgOp linalgOp, OperationFolder *folder) {
185   SmallVector<Value, 8> viewSizes = getViewSizes(builder, linalgOp);
186   AffineMap invertedMap =
187       inversePermutation(concatAffineMaps(linalgOp.getIndexingMaps()));
188   if (!invertedMap)
189     return {};
190   return applyMapToValues(builder, linalgOp.getLoc(), invertedMap, viewSizes,
191                           folder);
192 }
193 
194 /// Specialization to build an scf "for" nest.
195 template <>
196 void GenerateLoopNest<scf::ForOp>::doit(
197     ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
198     function_ref<void(ValueRange)> bodyBuilderFn,
199     Optional<LinalgLoopDistributionOptions>) {
200   SmallVector<Value, 4> lbs, ubs, steps;
201   unpackRanges(loopRanges, lbs, ubs, steps);
202   edsc::loopNestBuilder(lbs, ubs, steps, bodyBuilderFn);
203 }
204 
205 /// Specialization to build affine "for" nest.
206 template <>
207 void GenerateLoopNest<AffineForOp>::doit(
208     ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
209     function_ref<void(ValueRange)> bodyBuilderFn,
210     Optional<LinalgLoopDistributionOptions>) {
211   SmallVector<Value, 4> lbs, ubs, steps;
212   unpackRanges(loopRanges, lbs, ubs, steps);
213 
214   // Affine loops require constant steps.
215   SmallVector<int64_t, 4> constantSteps;
216   constantSteps.reserve(steps.size());
217   for (Value v : steps) {
218     auto op = v.getDefiningOp<ConstantIndexOp>();
219     assert(op && "Affine loops require constant steps");
220     constantSteps.push_back(op.getValue());
221   }
222 
223   edsc::affineLoopNestBuilder(lbs, ubs, constantSteps, bodyBuilderFn);
224 }
225 
226 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
227 static void updateBoundsForCyclicDistribution(OpBuilder &builder, Location loc,
228                                               Value procId, Value nprocs,
229                                               Value &lb, Value &ub,
230                                               Value &step) {
231   using edsc::op::operator+;
232   using edsc::op::operator*;
233   lb = lb + (procId * step);
234   step = nprocs * step;
235 }
236 
237 /// Generates a loop nest consisting of scf.parallel and scf.for, depending
238 /// on the `iteratorTypes.` Consecutive parallel loops create a single
239 /// scf.parallel operation; each sequential loop creates a new scf.for
240 /// operation. The body of the innermost loop is populated by
241 /// `bodyBuilderFn` that accepts a range of induction variables for all
242 /// loops. `ivStorage` is used to store the partial list of induction
243 /// variables.
244 // TODO: this function can be made iterative instead. However, it
245 // will have at most as many recursive calls as nested loops, which rarely
246 // exceeds 10.
247 static void
248 generateParallelLoopNest(ValueRange lbs, ValueRange ubs, ValueRange steps,
249                          ArrayRef<Attribute> iteratorTypes,
250                          function_ref<void(ValueRange)> bodyBuilderFn,
251                          SmallVectorImpl<Value> &ivStorage,
252                          ArrayRef<DistributionMethod> distributionMethod = {}) {
253   assert(lbs.size() == ubs.size());
254   assert(lbs.size() == steps.size());
255   assert(lbs.size() == iteratorTypes.size());
256 
257   // If there are no (more) loops to be generated, generate the body and be
258   // done with it.
259   if (iteratorTypes.empty())
260     return bodyBuilderFn(ivStorage);
261 
262   // Find the outermost parallel loops and drop their types from the list.
263   unsigned nLoops = iteratorTypes.size();
264   unsigned nOuterPar =
265       nLoops - iteratorTypes.drop_while(isParallelIteratorType).size();
266 
267   // If there are no outer parallel loops, generate one sequential loop and
268   // recurse. Note that we wouldn't have dropped anything from `iteratorTypes`
269   // in this case.
270   if (nOuterPar == 0) {
271     edsc::loopNestBuilder(lbs[0], ubs[0], steps[0], [&](Value iv) {
272       ivStorage.push_back(iv);
273       generateParallelLoopNest(lbs.drop_front(), ubs.drop_front(),
274                                steps.drop_front(), iteratorTypes.drop_front(),
275                                bodyBuilderFn, ivStorage, distributionMethod);
276     });
277     return;
278   }
279   if (distributionMethod.empty()) {
280     // Generate a single parallel loop-nest operation for all outermost
281     // parallel loops and recurse.
282     edsc::OperationBuilder<scf::ParallelOp>(
283         lbs.take_front(nOuterPar), ubs.take_front(nOuterPar),
284         steps.take_front(nOuterPar),
285         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
286           edsc::ScopedContext context(nestedBuilder, nestedLoc);
287           ivStorage.append(localIvs.begin(), localIvs.end());
288           generateParallelLoopNest(
289               lbs.drop_front(nOuterPar), ubs.drop_front(nOuterPar),
290               steps.drop_front(nOuterPar), iteratorTypes.drop_front(nOuterPar),
291               bodyBuilderFn, ivStorage,
292               (distributionMethod.size() < nOuterPar)
293                   ? ArrayRef<DistributionMethod>()
294                   : distributionMethod.drop_front(nOuterPar));
295         });
296     return;
297   }
298 
299   // Process all consecutive similarly distributed loops simultaneously.
300   DistributionMethod methodToUse = distributionMethod[0];
301   unsigned numProcessed = 1;
302   for (unsigned i = 1; i < nOuterPar && i < distributionMethod.size(); ++i) {
303     if (distributionMethod[i] != methodToUse)
304       break;
305     numProcessed++;
306   }
307 
308   switch (methodToUse) {
309   case DistributionMethod::Cyclic: {
310     // Generate a single parallel loop-nest operation for all outermost
311     // parallel loops and recurse.
312     edsc::OperationBuilder<scf::ParallelOp>(
313         lbs.take_front(numProcessed), ubs.take_front(numProcessed),
314         steps.take_front(numProcessed),
315         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
316           edsc::ScopedContext context(nestedBuilder, nestedLoc);
317           ivStorage.append(localIvs.begin(), localIvs.end());
318           generateParallelLoopNest(
319               lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
320               steps.drop_front(numProcessed),
321               iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
322               (distributionMethod.size() < numProcessed)
323                   ? ArrayRef<DistributionMethod>()
324                   : distributionMethod.drop_front(numProcessed));
325         });
326     return;
327   }
328   case DistributionMethod::CyclicNumProcsGeNumIters: {
329     // Check (for the processed loops) that the iteration is in-bounds.
330     using edsc::op::slt;
331     using edsc::op::operator&&;
332     Value cond = slt(lbs[0], ubs[0]);
333     for (unsigned i = 1; i < numProcessed; ++i)
334       cond = cond && slt(lbs[i], ubs[i]);
335     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
336     edsc::conditionBuilder(cond, [&]() {
337       generateParallelLoopNest(
338           lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
339           steps.drop_front(numProcessed),
340           iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
341           distributionMethod.drop_front(numProcessed));
342     });
343     return;
344   }
345   case DistributionMethod::CyclicNumProcsEqNumIters:
346     // No check/loops needed here. Set the `%iv` to be the `%lb` and proceed
347     // with inner loop generation.
348     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
349     generateParallelLoopNest(
350         lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
351         steps.drop_front(numProcessed), iteratorTypes.drop_front(numProcessed),
352         bodyBuilderFn, ivStorage, distributionMethod.drop_front(numProcessed));
353     return;
354   }
355 }
356 
357 /// Specialization for generating a mix of parallel and sequential scf loops.
358 template <>
359 void GenerateLoopNest<scf::ParallelOp>::doit(
360     ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
361     function_ref<void(ValueRange)> bodyBuilderFn,
362     Optional<LinalgLoopDistributionOptions> distributionOptions) {
363   // This function may be passed more iterator types than ranges.
364   assert(iteratorTypes.size() >= loopRanges.size() &&
365          "expected iterator type for all ranges");
366   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
367   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
368   unsigned numLoops = iteratorTypes.size();
369   ivs.reserve(numLoops);
370   lbsStorage.reserve(numLoops);
371   ubsStorage.reserve(numLoops);
372   stepsStorage.reserve(numLoops);
373 
374   // Get the loop lb, ub, and step.
375   unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
376 
377   // Modify the lb, ub, and step based on the distribution options.
378   SmallVector<DistributionMethod, 0> distributionMethod;
379   if (distributionOptions) {
380     auto &options = distributionOptions.getValue();
381     OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
382     Location loc = edsc::ScopedContext::getLocation();
383     distributionMethod.assign(distributionOptions->distributionMethod.begin(),
384                               distributionOptions->distributionMethod.end());
385     SmallVector<SubViewOp::Range, 2> parallelLoopRanges;
386     for (auto iteratorType : enumerate(iteratorTypes)) {
387       if (isParallelIteratorType(iteratorType.value()))
388         parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
389     }
390     if (distributionMethod.size() < parallelLoopRanges.size())
391       parallelLoopRanges.resize(distributionMethod.size());
392     SmallVector<ProcInfo, 2> procInfo =
393         options.procInfo(builder, loc, parallelLoopRanges);
394     unsigned index = 0;
395     for (auto iteratorType : enumerate(iteratorTypes)) {
396       if (index >= procInfo.size())
397         break;
398       if (isParallelIteratorType(iteratorType.value())) {
399         unsigned i = iteratorType.index();
400         updateBoundsForCyclicDistribution(builder, loc, procInfo[index].procId,
401                                           procInfo[index].nprocs, lbsStorage[i],
402                                           ubsStorage[i], stepsStorage[i]);
403         index++;
404       }
405     }
406   }
407   ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
408   generateParallelLoopNest(lbs, ubs, steps, iteratorTypes, bodyBuilderFn, ivs,
409                            distributionMethod);
410 
411   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
412 }
413 
414 } // namespace linalg
415 } // namespace mlir
416