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 #include "mlir/Dialect/Affine/IR/AffineOps.h"
15 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
16 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
17 #include "mlir/Dialect/SCF/EDSC/Builders.h"
18 #include "mlir/Dialect/SCF/SCF.h"
19 #include "mlir/Dialect/StandardOps/IR/Ops.h"
20 #include "mlir/IR/AffineExpr.h"
21 #include "mlir/IR/AffineMap.h"
22 #include "mlir/IR/Matchers.h"
23 #include "mlir/IR/OpImplementation.h"
24 #include "mlir/Pass/Pass.h"
25 #include "mlir/Transforms/FoldUtils.h"
26 
27 using namespace mlir;
28 using namespace mlir::linalg;
29 using namespace mlir::scf;
30 
31 Optional<RegionMatcher::BinaryOpKind>
32 RegionMatcher::matchAsScalarBinaryOp(GenericOp op) {
33   auto &region = op.region();
34   if (!llvm::hasSingleElement(region))
35     return llvm::None;
36 
37   Block &block = region.front();
38   if (block.getNumArguments() != 2 ||
39       !block.getArgument(0).getType().isSignlessIntOrFloat() ||
40       !block.getArgument(1).getType().isSignlessIntOrFloat())
41     return llvm::None;
42 
43   auto &ops = block.getOperations();
44   if (!llvm::hasSingleElement(block.without_terminator()))
45     return llvm::None;
46 
47   using mlir::matchers::m_Val;
48   auto a = m_Val(block.getArgument(0));
49   auto b = m_Val(block.getArgument(1));
50 
51   auto addPattern = m_Op<linalg::YieldOp>(m_Op<AddIOp>(a, b));
52   if (addPattern.match(&ops.back()))
53     return BinaryOpKind::IAdd;
54 
55   return llvm::None;
56 }
57 
58 static Value emitOrFoldComposedAffineApply(OpBuilder &b, Location loc,
59                                            AffineMap map,
60                                            ValueRange operandsRef,
61                                            OperationFolder *folder) {
62   SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
63   fullyComposeAffineMapAndOperands(&map, &operands);
64   canonicalizeMapAndOperands(&map, &operands);
65   return folder ? folder->create<AffineApplyOp>(b, loc, map, operands)
66                 : b.create<AffineApplyOp>(loc, map, operands);
67 }
68 
69 SmallVector<Value, 4> mlir::linalg::applyMapToValues(OpBuilder &b, Location loc,
70                                                      AffineMap map,
71                                                      ValueRange values,
72                                                      OperationFolder *folder) {
73   SmallVector<Value, 4> res;
74   res.reserve(map.getNumResults());
75   unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols();
76   // For each `expr` in `map`, applies the `expr` to the values extracted from
77   // ranges. If the resulting application can be folded into a Value, the
78   // folding occurs eagerly. Otherwise, an affine.apply operation is emitted.
79   for (auto expr : map.getResults()) {
80     AffineMap map = AffineMap::get(numDims, numSym, expr);
81     res.push_back(emitOrFoldComposedAffineApply(b, loc, map, values, folder));
82   }
83   return res;
84 }
85 
86 /// Returns all the operands of `linalgOp` that are not views.
87 /// Asserts that these operands are value types to allow transformations like
88 /// tiling to just use the values when cloning `linalgOp`.
89 SmallVector<Value, 4>
90 mlir::linalg::getAssumedNonViewOperands(LinalgOp linalgOp) {
91   auto *op = linalgOp.getOperation();
92   unsigned numViews = linalgOp.getNumInputsAndOutputs();
93   unsigned nOperands = op->getNumOperands() - numViews;
94   SmallVector<Value, 4> res;
95   res.reserve(nOperands);
96   for (unsigned i = 0; i < nOperands; ++i) {
97     res.push_back(op->getOperand(numViews + i));
98     auto t = res.back().getType();
99     (void)t;
100     assert((t.isSignlessIntOrIndexOrFloat() || t.isa<VectorType>()) &&
101            "expected scalar or vector type");
102   }
103   return res;
104 }
105 
106 bool mlir::linalg::isParallelIteratorType(Attribute attr) {
107   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
108     return strAttr.getValue() == getParallelIteratorTypeName();
109   }
110   return false;
111 }
112 
113 bool mlir::linalg::isReductionIteratorType(Attribute attr) {
114   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
115     return strAttr.getValue() == getReductionIteratorTypeName();
116   }
117   return false;
118 }
119 
120 bool mlir::linalg::isWindowIteratorType(Attribute attr) {
121   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
122     return strAttr.getValue() == getWindowIteratorTypeName();
123   }
124   return false;
125 }
126 
127 /// Explicit instantiation of loop nest generator for different loop types.
128 template struct mlir::linalg::GenerateLoopNest<scf::ForOp>;
129 template struct mlir::linalg::GenerateLoopNest<scf::ParallelOp>;
130 template struct mlir::linalg::GenerateLoopNest<AffineForOp>;
131 
132 /// Given a list of subview ranges, extract individual values for lower, upper
133 /// bounds and steps and put them into the corresponding vectors.
134 static void unpackRanges(ArrayRef<SubViewOp::Range> ranges,
135                          SmallVectorImpl<Value> &lbs,
136                          SmallVectorImpl<Value> &ubs,
137                          SmallVectorImpl<Value> &steps) {
138   for (SubViewOp::Range range : ranges) {
139     lbs.emplace_back(range.offset);
140     ubs.emplace_back(range.size);
141     steps.emplace_back(range.stride);
142   }
143 }
144 
145 namespace mlir {
146 namespace linalg {
147 
148 /// Specialization to build an scf "for" nest.
149 template <>
150 void GenerateLoopNest<scf::ForOp>::doit(
151     ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
152     function_ref<void(ValueRange)> bodyBuilderFn) {
153   SmallVector<Value, 4> lbs, ubs, steps;
154   unpackRanges(loopRanges, lbs, ubs, steps);
155   edsc::loopNestBuilder(lbs, ubs, steps, bodyBuilderFn);
156 }
157 
158 /// Specialization to build affine "for" nest.
159 template <>
160 void GenerateLoopNest<AffineForOp>::doit(
161     ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
162     function_ref<void(ValueRange)> bodyBuilderFn) {
163   SmallVector<Value, 4> lbs, ubs, steps;
164   unpackRanges(loopRanges, lbs, ubs, steps);
165 
166   // Affine loops require constant steps.
167   SmallVector<int64_t, 4> constantSteps;
168   constantSteps.reserve(steps.size());
169   for (Value v : steps) {
170     auto op = v.getDefiningOp<ConstantIndexOp>();
171     assert(op && "Affine loops require constant steps");
172     constantSteps.push_back(op.getValue());
173   }
174 
175   edsc::affineLoopNestBuilder(lbs, ubs, constantSteps, bodyBuilderFn);
176 }
177 
178 /// Generates a loop nest consisting of scf.parallel and scf.for, depending on
179 /// the `iteratorTypes.` Consecutive parallel loops create a single scf.parallel
180 /// operation; each sequential loop creates a new scf.for operation. The body
181 /// of the innermost loop is populated by `bodyBuilderFn` that accepts a range
182 /// of induction variables for all loops. `ivStorage` is used to store the
183 /// partial list of induction variables.
184 // TODO: this function can be made iterative instead. However, it
185 // will have at most as many recursive calls as nested loops, which rarely
186 // exceeds 10.
187 static void
188 generateParallelLoopNest(ValueRange lbs, ValueRange ubs, ValueRange steps,
189                          ArrayRef<Attribute> iteratorTypes,
190                          function_ref<void(ValueRange)> bodyBuilderFn,
191                          SmallVectorImpl<Value> &ivStorage) {
192   assert(lbs.size() == ubs.size());
193   assert(lbs.size() == steps.size());
194   assert(lbs.size() == iteratorTypes.size());
195 
196   // If there are no (more) loops to be generated, generate the body and be
197   // done with it.
198   if (iteratorTypes.empty())
199     return bodyBuilderFn(ivStorage);
200 
201   // Find the outermost parallel loops and drop their types from the list.
202   unsigned nLoops = iteratorTypes.size();
203   iteratorTypes = iteratorTypes.drop_while(isParallelIteratorType);
204   unsigned nOuterPar = nLoops - iteratorTypes.size();
205 
206   // If there are no outer parallel loops, generate one sequential loop and
207   // recurse. Note that we wouldn't have dropped anything from `iteratorTypes`
208   // in this case.
209   if (nOuterPar == 0) {
210     edsc::loopNestBuilder(lbs[0], ubs[0], steps[0], [&](Value iv) {
211       ivStorage.push_back(iv);
212       generateParallelLoopNest(lbs.drop_front(), ubs.drop_front(),
213                                steps.drop_front(), iteratorTypes.drop_front(),
214                                bodyBuilderFn, ivStorage);
215     });
216     return;
217   }
218 
219   // Generate a single parallel loop-nest operation for all outermost parallel
220   // loops and recurse.
221   edsc::OperationBuilder<scf::ParallelOp>(
222       lbs.take_front(nOuterPar), ubs.take_front(nOuterPar),
223       steps.take_front(nOuterPar),
224       [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
225         edsc::ScopedContext context(nestedBuilder, nestedLoc);
226         ivStorage.append(localIvs.begin(), localIvs.end());
227         generateParallelLoopNest(lbs.drop_front(nOuterPar),
228                                  ubs.drop_front(nOuterPar),
229                                  steps.drop_front(nOuterPar), iteratorTypes,
230                                  bodyBuilderFn, ivStorage);
231       });
232 }
233 
234 /// Specialization for generating a mix of parallel and sequential scf loops.
235 template <>
236 void GenerateLoopNest<scf::ParallelOp>::doit(
237     ArrayRef<SubViewOp::Range> loopRanges, ArrayRef<Attribute> iteratorTypes,
238     function_ref<void(ValueRange)> bodyBuilderFn) {
239   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
240   unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
241   ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
242 
243   // This function may be passed more iterator types than ranges.
244   assert(iteratorTypes.size() >= loopRanges.size() &&
245          "expected iterator type for all ranges");
246   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
247   ivs.reserve(iteratorTypes.size());
248   generateParallelLoopNest(lbs, ubs, steps, iteratorTypes, bodyBuilderFn, ivs);
249   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
250 }
251 
252 } // namespace linalg
253 } // namespace mlir
254