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