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   SmallVector<ProcInfo, 2> procInfo;
209   if (distributionOptions.hasValue())
210     procInfo = distributionOptions->procInfo(builder, loc, loopRanges);
211 
212   SmallVector<Value, 4> lbs, ubs, steps;
213   unpackRanges(loopRanges, lbs, ubs, steps);
214   LoopNest loopNest =
215       edsc::loopNestBuilder(lbs, ubs, steps, iterArgInitValues, bodyBuilderFn);
216 
217   if (!distributionOptions.hasValue() || loopNest.loops.empty())
218     return;
219 
220   // Only supports cyclic distribution for now.
221   for (auto it : llvm::zip(loopNest.loops, procInfo,
222                            distributionOptions->distributionMethod))
223     if (std::get<2>(it) == DistributionMethod::Cyclic)
224       mapLoopToProcessorIds(std::get<0>(it), std::get<1>(it).procId,
225                             std::get<1>(it).nprocs);
226 }
227 
228 /// Specialization to build affine "for" nest.
229 template <>
230 void GenerateLoopNest<AffineForOp>::doit(
231     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
232     ArrayRef<Attribute> iteratorTypes,
233     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
234     Optional<LinalgLoopDistributionOptions>) {
235   auto iterArgInitValues = linalgOp.getOutputTensors();
236   assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
237   SmallVector<Value, 4> lbs, ubs, steps;
238   unpackRanges(loopRanges, lbs, ubs, steps);
239 
240   // Affine loops require constant steps.
241   SmallVector<int64_t, 4> constantSteps;
242   constantSteps.reserve(steps.size());
243   for (Value v : steps) {
244     auto op = v.getDefiningOp<ConstantIndexOp>();
245     assert(op && "Affine loops require constant steps");
246     constantSteps.push_back(op.getValue());
247   }
248 
249   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
250     bodyBuilderFn(ivs, {});
251   };
252   edsc::affineLoopNestBuilder(lbs, ubs, constantSteps,
253                               bodyBuilderWithoutIterArgsFn);
254 }
255 
256 /// Specialization to build an linalg.tiled_loop
257 template <>
258 void GenerateLoopNest<TiledLoopOp>::doit(
259     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
260     ArrayRef<Attribute> iteratorTypes,
261     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
262     Optional<LinalgLoopDistributionOptions>) {
263   OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
264   Location loc = edsc::ScopedContext::getLocation();
265   SmallVector<ProcInfo, 2> procInfo;
266 
267   SmallVector<Value, 4> lbs, ubs, steps;
268   unpackRanges(loopRanges, lbs, ubs, steps);
269 
270   auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
271                               ValueRange ivs, ValueRange inputs,
272                               ValueRange outputs) {
273     ScopedContext context(nestedBuilder, nestedLoc);
274     scf::ValueVector results = bodyBuilderFn(ivs, linalgOp.getOutputTensors());
275     nestedBuilder.create<linalg::YieldOp>(nestedLoc, results);
276   };
277 
278   auto tiledLoop = builder.create<TiledLoopOp>(
279       loc, lbs, ubs, steps, linalgOp.getInputs(), linalgOp.getOutputs(),
280       builder.getArrayAttr(iteratorTypes), wrappedBuilderFn);
281 
282   // Replace inputs/outputs with the corresponding region args.
283   auto isInsideTiledLoop = [&](OpOperand &operand) {
284     return operand.getOwner()->getBlock() == tiledLoop.getBody();
285   };
286   for (auto it :
287        llvm::zip(linalgOp.getInputs(), tiledLoop.getRegionInputArgs()))
288     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop);
289   for (auto it :
290        llvm::zip(linalgOp.getOutputs(), tiledLoop.getRegionOutputArgs()))
291     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop);
292 }
293 
294 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
295 void updateBoundsForCyclicDistribution(OpBuilder &builder, Location loc,
296                                        Value procId, Value nprocs, Value &lb,
297                                        Value &ub, Value &step) {
298   using edsc::op::operator+;
299   using edsc::op::operator*;
300   lb = lb + (procId * step);
301   step = nprocs * step;
302 }
303 
304 /// Generates a loop nest consisting of scf.parallel and scf.for, depending
305 /// on the `iteratorTypes.` Consecutive parallel loops create a single
306 /// scf.parallel operation; each sequential loop creates a new scf.for
307 /// operation. The body of the innermost loop is populated by
308 /// `bodyBuilderFn` that accepts a range of induction variables for all
309 /// loops. `ivStorage` is used to store the partial list of induction
310 /// variables.
311 // TODO: this function can be made iterative instead. However, it
312 // will have at most as many recursive calls as nested loops, which rarely
313 // exceeds 10.
314 static void
315 generateParallelLoopNest(ValueRange lbs, ValueRange ubs, ValueRange steps,
316                          ArrayRef<Attribute> iteratorTypes,
317                          function_ref<void(ValueRange)> bodyBuilderFn,
318                          SmallVectorImpl<Value> &ivStorage,
319                          ArrayRef<DistributionMethod> distributionMethod = {}) {
320   assert(lbs.size() == ubs.size());
321   assert(lbs.size() == steps.size());
322   assert(lbs.size() == iteratorTypes.size());
323 
324   // If there are no (more) loops to be generated, generate the body and be
325   // done with it.
326   if (iteratorTypes.empty())
327     return bodyBuilderFn(ivStorage);
328 
329   // Find the outermost parallel loops and drop their types from the list.
330   unsigned nLoops = iteratorTypes.size();
331   unsigned nOuterPar =
332       nLoops - iteratorTypes.drop_while(isParallelIteratorType).size();
333 
334   // If there are no outer parallel loops, generate one sequential loop and
335   // recurse. Note that we wouldn't have dropped anything from `iteratorTypes`
336   // in this case.
337   if (nOuterPar == 0) {
338     edsc::loopNestBuilder(lbs[0], ubs[0], steps[0], [&](Value iv) {
339       ivStorage.push_back(iv);
340       generateParallelLoopNest(lbs.drop_front(), ubs.drop_front(),
341                                steps.drop_front(), iteratorTypes.drop_front(),
342                                bodyBuilderFn, ivStorage, distributionMethod);
343     });
344     return;
345   }
346   if (distributionMethod.empty()) {
347     // Generate a single parallel loop-nest operation for all outermost
348     // parallel loops and recurse.
349     edsc::OperationBuilder<scf::ParallelOp>(
350         lbs.take_front(nOuterPar), ubs.take_front(nOuterPar),
351         steps.take_front(nOuterPar),
352         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
353           edsc::ScopedContext context(nestedBuilder, nestedLoc);
354           ivStorage.append(localIvs.begin(), localIvs.end());
355           generateParallelLoopNest(
356               lbs.drop_front(nOuterPar), ubs.drop_front(nOuterPar),
357               steps.drop_front(nOuterPar), iteratorTypes.drop_front(nOuterPar),
358               bodyBuilderFn, ivStorage,
359               (distributionMethod.size() < nOuterPar)
360                   ? ArrayRef<DistributionMethod>()
361                   : distributionMethod.drop_front(nOuterPar));
362         });
363     return;
364   }
365 
366   // Process all consecutive similarly distributed loops simultaneously.
367   DistributionMethod methodToUse = distributionMethod[0];
368   unsigned numProcessed = 1;
369   for (unsigned i = 1; i < nOuterPar && i < distributionMethod.size(); ++i) {
370     if (distributionMethod[i] != methodToUse)
371       break;
372     numProcessed++;
373   }
374 
375   switch (methodToUse) {
376   case DistributionMethod::Cyclic: {
377     // Generate a single parallel loop-nest operation for all outermost
378     // parallel loops and recurse.
379     edsc::OperationBuilder<scf::ParallelOp>(
380         lbs.take_front(numProcessed), ubs.take_front(numProcessed),
381         steps.take_front(numProcessed),
382         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
383           edsc::ScopedContext context(nestedBuilder, nestedLoc);
384           ivStorage.append(localIvs.begin(), localIvs.end());
385           generateParallelLoopNest(
386               lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
387               steps.drop_front(numProcessed),
388               iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
389               (distributionMethod.size() < numProcessed)
390                   ? ArrayRef<DistributionMethod>()
391                   : distributionMethod.drop_front(numProcessed));
392         });
393     return;
394   }
395   case DistributionMethod::CyclicNumProcsGeNumIters: {
396     // Check (for the processed loops) that the iteration is in-bounds.
397     using edsc::op::slt;
398     using edsc::op::operator&&;
399     Value cond = slt(lbs[0], ubs[0]);
400     for (unsigned i = 1; i < numProcessed; ++i)
401       cond = cond && slt(lbs[i], ubs[i]);
402     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
403     edsc::conditionBuilder(cond, [&]() {
404       generateParallelLoopNest(
405           lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
406           steps.drop_front(numProcessed),
407           iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
408           distributionMethod.drop_front(numProcessed));
409     });
410     return;
411   }
412   case DistributionMethod::CyclicNumProcsEqNumIters:
413     // No check/loops needed here. Set the `%iv` to be the `%lb` and proceed
414     // with inner loop generation.
415     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
416     generateParallelLoopNest(
417         lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
418         steps.drop_front(numProcessed), iteratorTypes.drop_front(numProcessed),
419         bodyBuilderFn, ivStorage, distributionMethod.drop_front(numProcessed));
420     return;
421   }
422 }
423 
424 /// Specialization for generating a mix of parallel and sequential scf loops.
425 template <>
426 void GenerateLoopNest<scf::ParallelOp>::doit(
427     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
428     ArrayRef<Attribute> iteratorTypes,
429     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
430     Optional<LinalgLoopDistributionOptions> distributionOptions) {
431   auto iterArgInitValues = linalgOp.getOutputTensors();
432   assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");
433   // This function may be passed more iterator types than ranges.
434   assert(iteratorTypes.size() >= loopRanges.size() &&
435          "expected iterator type for all ranges");
436   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
437   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
438   unsigned numLoops = iteratorTypes.size();
439   ivs.reserve(numLoops);
440   lbsStorage.reserve(numLoops);
441   ubsStorage.reserve(numLoops);
442   stepsStorage.reserve(numLoops);
443 
444   // Get the loop lb, ub, and step.
445   unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
446 
447   // Modify the lb, ub, and step based on the distribution options.
448   SmallVector<DistributionMethod, 0> distributionMethod;
449   if (distributionOptions) {
450     auto &options = distributionOptions.getValue();
451     OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
452     Location loc = edsc::ScopedContext::getLocation();
453     distributionMethod.assign(distributionOptions->distributionMethod.begin(),
454                               distributionOptions->distributionMethod.end());
455     SmallVector<Range, 2> parallelLoopRanges;
456     for (auto iteratorType : enumerate(iteratorTypes)) {
457       if (isParallelIteratorType(iteratorType.value()))
458         parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
459     }
460     if (distributionMethod.size() < parallelLoopRanges.size())
461       parallelLoopRanges.resize(distributionMethod.size());
462     SmallVector<ProcInfo, 2> procInfo =
463         options.procInfo(builder, loc, parallelLoopRanges);
464     unsigned index = 0;
465     for (auto iteratorType : enumerate(iteratorTypes)) {
466       if (index >= procInfo.size())
467         break;
468       if (isParallelIteratorType(iteratorType.value())) {
469         unsigned i = iteratorType.index();
470         updateBoundsForCyclicDistribution(builder, loc, procInfo[index].procId,
471                                           procInfo[index].nprocs, lbsStorage[i],
472                                           ubsStorage[i], stepsStorage[i]);
473         index++;
474       }
475     }
476   }
477   ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
478   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
479     bodyBuilderFn(ivs, {});
480   };
481   generateParallelLoopNest(lbs, ubs, steps, iteratorTypes,
482                            bodyBuilderWithoutIterArgsFn, ivs,
483                            distributionMethod);
484 
485   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
486 }
487 
488 SmallVector<Value, 4> makeTiledShapes(OpBuilder &builder, Location loc,
489                                       LinalgOp linalgOp,
490                                       ArrayRef<Value> tiledOperands,
491                                       ValueRange ivs, ValueRange tileSizes,
492                                       ArrayRef<Value> sizeBounds) {
493   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
494                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
495                            [](Value v) { return !isZero(v); })) &&
496          "expected as many ivs as non-zero sizes");
497 
498   using namespace edsc::op;
499 
500   // Construct (potentially temporary) mins and maxes on which to apply maps
501   // that define tile subshapes.
502   SmallVector<Value, 8> lbs, subShapeSizes;
503   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
504     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n");
505     bool isTiled = !isZero(tileSizes[idx]);
506     lbs.push_back(isTiled ? ivs[idxIvs++] : (Value)std_constant_index(0));
507     // Before composing, we need to make range a closed interval.
508     Value size = isTiled ? tileSizes[idx] : sizeBounds[idx];
509     subShapeSizes.push_back(size - std_constant_index(1));
510     LLVM_DEBUG(llvm::dbgs() << "lb: " << lbs.back() << "\n");
511     LLVM_DEBUG(llvm::dbgs() << "size: " << subShapeSizes.back() << "\n");
512   }
513 
514   MLIRContext *context = builder.getContext();
515   SmallVector<Value, 4> tiledShapes;
516   tiledShapes.reserve(tiledOperands.size());
517   for (auto en : llvm::enumerate(tiledOperands)) {
518     Value shapedOp = en.value();
519     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
520     ShapedType shapedType = shapedOp.getType().cast<ShapedType>();
521     unsigned rank = shapedType.getRank();
522     AffineMap map = linalgOp.getIndexingMap(en.index());
523     // If the shape is not tiled, we can use it as is.
524     if (!isTiled(map, tileSizes)) {
525       tiledShapes.push_back(shapedOp);
526       LLVM_DEBUG(llvm::dbgs()
527                  << ": not tiled: use shape: " << shapedType << "\n");
528       continue;
529     }
530     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subshape...\n");
531 
532     // Construct a new subview / subtensor for the tile.
533     SmallVector<OpFoldResult, 4> offsets, sizes, strides;
534     offsets.reserve(rank);
535     sizes.reserve(rank);
536     strides.reserve(rank);
537     for (unsigned r = 0; r < rank; ++r) {
538       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for dim#" << r);
539       if (!isTiled(map.getSubMap({r}), tileSizes)) {
540         offsets.push_back(builder.getIndexAttr(0));
541         Value dim = memref_dim(shapedOp, r).value;
542         sizes.push_back(dim);
543         strides.push_back(builder.getIndexAttr(1));
544         LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n");
545         continue;
546       }
547       LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");
548 
549       // Tiling creates a new slice at the proper index, the slice step is 1
550       // (i.e. the op does not subsample, stepping occurs in the loop).
551       auto m = map.getSubMap({r});
552       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: submap: " << map << "\n");
553       auto offset = applyMapToValues(builder, loc, m, lbs).front();
554       offsets.push_back(offset);
555       auto closedIntSize =
556           applyMapToValues(builder, loc, m, subShapeSizes).front();
557       // Resulting size needs to be made half open interval again.
558       auto size = closedIntSize + std_constant_index(1);
559       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: raw size: " << size << "\n");
560 
561       // The size of the subview / subtensor should be trimmed to avoid
562       // out-of-bounds accesses, unless we statically know the subshape size
563       // divides the shape size evenly.
564       int64_t shapeSize = shapedType.getDimSize(r);
565       auto sizeCst = size.getDefiningOp<ConstantIndexOp>();
566       if (ShapedType::isDynamic(shapeSize) || !sizeCst ||
567           (shapeSize % sizeCst.getValue()) != 0) {
568         LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: shapeSize=" << shapeSize
569                                 << ", size: " << size
570                                 << ": make sure in bound with affine.min\n");
571         AffineExpr dim0, dim1, dim2;
572         bindDims(context, dim0, dim1, dim2);
573         // Compute min(size, dim - offset) to avoid out-of-bounds accesses.
574         auto minMap = AffineMap::get(
575             /*dimCount=*/3, /*symbolCount=*/0, {dim0, dim1 - dim2}, context);
576         Value d = memref_dim(shapedOp, r);
577         SmallVector<Value, 4> operands{size, d, offset};
578         fullyComposeAffineMapAndOperands(&minMap, &operands);
579         size = affine_min(builder.getIndexType(), minMap, operands);
580       }
581 
582       sizes.push_back(size);
583       LLVM_DEBUG(llvm::dbgs()
584                  << "makeTiledShapes: new offset: " << offset << "\n");
585       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: new size: " << size << "\n");
586       strides.push_back(builder.getIndexAttr(1));
587     }
588 
589     if (shapedType.isa<MemRefType>())
590       tiledShapes.push_back(builder.create<memref::SubViewOp>(
591           loc, shapedOp, offsets, sizes, strides));
592     else
593       tiledShapes.push_back(
594           builder.create<SubTensorOp>(loc, shapedOp, offsets, sizes, strides));
595   }
596 
597   return tiledShapes;
598 }
599 
600 } // namespace linalg
601 } // namespace mlir
602