1 //===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements utilities for the Linalg dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/Linalg/Utils/Utils.h"
14 
15 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
16 #include "mlir/Dialect/Affine/IR/AffineOps.h"
17 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
18 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
19 #include "mlir/Dialect/SCF/EDSC/Builders.h"
20 #include "mlir/Dialect/SCF/SCF.h"
21 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
22 #include "mlir/Dialect/StandardOps/IR/Ops.h"
23 #include "mlir/IR/AffineExpr.h"
24 #include "mlir/IR/AffineExprVisitor.h"
25 #include "mlir/IR/AffineMap.h"
26 #include "mlir/IR/Matchers.h"
27 #include "mlir/IR/OpImplementation.h"
28 #include "mlir/Pass/Pass.h"
29 #include "mlir/Transforms/LoopUtils.h"
30 #include "llvm/Support/Debug.h"
31 
32 #define DEBUG_TYPE "linalg-utils"
33 
34 using namespace mlir;
35 using namespace mlir::edsc;
36 using namespace mlir::edsc::intrinsics;
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 bool mlir::linalg::isParallelIteratorType(Attribute attr) {
121   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
122     return strAttr.getValue() == getParallelIteratorTypeName();
123   }
124   return false;
125 }
126 
127 bool mlir::linalg::isReductionIteratorType(Attribute attr) {
128   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
129     return strAttr.getValue() == getReductionIteratorTypeName();
130   }
131   return false;
132 }
133 
134 bool mlir::linalg::isWindowIteratorType(Attribute attr) {
135   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
136     return strAttr.getValue() == getWindowIteratorTypeName();
137   }
138   return false;
139 }
140 
141 /// Explicit instantiation of loop nest generator for different loop types.
142 template struct mlir::linalg::GenerateLoopNest<scf::ForOp>;
143 template struct mlir::linalg::GenerateLoopNest<scf::ParallelOp>;
144 template struct mlir::linalg::GenerateLoopNest<AffineForOp>;
145 template struct mlir::linalg::GenerateLoopNest<TiledLoopOp>;
146 
147 /// Given a list of subview ranges, extract individual values for lower, upper
148 /// bounds and steps and put them into the corresponding vectors.
149 static void unpackRanges(ArrayRef<Range> ranges, SmallVectorImpl<Value> &lbs,
150                          SmallVectorImpl<Value> &ubs,
151                          SmallVectorImpl<Value> &steps) {
152   for (Range range : ranges) {
153     lbs.emplace_back(range.offset);
154     ubs.emplace_back(range.size);
155     steps.emplace_back(range.stride);
156   }
157 }
158 
159 namespace mlir {
160 namespace linalg {
161 
162 /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp
163 /// is a constant then return a new value set to the smallest such constant.
164 /// Otherwise returngetSmallestBoundingIndex nullptr.
165 IntegerAttr getSmallestBoundingIndex(Value size) {
166   Optional<int64_t> boundingConst = {};
167   if (auto affineMinOp = size.getDefiningOp<AffineMinOp>()) {
168     for (auto e : affineMinOp.getAffineMap().getResults())
169       if (auto cst = e.dyn_cast<AffineConstantExpr>())
170         boundingConst = boundingConst
171                             ? std::min(boundingConst.getValue(), cst.getValue())
172                             : cst.getValue();
173   } else if (auto constIndexOp = size.getDefiningOp<ConstantOp>()) {
174     if (constIndexOp.getType().isa<IndexType>())
175       boundingConst = constIndexOp.value().cast<IntegerAttr>().getInt();
176   } else if (auto affineApplyOp = size.getDefiningOp<AffineApplyOp>()) {
177     if (auto cExpr = affineApplyOp.getAffineMap()
178                          .getResult(0)
179                          .dyn_cast<AffineConstantExpr>())
180       boundingConst = cExpr.getValue();
181   } else if (auto dimOp = size.getDefiningOp<memref::DimOp>()) {
182     auto shape = dimOp.memrefOrTensor().getType().dyn_cast<ShapedType>();
183     if (auto constOp = dimOp.index().getDefiningOp<ConstantOp>()) {
184       if (auto indexAttr = constOp.value().dyn_cast<IntegerAttr>()) {
185         auto dimIndex = indexAttr.getInt();
186         if (!shape.isDynamicDim(dimIndex)) {
187           boundingConst = shape.getShape()[dimIndex];
188         }
189       }
190     }
191   }
192   if (boundingConst && *boundingConst >= 0)
193     return Builder(size.getContext()).getIndexAttr(*boundingConst);
194   return nullptr;
195 }
196 
197 /// Specialization to build an scf "for" nest.
198 template <>
199 void GenerateLoopNest<scf::ForOp>::doit(
200     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
201     ArrayRef<Attribute> iteratorTypes,
202     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
203     Optional<LinalgLoopDistributionOptions> distributionOptions) {
204   auto iterArgInitValues = linalgOp.getOutputTensors();
205   // Create procInfo so it dominates loops, if appropriate.
206   OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
207   Location loc = edsc::ScopedContext::getLocation();
208 
209   SmallVector<ProcInfo, 4> procInfo;
210   SmallVector<DistributionMethod, 0> distributionMethod;
211   if (distributionOptions.hasValue()) {
212     // Collect loop ranges for parallel dimensions.
213     SmallVector<Range, 2> parallelLoopRanges;
214     for (auto iteratorType : enumerate(iteratorTypes))
215       if (isParallelIteratorType(iteratorType.value()))
216         parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
217 
218     // Get their distribution schemes.
219     distributionMethod = distributionOptions->distributionMethod;
220     if (distributionMethod.size() < parallelLoopRanges.size())
221       parallelLoopRanges.resize(distributionMethod.size());
222     procInfo = distributionOptions->procInfo(builder, loc, parallelLoopRanges);
223   }
224 
225   SmallVector<Value, 4> lbs, ubs, steps;
226   unpackRanges(loopRanges, lbs, ubs, steps);
227   LoopNest loopNest =
228       edsc::loopNestBuilder(lbs, ubs, steps, iterArgInitValues, bodyBuilderFn);
229 
230   if (!distributionOptions || loopNest.loops.empty())
231     return;
232 
233   // Filter out scf.for loops that were created out of parallel dimensions.
234   SmallVector<scf::ForOp, 4> loops;
235   for (auto iteratorType : enumerate(iteratorTypes))
236     if (isParallelIteratorType(iteratorType.value()))
237       loops.push_back(loopNest.loops[iteratorType.index()]);
238 
239   // Distribute - only supports cyclic distribution for now.
240   for (auto it : llvm::zip(loops, procInfo, distributionMethod))
241     if (std::get<2>(it) == DistributionMethod::Cyclic)
242       mapLoopToProcessorIds(std::get<0>(it), std::get<1>(it).procId,
243                             std::get<1>(it).nprocs);
244 }
245 
246 /// Specialization to build affine "for" nest.
247 template <>
248 void GenerateLoopNest<AffineForOp>::doit(
249     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
250     ArrayRef<Attribute> iteratorTypes,
251     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
252     Optional<LinalgLoopDistributionOptions>) {
253   auto iterArgInitValues = linalgOp.getOutputTensors();
254   assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
255   SmallVector<Value, 4> lbs, ubs, steps;
256   unpackRanges(loopRanges, lbs, ubs, steps);
257 
258   // Affine loops require constant steps.
259   SmallVector<int64_t, 4> constantSteps;
260   constantSteps.reserve(steps.size());
261   for (Value v : steps) {
262     auto op = v.getDefiningOp<ConstantIndexOp>();
263     assert(op && "Affine loops require constant steps");
264     constantSteps.push_back(op.getValue());
265   }
266 
267   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
268     bodyBuilderFn(ivs, {});
269   };
270   edsc::affineLoopNestBuilder(lbs, ubs, constantSteps,
271                               bodyBuilderWithoutIterArgsFn);
272 }
273 
274 /// Specialization to build an linalg.tiled_loop
275 template <>
276 void GenerateLoopNest<TiledLoopOp>::doit(
277     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
278     ArrayRef<Attribute> iteratorTypes,
279     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
280     Optional<LinalgLoopDistributionOptions>) {
281   OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
282   Location loc = edsc::ScopedContext::getLocation();
283   SmallVector<ProcInfo, 2> procInfo;
284 
285   SmallVector<Value, 4> lbs, ubs, steps;
286   unpackRanges(loopRanges, lbs, ubs, steps);
287 
288   auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
289                               ValueRange ivs, ValueRange inputs,
290                               ValueRange outputs) {
291     ScopedContext context(nestedBuilder, nestedLoc);
292     scf::ValueVector results = bodyBuilderFn(ivs, linalgOp.getOutputTensors());
293     nestedBuilder.create<linalg::YieldOp>(nestedLoc, results);
294   };
295 
296   auto tiledLoop = builder.create<TiledLoopOp>(
297       loc, lbs, ubs, steps, linalgOp.getInputs(), linalgOp.getOutputs(),
298       builder.getArrayAttr(iteratorTypes), wrappedBuilderFn);
299 
300   // Replace inputs/outputs with the corresponding region args.
301   auto isInsideTiledLoop = [&](OpOperand &operand) {
302     return operand.getOwner()->getBlock() == tiledLoop.getBody();
303   };
304   for (auto it :
305        llvm::zip(linalgOp.getInputs(), tiledLoop.getRegionInputArgs()))
306     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop);
307   for (auto it :
308        llvm::zip(linalgOp.getOutputs(), tiledLoop.getRegionOutputArgs()))
309     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop);
310 }
311 
312 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
313 void updateBoundsForCyclicDistribution(OpBuilder &builder, Location loc,
314                                        Value procId, Value nprocs, Value &lb,
315                                        Value &ub, Value &step) {
316   using edsc::op::operator+;
317   using edsc::op::operator*;
318   lb = lb + (procId * step);
319   step = nprocs * step;
320 }
321 
322 /// Generates a loop nest consisting of scf.parallel and scf.for, depending
323 /// on the `iteratorTypes.` Consecutive parallel loops create a single
324 /// scf.parallel operation; each sequential loop creates a new scf.for
325 /// operation. The body of the innermost loop is populated by
326 /// `bodyBuilderFn` that accepts a range of induction variables for all
327 /// loops. `ivStorage` is used to store the partial list of induction
328 /// variables.
329 // TODO: this function can be made iterative instead. However, it
330 // will have at most as many recursive calls as nested loops, which rarely
331 // exceeds 10.
332 static void
333 generateParallelLoopNest(ValueRange lbs, ValueRange ubs, ValueRange steps,
334                          ArrayRef<Attribute> iteratorTypes,
335                          function_ref<void(ValueRange)> bodyBuilderFn,
336                          SmallVectorImpl<Value> &ivStorage,
337                          ArrayRef<DistributionMethod> distributionMethod = {}) {
338   assert(lbs.size() == ubs.size());
339   assert(lbs.size() == steps.size());
340   assert(lbs.size() == iteratorTypes.size());
341 
342   // If there are no (more) loops to be generated, generate the body and be
343   // done with it.
344   if (iteratorTypes.empty())
345     return bodyBuilderFn(ivStorage);
346 
347   // Find the outermost parallel loops and drop their types from the list.
348   unsigned nLoops = iteratorTypes.size();
349   unsigned nOuterPar =
350       nLoops - iteratorTypes.drop_while(isParallelIteratorType).size();
351 
352   // If there are no outer parallel loops, generate one sequential loop and
353   // recurse. Note that we wouldn't have dropped anything from `iteratorTypes`
354   // in this case.
355   if (nOuterPar == 0) {
356     edsc::loopNestBuilder(lbs[0], ubs[0], steps[0], [&](Value iv) {
357       ivStorage.push_back(iv);
358       generateParallelLoopNest(lbs.drop_front(), ubs.drop_front(),
359                                steps.drop_front(), iteratorTypes.drop_front(),
360                                bodyBuilderFn, ivStorage, distributionMethod);
361     });
362     return;
363   }
364   if (distributionMethod.empty()) {
365     // Generate a single parallel loop-nest operation for all outermost
366     // parallel loops and recurse.
367     edsc::OperationBuilder<scf::ParallelOp>(
368         lbs.take_front(nOuterPar), ubs.take_front(nOuterPar),
369         steps.take_front(nOuterPar),
370         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
371           edsc::ScopedContext context(nestedBuilder, nestedLoc);
372           ivStorage.append(localIvs.begin(), localIvs.end());
373           generateParallelLoopNest(
374               lbs.drop_front(nOuterPar), ubs.drop_front(nOuterPar),
375               steps.drop_front(nOuterPar), iteratorTypes.drop_front(nOuterPar),
376               bodyBuilderFn, ivStorage,
377               (distributionMethod.size() < nOuterPar)
378                   ? ArrayRef<DistributionMethod>()
379                   : distributionMethod.drop_front(nOuterPar));
380         });
381     return;
382   }
383 
384   // Process all consecutive similarly distributed loops simultaneously.
385   DistributionMethod methodToUse = distributionMethod[0];
386   unsigned numProcessed = 1;
387   for (unsigned i = 1; i < nOuterPar && i < distributionMethod.size(); ++i) {
388     if (distributionMethod[i] != methodToUse)
389       break;
390     numProcessed++;
391   }
392 
393   switch (methodToUse) {
394   case DistributionMethod::Cyclic: {
395     // Generate a single parallel loop-nest operation for all outermost
396     // parallel loops and recurse.
397     edsc::OperationBuilder<scf::ParallelOp>(
398         lbs.take_front(numProcessed), ubs.take_front(numProcessed),
399         steps.take_front(numProcessed),
400         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
401           edsc::ScopedContext context(nestedBuilder, nestedLoc);
402           ivStorage.append(localIvs.begin(), localIvs.end());
403           generateParallelLoopNest(
404               lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
405               steps.drop_front(numProcessed),
406               iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
407               (distributionMethod.size() < numProcessed)
408                   ? ArrayRef<DistributionMethod>()
409                   : distributionMethod.drop_front(numProcessed));
410         });
411     return;
412   }
413   case DistributionMethod::CyclicNumProcsGeNumIters: {
414     // Check (for the processed loops) that the iteration is in-bounds.
415     using edsc::op::slt;
416     using edsc::op::operator&&;
417     Value cond = slt(lbs[0], ubs[0]);
418     for (unsigned i = 1; i < numProcessed; ++i)
419       cond = cond && slt(lbs[i], ubs[i]);
420     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
421     edsc::conditionBuilder(cond, [&]() {
422       generateParallelLoopNest(
423           lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
424           steps.drop_front(numProcessed),
425           iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
426           distributionMethod.drop_front(numProcessed));
427     });
428     return;
429   }
430   case DistributionMethod::CyclicNumProcsEqNumIters:
431     // No check/loops needed here. Set the `%iv` to be the `%lb` and proceed
432     // with inner loop generation.
433     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
434     generateParallelLoopNest(
435         lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
436         steps.drop_front(numProcessed), iteratorTypes.drop_front(numProcessed),
437         bodyBuilderFn, ivStorage, distributionMethod.drop_front(numProcessed));
438     return;
439   }
440 }
441 
442 /// Specialization for generating a mix of parallel and sequential scf loops.
443 template <>
444 void GenerateLoopNest<scf::ParallelOp>::doit(
445     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
446     ArrayRef<Attribute> iteratorTypes,
447     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
448     Optional<LinalgLoopDistributionOptions> distributionOptions) {
449   auto iterArgInitValues = linalgOp.getOutputTensors();
450   assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");
451   // This function may be passed more iterator types than ranges.
452   assert(iteratorTypes.size() >= loopRanges.size() &&
453          "expected iterator type for all ranges");
454   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
455   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
456   unsigned numLoops = iteratorTypes.size();
457   ivs.reserve(numLoops);
458   lbsStorage.reserve(numLoops);
459   ubsStorage.reserve(numLoops);
460   stepsStorage.reserve(numLoops);
461 
462   // Get the loop lb, ub, and step.
463   unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
464 
465   // Modify the lb, ub, and step based on the distribution options.
466   SmallVector<DistributionMethod, 0> distributionMethod;
467   if (distributionOptions) {
468     auto &options = distributionOptions.getValue();
469     OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
470     Location loc = edsc::ScopedContext::getLocation();
471     distributionMethod.assign(distributionOptions->distributionMethod.begin(),
472                               distributionOptions->distributionMethod.end());
473     SmallVector<Range, 2> parallelLoopRanges;
474     for (auto iteratorType : enumerate(iteratorTypes)) {
475       if (isParallelIteratorType(iteratorType.value()))
476         parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
477     }
478     if (distributionMethod.size() < parallelLoopRanges.size())
479       parallelLoopRanges.resize(distributionMethod.size());
480     SmallVector<ProcInfo, 2> procInfo =
481         options.procInfo(builder, loc, parallelLoopRanges);
482     unsigned index = 0;
483     for (auto iteratorType : enumerate(iteratorTypes)) {
484       if (index >= procInfo.size())
485         break;
486       if (isParallelIteratorType(iteratorType.value())) {
487         unsigned i = iteratorType.index();
488         updateBoundsForCyclicDistribution(builder, loc, procInfo[index].procId,
489                                           procInfo[index].nprocs, lbsStorage[i],
490                                           ubsStorage[i], stepsStorage[i]);
491         index++;
492       }
493     }
494   }
495   ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
496   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
497     bodyBuilderFn(ivs, {});
498   };
499   generateParallelLoopNest(lbs, ubs, steps, iteratorTypes,
500                            bodyBuilderWithoutIterArgsFn, ivs,
501                            distributionMethod);
502 
503   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
504 }
505 
506 SmallVector<Value, 4> makeTiledShapes(OpBuilder &builder, Location loc,
507                                       LinalgOp linalgOp,
508                                       ArrayRef<Value> tiledOperands,
509                                       ValueRange ivs, ValueRange tileSizes,
510                                       ArrayRef<Value> sizeBounds) {
511   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
512                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
513                            [](Value v) { return !isZero(v); })) &&
514          "expected as many ivs as non-zero sizes");
515 
516   using namespace edsc::op;
517 
518   // Construct (potentially temporary) mins and maxes on which to apply maps
519   // that define tile subshapes.
520   SmallVector<Value, 8> lbs, subShapeSizes;
521   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
522     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n");
523     bool isTiled = !isZero(tileSizes[idx]);
524     lbs.push_back(isTiled ? ivs[idxIvs++] : (Value)std_constant_index(0));
525     // Before composing, we need to make range a closed interval.
526     Value size = isTiled ? tileSizes[idx] : sizeBounds[idx];
527     subShapeSizes.push_back(size - std_constant_index(1));
528     LLVM_DEBUG(llvm::dbgs() << "lb: " << lbs.back() << "\n");
529     LLVM_DEBUG(llvm::dbgs() << "size: " << subShapeSizes.back() << "\n");
530   }
531 
532   MLIRContext *context = builder.getContext();
533   SmallVector<Value, 4> tiledShapes;
534   tiledShapes.reserve(tiledOperands.size());
535   for (auto en : llvm::enumerate(tiledOperands)) {
536     Value shapedOp = en.value();
537     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
538     ShapedType shapedType = shapedOp.getType().cast<ShapedType>();
539     unsigned rank = shapedType.getRank();
540     AffineMap map = linalgOp.getIndexingMap(en.index());
541     // If the shape is not tiled, we can use it as is.
542     if (!isTiled(map, tileSizes)) {
543       tiledShapes.push_back(shapedOp);
544       LLVM_DEBUG(llvm::dbgs()
545                  << ": not tiled: use shape: " << shapedType << "\n");
546       continue;
547     }
548     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subshape...\n");
549 
550     // Construct a new subview / subtensor for the tile.
551     SmallVector<OpFoldResult, 4> offsets, sizes, strides;
552     offsets.reserve(rank);
553     sizes.reserve(rank);
554     strides.reserve(rank);
555     for (unsigned r = 0; r < rank; ++r) {
556       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for dim#" << r);
557       if (!isTiled(map.getSubMap({r}), tileSizes)) {
558         offsets.push_back(builder.getIndexAttr(0));
559         Value dim = memref_dim(shapedOp, r).value;
560         sizes.push_back(dim);
561         strides.push_back(builder.getIndexAttr(1));
562         LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n");
563         continue;
564       }
565       LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");
566 
567       // Tiling creates a new slice at the proper index, the slice step is 1
568       // (i.e. the op does not subsample, stepping occurs in the loop).
569       auto m = map.getSubMap({r});
570       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: submap: " << map << "\n");
571       auto offset = applyMapToValues(builder, loc, m, lbs).front();
572       offsets.push_back(offset);
573       auto closedIntSize =
574           applyMapToValues(builder, loc, m, subShapeSizes).front();
575       // Resulting size needs to be made half open interval again.
576       auto size = closedIntSize + std_constant_index(1);
577       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: raw size: " << size << "\n");
578 
579       // The size of the subview / subtensor should be trimmed to avoid
580       // out-of-bounds accesses, unless we statically know the subshape size
581       // divides the shape size evenly.
582       int64_t shapeSize = shapedType.getDimSize(r);
583       auto sizeCst = size.getDefiningOp<ConstantIndexOp>();
584       if (ShapedType::isDynamic(shapeSize) || !sizeCst ||
585           (shapeSize % sizeCst.getValue()) != 0) {
586         LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: shapeSize=" << shapeSize
587                                 << ", size: " << size
588                                 << ": make sure in bound with affine.min\n");
589         AffineExpr dim0, dim1, dim2;
590         bindDims(context, dim0, dim1, dim2);
591         // Compute min(size, dim - offset) to avoid out-of-bounds accesses.
592         auto minMap = AffineMap::get(
593             /*dimCount=*/3, /*symbolCount=*/0, {dim0, dim1 - dim2}, context);
594         Value d = memref_dim(shapedOp, r);
595         SmallVector<Value, 4> operands{size, d, offset};
596         fullyComposeAffineMapAndOperands(&minMap, &operands);
597         size = affine_min(builder.getIndexType(), minMap, operands);
598       }
599 
600       sizes.push_back(size);
601       LLVM_DEBUG(llvm::dbgs()
602                  << "makeTiledShapes: new offset: " << offset << "\n");
603       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: new size: " << size << "\n");
604       strides.push_back(builder.getIndexAttr(1));
605     }
606 
607     if (shapedType.isa<MemRefType>())
608       tiledShapes.push_back(builder.create<memref::SubViewOp>(
609           loc, shapedOp, offsets, sizes, strides));
610     else
611       tiledShapes.push_back(
612           builder.create<SubTensorOp>(loc, shapedOp, offsets, sizes, strides));
613   }
614 
615   return tiledShapes;
616 }
617 
618 } // namespace linalg
619 } // namespace mlir
620