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
18namespace mlir {
19
20// Forward declarations.
21class Value;
22namespace utils {
23enum class IteratorType : uint32_t;
24} // namespace utils
25namespace linalg {
26class GenericOp;
27} // namespace linalg
28
29namespace sparse_tensor {
30
31/// Iteration graph sorting mask,
32enum 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
42class IterationGraphSorter {
43public:
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
56private:
57 // Private constructor.
58 IterationGraphSorter(SmallVector<Value> &&ins,
59 SmallVector<AffineMap> &&loop2InsLvl, Value out,
60 AffineMap loop2OutLvl,
61 SmallVector<utils::IteratorType> &&iterTypes);
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.
72 SmallVector<Value> ins;
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;
80 SmallVector<utils::IteratorType> iterTypes;
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_
93

source code of mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h