MLIR  20.0.0git
IterationGraphSorter.h
Go to the documentation of this file.
1 //===- IterationGraphSorter.h -----------------------------------*- C++ -*-===//
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 header file defines the iteration graph sorter (top-sort scheduling).
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_ITERATIONGRAPHSORTER_H_
14 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_ITERATIONGRAPHSORTER_H_
15 
16 #include "mlir/IR/AffineMap.h"
17 
18 namespace mlir {
19 
20 // Forward declarations.
21 class Value;
22 namespace utils {
23 enum class IteratorType : uint32_t;
24 } // namespace utils
25 namespace linalg {
26 class GenericOp;
27 } // namespace linalg
28 
29 namespace sparse_tensor {
30 
31 /// Iteration graph sorting mask,
32 enum class SortMask : unsigned {
33  // The individual mask bits.
34  kIncludeDenseOutput = 0x1, // b001
35  kIncludeDenseInput = 0x2, // b010
36  // The subsets of mask bits.
37  kIncludeAll = 0x7, // b111
38  kIncludeDense = 0x3, // b011
39  kSparseOnly = 0x0, // b000
40 };
41 
43 public:
44  /// Factory method that construct an iteration graph sorter
45  /// for the given linalg.generic operation.
46  static IterationGraphSorter fromGenericOp(linalg::GenericOp genericOp);
47 
48  /// Returns a permutation that represents the scheduled loop order.
49  /// Note that the returned AffineMap could be null if the kernel
50  /// cannot be scheduled due to cyclic iteration graph.
51  [[nodiscard]] AffineMap sort(SortMask mask, Value ignored = nullptr);
52 
53  /// Returns the number of loops in the iteration graph.
54  unsigned getNumLoops() const { return loop2OutLvl.getNumDims(); }
55 
56 private:
57  // Private constructor.
59  SmallVector<AffineMap> &&loop2InsLvl, Value out,
60  AffineMap loop2OutLvl,
62 
63  // Adds all the constraints in the given loop to level map.
64  void addConstraints(Value t, AffineMap loop2LvlMap);
65 
66  /// A helper to compute a topological sort. The method has an
67  /// O(n^2) time complexity since we use an adjacency matrix
68  /// representation for the iteration graph.
69  AffineMap topoSort();
70 
71  // Input tensors and associated loop to level maps.
73  SmallVector<AffineMap> loop2InsLvl;
74 
75  // Output tensor and associated loop to level map.
76  Value out;
77  AffineMap loop2OutLvl;
78 
79  // Loop itation types;
81 
82  // Adjacency matrix that represents the iteration graph.
83  std::vector<std::vector<bool>> itGraph;
84 
85  // InDegree used for topo sort.
86  std::vector<unsigned> inDegree;
87 };
88 
89 } // namespace sparse_tensor
90 } // namespace mlir
91 
92 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_ITERATIONGRAPHSORTER_H_
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
unsigned getNumLoops() const
Returns the number of loops in the iteration graph.
SortMask
Iteration graph sorting mask,.
Include the generated interface declarations.