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 
104  // Select loop based on strategy.
105  unsigned src;
106  switch (strategy) {
108  src = it.back();
109  break;
110  }
111 
112  loopOrder.push_back(src);
113  it.pop_back();
114  // Update in-degree, and push 0-degree node into worklist.
115  for (unsigned dst = 0; dst < numLoops; dst++) {
116  if (itGraph[src][dst] && --inDegree[dst] == 0) {
117  if (iterTypes[dst] == utils::IteratorType::reduction)
118  redIt.push_back(dst);
119  else
120  parIt.push_back(dst);
121  }
122  }
123  }
124 
125  // Return the topological sort on success.
126  if (loopOrder.size() == numLoops)
127  return AffineMap::getPermutationMap(loopOrder, out.getContext());
128 
129  // Cycle detected.
130  return AffineMap();
131 }
132 
134  linalg::GenericOp genericOp, sparse_tensor::LoopOrderingStrategy strategy) {
135  // Must be a demapped sparse kernel.
136  assert(!hasAnyNonIdentityOperandsOrResults(genericOp) &&
137  hasAnySparseOperandOrResult(genericOp) &&
138  genericOp.getNumDpsInits() == 1);
139 
140  SmallVector<AffineMap> loopMap = genericOp.getIndexingMapsArray();
141  SmallVector<Value> ins = genericOp.getDpsInputs();
142 
143  AffineMap outMap = loopMap.back();
144  loopMap.pop_back();
145 
146  Value out = genericOp.getDpsInitOperand(0)->get();
148  genericOp.getIteratorTypesArray();
149 
150  return IterationGraphSorter(std::move(ins), std::move(loopMap), out, outMap,
151  std::move(iterTypes), strategy);
152 }
153 
154 IterationGraphSorter::IterationGraphSorter(
155  SmallVector<Value> &&ins, SmallVector<AffineMap> &&loop2InsLvl, Value out,
156  AffineMap loop2OutLvl, SmallVector<utils::IteratorType> &&iterTypes,
158  : ins(std::move(ins)), loop2InsLvl(std::move(loop2InsLvl)), out(out),
159  loop2OutLvl(loop2OutLvl), iterTypes(std::move(iterTypes)),
160  strategy(strategy) {
161  // One map per tensor.
162  assert(loop2InsLvl.size() == ins.size());
163  // All the affine maps have the same number of dimensions (loops).
164  assert(llvm::all_equal(llvm::map_range(
165  loop2InsLvl, [](AffineMap m) { return m.getNumDims(); })));
166  // The number of results of the map should match the rank of the tensor.
167  assert(llvm::all_of(llvm::zip(loop2InsLvl, ins), [](auto mvPair) {
168  auto [m, v] = mvPair;
169  return m.getNumResults() == cast<ShapedType>(v.getType()).getRank();
170  }));
171 
172  itGraph.resize(getNumLoops(), std::vector<bool>(getNumLoops(), false));
173  inDegree.resize(getNumLoops());
174 }
175 
177  // Reset the adjacency matrix that represents the iteration graph.
178  for (auto &row : itGraph)
179  llvm::fill(row, false);
180 
181  // Reset in-degree.
182  llvm::fill(inDegree, 0);
183 
184  // Add the constraints for the loop to level map.
185  for (auto [in, map] : llvm::zip(ins, loop2InsLvl)) {
186  // Get map and encoding.
187  const auto enc = getSparseTensorEncoding(in.getType());
188  // Skip dense inputs when not requested.
189  if ((!enc && !includesDenseInput(mask)) || in == ignored)
190  continue;
191  addConstraints(in, map);
192  }
193 
194  // Add the constraints for the output map.
195  const auto enc = getSparseTensorEncoding(out.getType());
196  if ((enc || includesDenseOutput(mask)) && out != ignored)
197  addConstraints(out, loop2OutLvl);
198 
199  // Return the topological sort (empty for cyclic).
200  return topoSort();
201 }
202 
203 void IterationGraphSorter::addConstraints(Value t, AffineMap loop2LvlMap) {
204  auto addIterOrdering = [this](unsigned f, unsigned t) {
205  if (!itGraph[f][t] && f != t) {
206  itGraph[f][t] = true;
207  inDegree[t]++;
208  }
209  };
210 
211  // Set up a reduction finder.
212  AffineDimFinder finder(iterTypes);
213  finder.setPickedIterType(utils::IteratorType::reduction);
214 
215  // To compute iteration graph for tensor[d0 + d1 + d3, d4 + d5 + d6],
216  // we require there exist d_x \in {d0, d1, d3} and d_y \in {d4, d5, d6},
217  // and d_x > d_y && {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
218  const Level lvlRank = loop2LvlMap.getNumResults();
219  for (Level lvl = 1; lvl < lvlRank; lvl++) {
220  const AffineExpr fa = loop2LvlMap.getResult(lvl - 1);
221  const AffineExpr ta = loop2LvlMap.getResult(lvl);
222 
223  if (llvm::isa<AffineDimExpr>(fa) || llvm::isa<AffineDimExpr>(ta)) {
224  // Special case when at least one loop2LvlExp is a simple AffineDimExpr
225  // (say, d0) and we require d0 > {d1, d2, ...} or {d1, d2, ...} > d0
226  AffineDimCollector fCollector;
227  fCollector.walkPostOrder(fa);
228  AffineDimCollector tCollector;
229  tCollector.walkPostOrder(ta);
230 
231  for (auto fd : fCollector.dims) {
232  for (auto td : tCollector.dims) {
233  const unsigned f = fd.getPosition();
234  const unsigned t = td.getPosition();
235  addIterOrdering(f, t);
236  }
237  }
238  continue;
239  }
240 
241  // When both loop2LvlExpr is compound, we pick an abitrary reduction loop
242  // from lhs and rhs and use them as d_x and d_y.
243  finder.walkPostOrder(fa);
244  const AffineDimExpr fexp = finder.getDimExpr();
245  const unsigned fldx = fexp.getPosition();
246 
247  finder.walkPostOrder(ta);
248  const AffineDimExpr texp = finder.getDimExpr();
249  const unsigned tldx = texp.getPosition();
250 
251  // d_x > d_y
252  addIterOrdering(fldx, tldx);
253 
254  AffineDimCollector fCollector;
255  fCollector.walkPostOrder(fa);
256  AffineDimCollector tCollector;
257  tCollector.walkPostOrder(ta);
258 
259  // Make sure dx and dy is the last.
260  for (auto fd : fCollector.dims) {
261  const unsigned f = fd.getPosition();
262  addIterOrdering(f, fldx);
263  }
264  for (auto td : tCollector.dims) {
265  const unsigned t = td.getPosition();
266  addIterOrdering(t, tldx);
267  }
268  // {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
269  // This is to ensure that the affine expressions are reduced in sparse
270  // tensor level ordering.
271  for (auto fd : fCollector.dims) {
272  const unsigned f = fd.getPosition();
273  if (f == fldx) // skip d_x
274  continue;
275  for (auto td : tCollector.dims) {
276  const unsigned t = td.getPosition();
277  if (t == tldx) // skip d_y
278  continue;
279  addIterOrdering(f, t);
280  }
281  }
282  }
283 }
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
static IterationGraphSorter fromGenericOp(linalg::GenericOp genericOp, sparse_tensor::LoopOrderingStrategy strategy)
Factory method that constructs an iteration graph sorter for the given linalg.generic operation with ...
unsigned getNumLoops() const
Returns the number of loops in the iteration graph.
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
LoopOrderingStrategy
Defines a strategy for loop ordering during sparse code generation.
Definition: Passes.h:61
@ kDefault
Default strategy (eagerly selects last loop in topological sort).
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.