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