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