MLIR  22.0.0git
IterationGraphSorter.cpp
Go to the documentation of this file.
1 //===- IterationGraphSorter.cpp -------------------------------------------===//
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 #include "IterationGraphSorter.h"
10 
15 #include "mlir/IR/BuiltinTypes.h"
16 
17 using namespace mlir;
18 using namespace mlir::sparse_tensor;
19 
20 namespace {
21 
22 /// A helper class that visits an affine expression and tries to find
23 /// an AffineDimExpr to which the corresponding iterator from a GenericOp
24 /// matches the desired iterator type. If there is no matched iterator
25 /// type, the method returns the first DimExpr in the expression.
26 class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
27 public:
28  explicit AffineDimFinder(ArrayRef<utils::IteratorType> itTypes)
29  : iterTypes(itTypes) {}
30 
31  /// Overrides the visit method from AffineExprVisitor.
32  void visitDimExpr(AffineDimExpr expr) {
33  if (pickedDim == nullptr || pickIterType == iterTypes[expr.getPosition()])
34  pickedDim = expr;
35  }
36 
37  /// Sets the desired iterator type that we want to pick.
38  void setPickedIterType(utils::IteratorType iterType) {
39  pickIterType = iterType;
40  }
41 
42  /// Gets the desired AffineDimExpr.
43  AffineDimExpr getDimExpr() const {
44  return llvm::cast<AffineDimExpr>(pickedDim);
45  }
46 
47  /// Walks the graph in post order to find dim expr.
48  void walkPostOrder(AffineExpr expr) {
49  pickedDim = nullptr;
51  }
52 
53 private:
54  /// The picked AffineDimExpr after visit.
55  AffineExpr pickedDim;
56  /// The iterator type that we want.
57  utils::IteratorType pickIterType;
58  /// The mapping between levels and iterator types.
60 };
61 
62 /// Flattens an affine expression into a list of AffineDimExprs.
63 struct AffineDimCollector : public AffineExprVisitor<AffineDimCollector> {
64  // Overrides method from AffineExprVisitor.
65  void visitDimExpr(AffineDimExpr expr) { dims.push_back(expr); }
67 };
68 
69 } // namespace
70 
71 inline static bool includesAny(SortMask mask1, SortMask mask2) {
72  return static_cast<unsigned>(mask1) & static_cast<unsigned>(mask2);
73 }
74 
75 inline static bool includesDenseInput(SortMask mask) {
77 }
78 
79 inline static bool includesDenseOutput(SortMask mask) {
81 }
82 
83 AffineMap IterationGraphSorter::topoSort() {
84  // The sorted result will put the first Reduction iterator to the
85  // latest possible position.
86  std::vector<unsigned> redIt; // reduce iterator with 0 degree
87  std::vector<unsigned> parIt; // parallel iterator with 0 degree
88  const unsigned numLoops = getNumLoops();
89  for (unsigned i = 0; i < numLoops; i++) {
90  if (inDegree[i] == 0) {
91  if (iterTypes[i] == utils::IteratorType::reduction)
92  redIt.push_back(i);
93  else
94  parIt.push_back(i);
95  }
96  }
97 
98  SmallVector<unsigned> loopOrder;
99  while (!redIt.empty() || !parIt.empty()) {
100  // We always prefer a parallel loop over a reduction loop because putting
101  // a reduction loop early might make the loop sequence inadmissible.
102  auto &it = !parIt.empty() ? parIt : redIt;
103  auto src = it.back();
104  loopOrder.push_back(src);
105  it.pop_back();
106  // Update in-degree, and push 0-degree node into worklist.
107  for (unsigned dst = 0; dst < numLoops; dst++) {
108  if (itGraph[src][dst] && --inDegree[dst] == 0) {
109  if (iterTypes[dst] == utils::IteratorType::reduction)
110  redIt.push_back(dst);
111  else
112  parIt.push_back(dst);
113  }
114  }
115  }
116 
117  // Return the topological sort on success.
118  if (loopOrder.size() == numLoops)
119  return AffineMap::getPermutationMap(loopOrder, out.getContext());
120 
121  // Cycle detected.
122  return AffineMap();
123 }
124 
126 IterationGraphSorter::fromGenericOp(linalg::GenericOp genericOp) {
127  // Must be a demapped sparse kernel.
128  assert(!hasAnyNonIdentityOperandsOrResults(genericOp) &&
129  hasAnySparseOperandOrResult(genericOp) &&
130  genericOp.getNumDpsInits() == 1);
131 
132  SmallVector<AffineMap> loopMap = genericOp.getIndexingMapsArray();
133  SmallVector<Value> ins = genericOp.getDpsInputs();
134 
135  AffineMap outMap = loopMap.back();
136  loopMap.pop_back();
137 
138  Value out = genericOp.getDpsInitOperand(0)->get();
140  genericOp.getIteratorTypesArray();
141 
142  return IterationGraphSorter(std::move(ins), std::move(loopMap), out, outMap,
143  std::move(iterTypes));
144 }
145 
146 IterationGraphSorter::IterationGraphSorter(
147  SmallVector<Value> &&ins, SmallVector<AffineMap> &&loop2InsLvl, Value out,
148  AffineMap loop2OutLvl, SmallVector<utils::IteratorType> &&iterTypes)
149  : ins(std::move(ins)), loop2InsLvl(std::move(loop2InsLvl)), out(out),
150  loop2OutLvl(loop2OutLvl), iterTypes(std::move(iterTypes)) {
151  // One map per tensor.
152  assert(loop2InsLvl.size() == ins.size());
153  // All the affine maps have the same number of dimensions (loops).
154  assert(llvm::all_equal(llvm::map_range(
155  loop2InsLvl, [](AffineMap m) { return m.getNumDims(); })));
156  // The number of results of the map should match the rank of the tensor.
157  assert(llvm::all_of(llvm::zip(loop2InsLvl, ins), [](auto mvPair) {
158  auto [m, v] = mvPair;
159  return m.getNumResults() == cast<ShapedType>(v.getType()).getRank();
160  }));
161 
162  itGraph.resize(getNumLoops(), std::vector<bool>(getNumLoops(), false));
163  inDegree.resize(getNumLoops());
164 }
165 
167  // Reset the adjacency matrix that represents the iteration graph.
168  for (auto &row : itGraph)
169  llvm::fill(row, false);
170 
171  // Reset in-degree.
172  llvm::fill(inDegree, 0);
173 
174  // Add the constraints for the loop to level map.
175  for (auto [in, map] : llvm::zip(ins, loop2InsLvl)) {
176  // Get map and encoding.
177  const auto enc = getSparseTensorEncoding(in.getType());
178  // Skip dense inputs when not requested.
179  if ((!enc && !includesDenseInput(mask)) || in == ignored)
180  continue;
181  addConstraints(in, map);
182  }
183 
184  // Add the constraints for the output map.
185  const auto enc = getSparseTensorEncoding(out.getType());
186  if ((enc || includesDenseOutput(mask)) && out != ignored)
187  addConstraints(out, loop2OutLvl);
188 
189  // Return the topological sort (empty for cyclic).
190  return topoSort();
191 }
192 
193 void IterationGraphSorter::addConstraints(Value t, AffineMap loop2LvlMap) {
194  auto addIterOrdering = [this](unsigned f, unsigned t) {
195  if (!itGraph[f][t] && f != t) {
196  itGraph[f][t] = true;
197  inDegree[t]++;
198  }
199  };
200 
201  // Set up a reduction finder.
202  AffineDimFinder finder(iterTypes);
203  finder.setPickedIterType(utils::IteratorType::reduction);
204 
205  // To compute iteration graph for tensor[d0 + d1 + d3, d4 + d5 + d6],
206  // we require there exist d_x \in {d0, d1, d3} and d_y \in {d4, d5, d6},
207  // and d_x > d_y && {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
208  const Level lvlRank = loop2LvlMap.getNumResults();
209  for (Level lvl = 1; lvl < lvlRank; lvl++) {
210  const AffineExpr fa = loop2LvlMap.getResult(lvl - 1);
211  const AffineExpr ta = loop2LvlMap.getResult(lvl);
212 
213  if (llvm::isa<AffineDimExpr>(fa) || llvm::isa<AffineDimExpr>(ta)) {
214  // Special case when at least one loop2LvlExp is a simple AffineDimExpr
215  // (say, d0) and we require d0 > {d1, d2, ...} or {d1, d2, ...} > d0
216  AffineDimCollector fCollector;
217  fCollector.walkPostOrder(fa);
218  AffineDimCollector tCollector;
219  tCollector.walkPostOrder(ta);
220 
221  for (auto fd : fCollector.dims) {
222  for (auto td : tCollector.dims) {
223  const unsigned f = fd.getPosition();
224  const unsigned t = td.getPosition();
225  addIterOrdering(f, t);
226  }
227  }
228  continue;
229  }
230 
231  // When both loop2LvlExpr is compound, we pick an abitrary reduction loop
232  // from lhs and rhs and use them as d_x and d_y.
233  finder.walkPostOrder(fa);
234  const AffineDimExpr fexp = finder.getDimExpr();
235  const unsigned fldx = fexp.getPosition();
236 
237  finder.walkPostOrder(ta);
238  const AffineDimExpr texp = finder.getDimExpr();
239  const unsigned tldx = texp.getPosition();
240 
241  // d_x > d_y
242  addIterOrdering(fldx, tldx);
243 
244  AffineDimCollector fCollector;
245  fCollector.walkPostOrder(fa);
246  AffineDimCollector tCollector;
247  tCollector.walkPostOrder(ta);
248 
249  // Make sure dx and dy is the last.
250  for (auto fd : fCollector.dims) {
251  const unsigned f = fd.getPosition();
252  addIterOrdering(f, fldx);
253  }
254  for (auto td : tCollector.dims) {
255  const unsigned t = td.getPosition();
256  addIterOrdering(t, tldx);
257  }
258  // {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
259  // This is to ensure that the affine expressions are reduced in sparse
260  // tensor level ordering.
261  for (auto fd : fCollector.dims) {
262  const unsigned f = fd.getPosition();
263  if (f == fldx) // skip d_x
264  continue;
265  for (auto td : tCollector.dims) {
266  const unsigned t = td.getPosition();
267  if (t == tldx) // skip d_y
268  continue;
269  addIterOrdering(f, t);
270  }
271  }
272  }
273 }
static bool includesDenseOutput(SortMask mask)
static bool includesAny(SortMask mask1, SortMask mask2)
static bool includesDenseInput(SortMask mask)
A dimensional identifier appearing in an affine expression.
Definition: AffineExpr.h:223
unsigned getPosition() const
Definition: AffineExpr.cpp:346
See documentation for AffineExprVisitorBase.
RetTy walkPostOrder(AffineExpr expr)
Base type for affine expression.
Definition: AffineExpr.h:68
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
unsigned getNumDims() const
Definition: AffineMap.cpp:390
unsigned getNumResults() const
Definition: AffineMap.cpp:398
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.cpp:407
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
Definition: AffineMap.cpp:260
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
MLIRContext * getContext() const
Utility to get the associated MLIRContext that this value is defined in.
Definition: Value.h:108
Type getType() const
Return the type of this value.
Definition: Value.h:105
unsigned getNumLoops() const
Returns the number of loops in the iteration graph.
static IterationGraphSorter fromGenericOp(linalg::GenericOp genericOp)
Factory method that construct an iteration graph sorter for the given linalg.generic operation.
AffineMap sort(SortMask mask, Value ignored=nullptr)
Returns a permutation that represents the scheduled loop order.
bool hasAnySparseOperandOrResult(Operation *op)
Returns true iff MLIR operand has any sparse operand or result.
Definition: SparseTensor.h:196
uint64_t Level
The type of level identifiers and level-ranks.
Definition: SparseTensor.h:42
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
bool hasAnyNonIdentityOperandsOrResults(Operation *op)
Returns true iff MLIR operation has any sparse tensor with non-identity dim2lvl maps.
SortMask
Iteration graph sorting mask,.
Include the generated interface declarations.