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/IR/AffineOps.h"
16 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
17 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
18 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
19 #include "mlir/Dialect/MemRef/IR/MemRef.h"
20 #include "mlir/Dialect/SCF/SCF.h"
21 #include "mlir/Dialect/StandardOps/IR/Ops.h"
22 #include "mlir/Dialect/StandardOps/Utils/Utils.h"
23 #include "mlir/Dialect/Tensor/IR/Tensor.h"
24 #include "mlir/Dialect/Utils/StaticValueUtils.h"
25 #include "mlir/IR/AffineExpr.h"
26 #include "mlir/IR/AffineExprVisitor.h"
27 #include "mlir/IR/AffineMap.h"
28 #include "mlir/IR/Matchers.h"
29 #include "mlir/IR/OpImplementation.h"
30 #include "mlir/Pass/Pass.h"
31 #include "mlir/Transforms/LoopUtils.h"
32 #include "llvm/ADT/TypeSwitch.h"
33 #include "llvm/Support/Debug.h"
34 
35 #define DEBUG_TYPE "linalg-utils"
36 
37 using namespace mlir;
38 using namespace mlir::linalg;
39 using namespace mlir::scf;
40 
41 static bool isZero(Value v) {
42   if (auto cst = v.getDefiningOp<arith::ConstantIndexOp>())
43     return cst.value() == 0;
44   return false;
45 }
46 
47 namespace {
48 
49 // Helper visitor to determine whether an AffineExpr is tiled.
50 // This is achieved by traversing every AffineDimExpr with position `pos` and
51 // checking whether the corresponding `tileSizes[pos]` is non-zero.
52 // This also enforces only positive coefficients occur in multiplications.
53 //
54 // Example:
55 //   `d0 + 2 * d1 + d3` is tiled by [0, 0, 0, 2] but not by [0, 0, 2, 0]
56 //
57 struct TileCheck : public AffineExprVisitor<TileCheck> {
58   TileCheck(ValueRange tileSizes) : isTiled(false), tileSizes(tileSizes) {}
59 
60   void visitDimExpr(AffineDimExpr expr) {
61     isTiled |= !isZero(tileSizes[expr.getPosition()]);
62   }
63   void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) {
64     visit(expr.getLHS());
65     visit(expr.getRHS());
66     if (expr.getKind() == mlir::AffineExprKind::Mul)
67       assert(expr.getRHS().cast<AffineConstantExpr>().getValue() > 0 &&
68              "nonpositive multiplying coefficient");
69   }
70   bool isTiled;
71   ValueRange tileSizes;
72 };
73 
74 } // namespace
75 
76 static bool isTiled(AffineExpr expr, ValueRange tileSizes) {
77   if (!expr)
78     return false;
79   TileCheck t(tileSizes);
80   t.visit(expr);
81   return t.isTiled;
82 }
83 
84 // Checks whether the `map  varies with respect to a non-zero `tileSize`.
85 static bool isTiled(AffineMap map, ValueRange tileSizes) {
86   if (!map)
87     return false;
88   for (unsigned r = 0; r < map.getNumResults(); ++r)
89     if (isTiled(map.getResult(r), tileSizes))
90       return true;
91   return false;
92 }
93 
94 Optional<RegionMatcher::BinaryOpKind>
95 RegionMatcher::matchAsScalarBinaryOp(GenericOp op) {
96   auto &region = op.region();
97   if (!llvm::hasSingleElement(region))
98     return llvm::None;
99 
100   Block &block = region.front();
101   if (block.getNumArguments() != 2 ||
102       !block.getArgument(0).getType().isSignlessIntOrFloat() ||
103       !block.getArgument(1).getType().isSignlessIntOrFloat())
104     return llvm::None;
105 
106   auto &ops = block.getOperations();
107   if (!llvm::hasSingleElement(block.without_terminator()))
108     return llvm::None;
109 
110   using mlir::matchers::m_Val;
111   auto a = m_Val(block.getArgument(0));
112   auto b = m_Val(block.getArgument(1));
113 
114   auto addPattern = m_Op<linalg::YieldOp>(m_Op<arith::AddIOp>(a, b));
115   if (addPattern.match(&ops.back()))
116     return BinaryOpKind::IAdd;
117 
118   return llvm::None;
119 }
120 
121 /// Explicit instantiation of loop nest generator for different loop types.
122 template struct mlir::linalg::GenerateLoopNest<scf::ForOp>;
123 template struct mlir::linalg::GenerateLoopNest<scf::ParallelOp>;
124 template struct mlir::linalg::GenerateLoopNest<AffineForOp>;
125 template struct mlir::linalg::GenerateLoopNest<TiledLoopOp>;
126 
127 /// Given a list of subview ranges, extract individual values for lower, upper
128 /// bounds and steps and put them into the corresponding vectors.
129 static void unpackRanges(ArrayRef<Range> ranges, SmallVectorImpl<Value> &lbs,
130                          SmallVectorImpl<Value> &ubs,
131                          SmallVectorImpl<Value> &steps) {
132   for (Range range : ranges) {
133     lbs.emplace_back(range.offset);
134     ubs.emplace_back(range.size);
135     steps.emplace_back(range.stride);
136   }
137 }
138 
139 namespace mlir {
140 namespace linalg {
141 
142 bool isPermutation(ArrayRef<int64_t> permutation) {
143   // Count the number of appearances for all indices.
144   SmallVector<int64_t> indexCounts(permutation.size(), 0);
145   for (auto index : permutation) {
146     // Exit if the index is out-of-range.
147     if (index < 0 || index >= static_cast<int64_t>(permutation.size()))
148       return false;
149     indexCounts[index]++;
150   }
151   // Return true if all indices appear once.
152   return count(indexCounts, 1) == static_cast<int64_t>(permutation.size());
153 }
154 
155 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
156 /// the type of `source`.
157 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim) {
158   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
159     return b.createOrFold<memref::DimOp>(loc, source, dim);
160   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
161     return b.createOrFold<tensor::DimOp>(loc, source, dim);
162   llvm_unreachable("Expected MemRefType or TensorType");
163 }
164 
165 /// Given an operation, retrieves the value of each dynamic dimension through
166 /// constructing the necessary DimOp operators.
167 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b) {
168   SmallVector<Value, 4> dynOperands;
169   auto shapedType = val.getType().cast<ShapedType>();
170   for (auto dim : llvm::enumerate(shapedType.getShape())) {
171     if (dim.value() == ShapedType::kDynamicSize)
172       dynOperands.push_back(createOrFoldDimOp(b, loc, val, dim.index()));
173   }
174   return dynOperands;
175 }
176 
177 /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp
178 /// is a constant then return a new value set to the smallest such constant.
179 /// Otherwise returngetSmallestBoundingIndex nullptr.
180 IntegerAttr getSmallestBoundingIndex(Value size) {
181   Optional<int64_t> boundingConst = {};
182   if (auto affineMinOp = size.getDefiningOp<AffineMinOp>()) {
183     for (auto e : affineMinOp.getAffineMap().getResults())
184       if (auto cst = e.dyn_cast<AffineConstantExpr>())
185         boundingConst = boundingConst
186                             ? std::min(boundingConst.getValue(), cst.getValue())
187                             : cst.getValue();
188   } else if (auto constIndexOp = size.getDefiningOp<arith::ConstantOp>()) {
189     if (constIndexOp.getType().isa<IndexType>())
190       boundingConst = constIndexOp.getValue().cast<IntegerAttr>().getInt();
191   } else if (auto affineApplyOp = size.getDefiningOp<AffineApplyOp>()) {
192     if (auto cExpr = affineApplyOp.getAffineMap()
193                          .getResult(0)
194                          .dyn_cast<AffineConstantExpr>())
195       boundingConst = cExpr.getValue();
196   } else if (auto dimOp = size.getDefiningOp<tensor::DimOp>()) {
197     auto shape = dimOp.source().getType().dyn_cast<ShapedType>();
198     if (auto constOp = dimOp.index().getDefiningOp<arith::ConstantOp>()) {
199       if (auto indexAttr = constOp.getValue().dyn_cast<IntegerAttr>()) {
200         auto dimIndex = indexAttr.getInt();
201         if (!shape.isDynamicDim(dimIndex)) {
202           boundingConst = shape.getShape()[dimIndex];
203         }
204       }
205     }
206   }
207   if (boundingConst && *boundingConst >= 0)
208     return Builder(size.getContext()).getIndexAttr(*boundingConst);
209   return nullptr;
210 }
211 
212 tensor::ExtractSliceOp makeComposedExtractSliceOp(
213     OpBuilder &b, Location loc, Value source, ArrayRef<OpFoldResult> offsets,
214     ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) {
215   assert(source && "expect source to be nonzero");
216 
217   // Do not fold if the producer is not an ExtractSliceOp.
218   auto producerOp = source.getDefiningOp<tensor::ExtractSliceOp>();
219   if (!producerOp)
220     return b.create<tensor::ExtractSliceOp>(loc, source, offsets, sizes,
221                                             strides);
222 
223   // Do not fold if the producer is rank reducing or if there are any non-unit
224   // strides. Supporting non-unit strides complicates the offset computation
225   // since the consumer offsets need to be multiplied by the producer strides.
226   // TODO: support non-unit strides once there are use cases.
227   SmallVector<OpFoldResult> allStrides = producerOp.getMixedStrides();
228   allStrides.append(strides.begin(), strides.end());
229   bool hasNonUnitStride = any_of(allStrides, [](OpFoldResult ofr) {
230     return getConstantIntValue(ofr) != static_cast<int64_t>(1);
231   });
232   if (hasNonUnitStride ||
233       producerOp.getSourceType().getRank() !=
234           producerOp.getResult().getType().cast<ShapedType>().getRank())
235     return b.create<tensor::ExtractSliceOp>(loc, source, offsets, sizes,
236                                             strides);
237 
238   // Fold the producer by adding the offests and extracting the slice directly
239   // from the producer source tensor.
240   SmallVector<OpFoldResult> foldedOffsets(offsets.begin(), offsets.end());
241   AffineExpr dim1, dim2;
242   bindDims(b.getContext(), dim1, dim2);
243   for (auto en : enumerate(producerOp.getMixedOffsets())) {
244     SmallVector<Value> offsetValues = {
245         getValueOrCreateConstantIndexOp(b, loc, foldedOffsets[en.index()]),
246         getValueOrCreateConstantIndexOp(b, loc, en.value())};
247     foldedOffsets[en.index()] =
248         makeComposedAffineApply(b, loc, dim1 + dim2, offsetValues).getResult();
249   }
250   return b.create<tensor::ExtractSliceOp>(loc, producerOp.source(),
251                                           foldedOffsets, sizes, strides);
252 }
253 
254 /// Specialization to build an scf "for" nest.
255 template <>
256 void GenerateLoopNest<scf::ForOp>::doit(
257     OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
258     ArrayRef<Attribute> iteratorTypes,
259     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
260                                   ValueRange)>
261         bodyBuilderFn,
262     Optional<LinalgLoopDistributionOptions> distributionOptions,
263     ArrayRef<StringRef> distributionTypes) {
264   SmallVector<Value> iterArgInitValues = linalgOp.getOutputTensorOperands();
265   // Create procInfo so it dominates loops, if appropriate.
266   SmallVector<ProcInfo, 4> procInfo;
267   SmallVector<DistributionMethod, 0> distributionMethod;
268   if (distributionOptions.hasValue()) {
269     // Collect loop ranges for parallel dimensions.
270     SmallVector<Range, 2> parallelLoopRanges;
271     for (auto iteratorType : enumerate(iteratorTypes))
272       if (isParallelIterator(iteratorType.value()))
273         parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
274 
275     // Get their distribution schemes.
276     distributionMethod = distributionOptions->distributionMethod;
277     if (distributionMethod.size() < parallelLoopRanges.size())
278       parallelLoopRanges.resize(distributionMethod.size());
279     procInfo = distributionOptions->procInfo(b, loc, parallelLoopRanges);
280   }
281 
282   SmallVector<Value, 4> lbs, ubs, steps;
283   unpackRanges(loopRanges, lbs, ubs, steps);
284   LoopNest loopNest = mlir::scf::buildLoopNest(
285       b, loc, lbs, ubs, steps, iterArgInitValues,
286       [&](OpBuilder &b, Location loc, ValueRange ivs, ValueRange iterArgs) {
287         assert(iterArgs.size() == linalgOp.getOutputTensorOperands().size() &&
288                "expect the number of output tensors and iter args to match");
289         SmallVector<Value> operandValuesToUse =
290             linalgOp.getInputAndOutputOperands();
291         if (!iterArgs.empty()) {
292           operandValuesToUse = linalgOp.getInputOperands();
293           operandValuesToUse.append(iterArgs.begin(), iterArgs.end());
294         }
295         return bodyBuilderFn(b, loc, ivs, operandValuesToUse);
296       });
297 
298   if (!distributionOptions || loopNest.loops.empty())
299     return;
300 
301   // Filter out scf.for loops that were created out of parallel dimensions.
302   SmallVector<scf::ForOp, 4> loops;
303   for (auto iteratorType : enumerate(iteratorTypes))
304     if (isParallelIterator(iteratorType.value()))
305       loops.push_back(loopNest.loops[iteratorType.index()]);
306 
307   // Distribute - only supports cyclic distribution for now.
308   for (auto it : llvm::zip(loops, procInfo, distributionMethod))
309     if (std::get<2>(it) == DistributionMethod::Cyclic)
310       mapLoopToProcessorIds(std::get<0>(it), std::get<1>(it).procId,
311                             std::get<1>(it).nprocs);
312 }
313 
314 /// Specialization to build affine "for" nest.
315 template <>
316 void GenerateLoopNest<AffineForOp>::doit(
317     OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
318     ArrayRef<Attribute> iteratorTypes,
319     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
320                                   ValueRange)>
321         bodyBuilderFn,
322     Optional<LinalgLoopDistributionOptions>, ArrayRef<StringRef>) {
323   SmallVector<Value> iterArgInitValues = linalgOp.getOutputTensorOperands();
324   assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
325   SmallVector<Value, 4> lbs, ubs, steps;
326   unpackRanges(loopRanges, lbs, ubs, steps);
327 
328   // Affine loops require constant steps.
329   SmallVector<int64_t, 4> constantSteps;
330   constantSteps.reserve(steps.size());
331   for (Value v : steps) {
332     auto op = v.getDefiningOp<arith::ConstantIndexOp>();
333     assert(op && "Affine loops require constant steps");
334     constantSteps.push_back(op.value());
335   }
336 
337   mlir::buildAffineLoopNest(b, loc, lbs, ubs, constantSteps,
338                             [&](OpBuilder &b, Location loc, ValueRange ivs) {
339                               SmallVector<Value> operandValuesToUse =
340                                   linalgOp.getInputAndOutputOperands();
341                               bodyBuilderFn(b, loc, ivs, operandValuesToUse);
342                             });
343 }
344 
345 /// Specialization to build an linalg.tiled_loop
346 template <>
347 void GenerateLoopNest<TiledLoopOp>::doit(
348     OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
349     ArrayRef<Attribute> iteratorTypes,
350     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
351                                   ValueRange)>
352         bodyBuilderFn,
353     Optional<LinalgLoopDistributionOptions> distributionOptions,
354     ArrayRef<StringRef> distributionTypes) {
355   SmallVector<ProcInfo, 2> procInfo;
356   SmallVector<Value, 4> lbs, ubs, steps;
357   unpackRanges(loopRanges, lbs, ubs, steps);
358 
359   auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
360                               ValueRange ivs, ValueRange inputs,
361                               ValueRange outputs) {
362     SmallVector<Value> operandValuesToUse = inputs;
363     operandValuesToUse.append(outputs.begin(), outputs.end());
364     scf::ValueVector results =
365         bodyBuilderFn(nestedBuilder, nestedLoc, ivs, operandValuesToUse);
366     nestedBuilder.create<linalg::YieldOp>(nestedLoc, results);
367   };
368 
369   SmallVector<Value> inputOperands = linalgOp.getInputOperands();
370   SmallVector<Value> outputOperands = linalgOp.getOutputOperands();
371   auto tiledLoop =
372       b.create<TiledLoopOp>(loc, lbs, ubs, steps, inputOperands, outputOperands,
373                             b.getArrayAttr(iteratorTypes), wrappedBuilderFn);
374   if (!distributionTypes.empty())
375     tiledLoop.setDistributionTypes(b, distributionTypes);
376 }
377 
378 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
379 void updateBoundsForCyclicDistribution(OpBuilder &b, Location loc, Value procId,
380                                        Value nprocs, Value &lb, Value &ub,
381                                        Value &step) {
382   AffineExpr d0, d1;
383   bindDims(b.getContext(), d0, d1);
384   AffineExpr s0 = getAffineSymbolExpr(0, b.getContext());
385   lb = makeComposedAffineApply(b, loc, d0 + d1 * s0, {lb, procId, step});
386   step = makeComposedAffineApply(b, loc, d0 * s0, {nprocs, step});
387 }
388 
389 /// Generates a loop nest consisting of scf.parallel and scf.for, depending
390 /// on the `iteratorTypes.` Consecutive parallel loops create a single
391 /// scf.parallel operation; each sequential loop creates a new scf.for
392 /// operation. The body of the innermost loop is populated by
393 /// `bodyBuilderFn` that accepts a range of induction variables for all
394 /// loops. `ivStorage` is used to store the partial list of induction
395 /// variables.
396 // TODO: this function can be made iterative instead. However, it
397 // will have at most as many recursive calls as nested loops, which rarely
398 // exceeds 10.
399 static void generateParallelLoopNest(
400     OpBuilder &b, Location loc, ValueRange lbs, ValueRange ubs,
401     ValueRange steps, ArrayRef<Attribute> iteratorTypes,
402     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn,
403     SmallVectorImpl<Value> &ivStorage,
404     ArrayRef<DistributionMethod> distributionMethod = {}) {
405   assert(lbs.size() == ubs.size());
406   assert(lbs.size() == steps.size());
407   assert(lbs.size() == iteratorTypes.size());
408 
409   // If there are no (more) loops to be generated, generate the body and be
410   // done with it.
411   if (iteratorTypes.empty()) {
412     bodyBuilderFn(b, loc, ivStorage);
413     return;
414   }
415 
416   // Find the outermost parallel loops and drop their types from the list.
417   unsigned nLoops = iteratorTypes.size();
418   unsigned nOuterPar =
419       nLoops - iteratorTypes.drop_while(isParallelIterator).size();
420 
421   // If there are no outer parallel loops, generate one sequential loop and
422   // recurse. Note that we wouldn't have dropped anything from `iteratorTypes`
423   // in this case.
424   if (nOuterPar == 0) {
425     LoopNest singleLoop = buildLoopNest(
426         b, loc, lbs.take_front(), ubs.take_front(), steps.take_front(),
427         [&](OpBuilder &b, Location loc, ValueRange ivs) {
428           ivStorage.append(ivs.begin(), ivs.end());
429           generateParallelLoopNest(b, loc, lbs.drop_front(), ubs.drop_front(),
430                                    steps.drop_front(),
431                                    iteratorTypes.drop_front(), bodyBuilderFn,
432                                    ivStorage, distributionMethod);
433         });
434     return;
435   }
436   if (distributionMethod.empty()) {
437     // Generate a single parallel loop-nest operation for all outermost
438     // parallel loops and recurse.
439     b.create<scf::ParallelOp>(
440         loc, lbs.take_front(nOuterPar), ubs.take_front(nOuterPar),
441         steps.take_front(nOuterPar),
442         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
443           ivStorage.append(localIvs.begin(), localIvs.end());
444           generateParallelLoopNest(
445               nestedBuilder, nestedLoc, lbs.drop_front(nOuterPar),
446               ubs.drop_front(nOuterPar), steps.drop_front(nOuterPar),
447               iteratorTypes.drop_front(nOuterPar), bodyBuilderFn, ivStorage,
448               (distributionMethod.size() < nOuterPar)
449                   ? ArrayRef<DistributionMethod>()
450                   : distributionMethod.drop_front(nOuterPar));
451         });
452     return;
453   }
454 
455   // Process all consecutive similarly distributed loops simultaneously.
456   DistributionMethod methodToUse = distributionMethod[0];
457   unsigned numProcessed = 1;
458   for (unsigned i = 1; i < nOuterPar && i < distributionMethod.size(); ++i) {
459     if (distributionMethod[i] != methodToUse)
460       break;
461     numProcessed++;
462   }
463 
464   switch (methodToUse) {
465   case DistributionMethod::Cyclic: {
466     // Generate a single parallel loop-nest operation for all outermost
467     // parallel loops and recurse.
468     b.create<scf::ParallelOp>(
469         loc, lbs.take_front(numProcessed), ubs.take_front(numProcessed),
470         steps.take_front(numProcessed),
471         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
472           ivStorage.append(localIvs.begin(), localIvs.end());
473           generateParallelLoopNest(
474               nestedBuilder, nestedLoc, lbs.drop_front(numProcessed),
475               ubs.drop_front(numProcessed), steps.drop_front(numProcessed),
476               iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
477               (distributionMethod.size() < numProcessed)
478                   ? ArrayRef<DistributionMethod>()
479                   : distributionMethod.drop_front(numProcessed));
480         });
481     return;
482   }
483   case DistributionMethod::CyclicNumProcsGeNumIters: {
484     // Check (for the processed loops) that the iteration is in-bounds.
485     ArithBuilder ab(b, loc);
486     Value cond = ab.slt(lbs[0], ubs[0]);
487     for (unsigned i = 1; i < numProcessed; ++i)
488       cond = ab._and(cond, ab.slt(lbs[i], ubs[i]));
489     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
490     b.create<scf::IfOp>(loc, cond, [&](OpBuilder &b, Location loc) {
491       generateParallelLoopNest(
492           b, loc, lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
493           steps.drop_front(numProcessed),
494           iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
495           distributionMethod.drop_front(numProcessed));
496       b.create<scf::YieldOp>(loc, ValueRange{});
497     });
498     return;
499   }
500   case DistributionMethod::CyclicNumProcsEqNumIters:
501     // No check/loops needed here. Set the `%iv` to be the `%lb` and proceed
502     // with inner loop generation.
503     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
504     generateParallelLoopNest(
505         b, loc, lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
506         steps.drop_front(numProcessed), iteratorTypes.drop_front(numProcessed),
507         bodyBuilderFn, ivStorage, distributionMethod.drop_front(numProcessed));
508     return;
509   }
510 }
511 
512 /// Specialization for generating a mix of parallel and sequential scf loops.
513 template <>
514 void GenerateLoopNest<scf::ParallelOp>::doit(
515     OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
516     ArrayRef<Attribute> iteratorTypes,
517     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
518                                   ValueRange)>
519         bodyBuilderFn,
520     Optional<LinalgLoopDistributionOptions> distributionOptions,
521     ArrayRef<StringRef> distributionTypes) {
522   SmallVector<Value> iterArgInitValues = linalgOp.getOutputTensorOperands();
523   assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");
524   // This function may be passed more iterator types than ranges.
525   assert(iteratorTypes.size() >= loopRanges.size() &&
526          "expected iterator type for all ranges");
527   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
528   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
529   unsigned numLoops = iteratorTypes.size();
530   ivs.reserve(numLoops);
531   lbsStorage.reserve(numLoops);
532   ubsStorage.reserve(numLoops);
533   stepsStorage.reserve(numLoops);
534 
535   // Get the loop lb, ub, and step.
536   unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
537 
538   // Modify the lb, ub, and step based on the distribution options.
539   SmallVector<DistributionMethod, 0> distributionMethod;
540   if (distributionOptions) {
541     auto &options = distributionOptions.getValue();
542     distributionMethod.assign(distributionOptions->distributionMethod.begin(),
543                               distributionOptions->distributionMethod.end());
544     SmallVector<Range, 2> parallelLoopRanges;
545     for (auto iteratorType : enumerate(iteratorTypes)) {
546       if (isParallelIterator(iteratorType.value()))
547         parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
548     }
549     if (distributionMethod.size() < parallelLoopRanges.size())
550       parallelLoopRanges.resize(distributionMethod.size());
551     SmallVector<ProcInfo, 2> procInfo =
552         options.procInfo(b, loc, parallelLoopRanges);
553     unsigned index = 0;
554     for (auto iteratorType : enumerate(iteratorTypes)) {
555       if (index >= procInfo.size())
556         break;
557       if (isParallelIterator(iteratorType.value())) {
558         unsigned i = iteratorType.index();
559         updateBoundsForCyclicDistribution(b, loc, procInfo[index].procId,
560                                           procInfo[index].nprocs, lbsStorage[i],
561                                           ubsStorage[i], stepsStorage[i]);
562         index++;
563       }
564     }
565   }
566   ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
567   generateParallelLoopNest(
568       b, loc, lbs, ubs, steps, iteratorTypes,
569       [&](OpBuilder &b, Location loc, ValueRange ivs) {
570         SmallVector<Value> operandValuesToUse =
571             linalgOp.getInputAndOutputOperands();
572         bodyBuilderFn(b, loc, ivs, operandValuesToUse);
573       },
574       ivs, distributionMethod);
575 
576   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
577 }
578 
579 static Value fullyComposeAndAffineApply(OpBuilder &b, Location loc,
580                                         AffineExpr expr, ValueRange operands) {
581   AffineMap map = AffineMap::inferFromExprList({expr}).front();
582   SmallVector<Value> normalizedOperands(operands.begin(), operands.end());
583   mlir::fullyComposeAffineMapAndOperands(&map, &normalizedOperands);
584   canonicalizeMapAndOperands(&map, &normalizedOperands);
585   return b.createOrFold<AffineApplyOp>(loc, map, normalizedOperands);
586 }
587 
588 Value makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile,
589                      ValueRange tileSizes, AffineMap map, ValueRange lbs,
590                      ValueRange ubs, ValueRange subShapeSizes) {
591   auto shapedType = valueToTile.getType().dyn_cast<ShapedType>();
592   assert(shapedType && "only shaped types can be tiled");
593   ArrayRef<int64_t> shape = shapedType.getShape();
594   int64_t rank = shapedType.getRank();
595 
596   // Construct a new subview / extract_slice for the tile.
597   SmallVector<OpFoldResult, 4> offsets, sizes, strides;
598   offsets.reserve(rank);
599   sizes.reserve(rank);
600   strides.reserve(rank);
601   for (unsigned r = 0; r < rank; ++r) {
602     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: for dim#" << r);
603     if (!isTiled(map.getSubMap({r}), tileSizes)) {
604       offsets.push_back(builder.getIndexAttr(0));
605       Value dim = createOrFoldDimOp(builder, loc, valueToTile, r);
606       sizes.push_back(getAsOpFoldResult(dim));
607       strides.push_back(builder.getIndexAttr(1));
608       LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n");
609       continue;
610     }
611     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");
612 
613     // Tiling creates a new slice at the proper index, the slice step is 1
614     // (i.e. the op does not subsample, stepping occurs in the loop).
615     auto m = map.getSubMap({r});
616     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: submap: " << m << "\n");
617     auto offset = applyMapToValues(builder, loc, m, lbs).front();
618     offsets.push_back(offset);
619     auto closedIntSize =
620         applyMapToValues(builder, loc, m, subShapeSizes).front();
621     // Resulting size needs to be made half open interval again.
622     AffineExpr s0 = getAffineSymbolExpr(0, builder.getContext());
623     Value size =
624         fullyComposeAndAffineApply(builder, loc, s0 + 1, closedIntSize);
625     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: raw size: " << size << "\n");
626 
627     // The size of the subview / extract_slice should be trimmed to avoid
628     // out-of-bounds accesses, unless:
629     // a. We statically know the subshape size divides the shape size evenly.
630     // b. The subshape size is 1. According to the way the loops are set up,
631     //    tensors with "0" dimensions would never be constructed.
632     int64_t shapeSize = shape[r];
633     auto sizeCst = size.getDefiningOp<arith::ConstantIndexOp>();
634     auto hasTileSizeOne = sizeCst && sizeCst.value() == 1;
635     auto dividesEvenly = sizeCst && !ShapedType::isDynamic(shapeSize) &&
636                          ((shapeSize % sizeCst.value()) == 0);
637     if (!hasTileSizeOne && !dividesEvenly) {
638       LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: shapeSize=" << shapeSize
639                               << ", size: " << size
640                               << ": make sure in bound with affine.min\n");
641 
642       AffineExpr dim0, dim1, dim2;
643       bindDims(builder.getContext(), dim0, dim1, dim2);
644 
645       // Get the dimension size for this dimension. We need to first calculate
646       // the max index and then plus one. This is important because for
647       // convolution ops, we have its input window dimension's affine map of the
648       // form `(d0 * s0 + d1)`, where `d0`/`d1 is an output/filter window
649       // dimension and `s0` is stride. Directly use the dimension size of
650       // output/filer window dimensions will cause incorrect calculation.
651       AffineMap minusOneMap =
652           AffineMap::inferFromExprList({ArrayRef<AffineExpr>{dim0 - 1}})
653               .front();
654       AffineMap plusOneMap =
655           AffineMap::inferFromExprList({ArrayRef<AffineExpr>{dim0 + 1}})
656               .front();
657       auto maxIndices = llvm::to_vector<8>(llvm::map_range(ubs, [&](Value ub) {
658         return makeComposedAffineApply(builder, loc, minusOneMap, {ub})
659             .getResult();
660       }));
661       Value maxIndex = applyMapToValues(builder, loc, m, maxIndices).front();
662       Value d = makeComposedAffineApply(builder, loc, plusOneMap, {maxIndex});
663 
664       // Compute min(size, dim - offset) to avoid out-of-bounds accesses.
665       AffineMap minMap = AffineMap::inferFromExprList(
666                              {ArrayRef<AffineExpr>{dim0, dim1 - dim2}})
667                              .front();
668       SmallVector<Value, 4> operands{size, d, offset};
669       fullyComposeAffineMapAndOperands(&minMap, &operands);
670       canonicalizeMapAndOperands(&minMap, &operands);
671       size = builder.create<AffineMinOp>(loc, builder.getIndexType(), minMap,
672                                          operands);
673     }
674 
675     sizes.push_back(size);
676     LLVM_DEBUG(llvm::dbgs()
677                << "makeTiledShape: new offset: " << offset << "\n");
678     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: new size: " << size << "\n");
679     strides.push_back(builder.getIndexAttr(1));
680   }
681 
682   auto *sliceOp = TypeSwitch<ShapedType, Operation *>(shapedType)
683                       .Case([&](MemRefType) {
684                         return builder.create<memref::SubViewOp>(
685                             loc, valueToTile, offsets, sizes, strides);
686                       })
687                       .Case([&](RankedTensorType) {
688                         return makeComposedExtractSliceOp(
689                             builder, loc, valueToTile, offsets, sizes, strides);
690                       })
691                       .Default([](ShapedType) -> Operation * {
692                         llvm_unreachable("Unexpected shaped type");
693                       });
694   return sliceOp->getResult(0);
695 }
696 
697 SmallVector<Value> computeTileOffsets(OpBuilder &b, Location loc,
698                                       ValueRange ivs, ValueRange tileSizes) {
699   SmallVector<Value> offsets;
700   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
701     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n");
702     bool isTiled = !isZero(tileSizes[idx]);
703     offsets.push_back(
704         isTiled ? ivs[idxIvs++]
705                 : b.create<arith::ConstantIndexOp>(loc, 0).getResult());
706     LLVM_DEBUG(llvm::dbgs()
707                << "computeTileOffsets: " << offsets.back() << "\n");
708   }
709   return offsets;
710 }
711 
712 SmallVector<Value> computeTileSizes(OpBuilder &b, Location loc, ValueRange ivs,
713                                     ValueRange tileSizes,
714                                     ArrayRef<Value> sizeBounds) {
715   SmallVector<Value> sizes;
716   for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx) {
717     bool isTiled = !isZero(tileSizes[idx]);
718     // Before composing, we need to make range a closed interval.
719     Value size = isTiled ? tileSizes[idx] : sizeBounds[idx];
720     AffineExpr d0 = getAffineDimExpr(0, b.getContext());
721     sizes.push_back(fullyComposeAndAffineApply(b, loc, d0 - 1, size));
722     LLVM_DEBUG(llvm::dbgs() << "computeTileSizes: " << sizes.back() << "\n");
723   }
724   return sizes;
725 }
726 
727 SmallVector<Value, 4> makeTiledShapes(OpBuilder &b, Location loc,
728                                       LinalgOp linalgOp,
729                                       ArrayRef<Value> valuesToTile,
730                                       ValueRange ivs, ValueRange tileSizes,
731                                       ArrayRef<Value> sizeBounds) {
732   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
733                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
734                            [](Value v) { return !isZero(v); })) &&
735          "expected as many ivs as non-zero sizes");
736 
737   // Construct (potentially temporary) mins and maxes on which to apply maps
738   // that define tile subshapes.
739   SmallVector<Value> lbs = computeTileOffsets(b, loc, ivs, tileSizes);
740   SmallVector<Value> subShapeSizes =
741       computeTileSizes(b, loc, ivs, tileSizes, sizeBounds);
742 
743   assert(static_cast<int64_t>(valuesToTile.size()) ==
744              linalgOp.getNumInputsAndOutputs() &&
745          "expected one value to tile for every operand");
746   SmallVector<Value, 4> tiledShapes;
747   tiledShapes.reserve(valuesToTile.size());
748   for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
749     Value shapedOp = valuesToTile[opOperand->getOperandNumber()];
750     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
751     AffineMap map = linalgOp.getTiedIndexingMap(opOperand);
752     // If the shape is not tiled, we can use it as is.
753     if (!isTiled(map, tileSizes)) {
754       tiledShapes.push_back(shapedOp);
755       LLVM_DEBUG(llvm::dbgs() << ": not tiled: use shape: "
756                               << opOperand->get().getType() << "\n");
757       continue;
758     }
759     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subshape...\n");
760 
761     tiledShapes.push_back(makeTiledShape(b, loc, shapedOp, tileSizes, map, lbs,
762                                          sizeBounds, subShapeSizes));
763   }
764 
765   return tiledShapes;
766 }
767 
768 void addTileLoopIvsToIndexOpResults(OpBuilder &b, LinalgOp tiledOp,
769                                     ArrayRef<Value> ivs) {
770   if (tiledOp.hasIndexSemantics()) {
771     for (IndexOp indexOp : tiledOp.getBlock()->getOps<IndexOp>()) {
772       if (ivs[indexOp.dim()] == nullptr)
773         continue;
774       OpBuilder::InsertionGuard guard(b);
775       b.setInsertionPointAfter(indexOp);
776       AffineExpr index, offset;
777       bindDims(b.getContext(), index, offset);
778       AffineApplyOp applyOp = makeComposedAffineApply(
779           b, indexOp.getLoc(), index + offset,
780           ValueRange{indexOp.getResult(), ivs[indexOp.dim()]});
781       indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp);
782     }
783   }
784 }
785 
786 } // namespace linalg
787 } // namespace mlir
788