MLIR  22.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 
17 #include "mlir/IR/AffineMap.h"
18 
19 namespace mlir {
20 
21 // Forward declarations.
22 class Value;
23 namespace utils {
24 enum class IteratorType : uint32_t;
25 } // namespace utils
26 namespace linalg {
27 class GenericOp;
28 } // namespace linalg
29 
30 namespace sparse_tensor {
31 
32 /// Iteration graph sorting mask,
33 enum class SortMask : unsigned {
34  // The individual mask bits.
35  kIncludeDenseOutput = 0x1, // b001
36  kIncludeDenseInput = 0x2, // b010
37  // The subsets of mask bits.
38  kIncludeAll = 0x7, // b111
39  kIncludeDense = 0x3, // b011
40  kSparseOnly = 0x0, // b000
41 };
42 
44 public:
45  /// Factory method that constructs an iteration graph sorter
46  /// for the given linalg.generic operation with a specific loop ordering
47  /// strategy.
49  fromGenericOp(linalg::GenericOp genericOp,
51 
52  /// Returns a permutation that represents the scheduled loop order.
53  /// Note that the returned AffineMap could be null if the kernel
54  /// cannot be scheduled due to cyclic iteration graph.
55  [[nodiscard]] AffineMap sort(SortMask mask, Value ignored = nullptr);
56 
57  /// Returns the number of loops in the iteration graph.
58  unsigned getNumLoops() const { return loop2OutLvl.getNumDims(); }
59 
60 private:
61  // Private constructor.
63  SmallVector<AffineMap> &&loop2InsLvl, Value out,
64  AffineMap loop2OutLvl,
68 
69  // Adds all the constraints in the given loop to level map.
70  void addConstraints(Value t, AffineMap loop2LvlMap);
71 
72  /// A helper to compute a topological sort. The method has an
73  /// O(n^2) time complexity since we use an adjacency matrix
74  /// representation for the iteration graph.
75  AffineMap topoSort();
76 
77  // Input tensors and associated loop to level maps.
79  SmallVector<AffineMap> loop2InsLvl;
80 
81  // Output tensor and associated loop to level map.
82  Value out;
83  AffineMap loop2OutLvl;
84 
85  // Loop itation types;
87 
88  // Adjacency matrix that represents the iteration graph.
89  std::vector<std::vector<bool>> itGraph;
90 
91  // InDegree used for topo sort.
92  std::vector<unsigned> inDegree;
93 
94  // Loop ordering strategy.
96 };
97 
98 } // namespace sparse_tensor
99 } // namespace mlir
100 
101 #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.
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).
SortMask
Iteration graph sorting mask,.
Include the generated interface declarations.