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