1 | //===- IntegerSet.h - MLIR Integer Set Class --------------------*- 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 | // Integer sets are sets of points from the integer lattice constrained by |
10 | // affine equality/inequality constraints. This class is meant to represent |
11 | // integer sets in the IR - for 'affine.if' operations and as attributes of |
12 | // other operations. It is typically expected to contain only a handful of |
13 | // affine constraints, and is immutable like an affine map. Integer sets are not |
14 | // unique'd unless the number of constraints they contain are below a certain |
15 | // threshold - although affine expressions that make up its equalities and |
16 | // inequalities are themselves unique. |
17 | |
18 | // This class is not meant for affine analysis and operations like set |
19 | // operations, emptiness checks, or other math operations for analysis and |
20 | // transformation. For the latter, use FlatAffineValueConstraints. |
21 | // |
22 | //===----------------------------------------------------------------------===// |
23 | |
24 | #ifndef MLIR_IR_INTEGERSET_H |
25 | #define MLIR_IR_INTEGERSET_H |
26 | |
27 | #include "mlir/IR/AffineExpr.h" |
28 | #include "llvm/ADT/ArrayRef.h" |
29 | |
30 | namespace mlir { |
31 | |
32 | namespace detail { |
33 | struct IntegerSetStorage; |
34 | } // namespace detail |
35 | |
36 | class MLIRContext; |
37 | |
38 | /// An integer set representing a conjunction of one or more affine equalities |
39 | /// and inequalities. An integer set in the IR is immutable like the affine map, |
40 | /// but integer sets are not unique'd unless the number of constraints in them |
41 | /// is below `kUniquingThreshold`. The affine expressions that make up the |
42 | /// equalities and inequalities of an integer set are themselves unique and are |
43 | /// allocated by the bump pointer allocator. |
44 | class IntegerSet { |
45 | public: |
46 | using ImplType = detail::IntegerSetStorage; |
47 | |
48 | constexpr IntegerSet() = default; |
49 | explicit IntegerSet(ImplType *set) : set(set) {} |
50 | |
51 | static IntegerSet get(unsigned dimCount, unsigned symbolCount, |
52 | ArrayRef<AffineExpr> constraints, |
53 | ArrayRef<bool> eqFlags); |
54 | |
55 | // Returns the canonical empty IntegerSet (i.e. a set with no integer points). |
56 | static IntegerSet getEmptySet(unsigned numDims, unsigned numSymbols, |
57 | MLIRContext *context) { |
58 | auto one = getAffineConstantExpr(constant: 1, context); |
59 | // 1 == 0. |
60 | return get(dimCount: numDims, symbolCount: numSymbols, constraints: one, eqFlags: true); |
61 | } |
62 | |
63 | /// Returns true if this is the canonical integer set. |
64 | bool isEmptyIntegerSet() const; |
65 | |
66 | /// This method substitutes any uses of dimensions and symbols (e.g. |
67 | /// dim#0 with dimReplacements[0]) in subexpressions and returns the modified |
68 | /// integer set. Because this can be used to eliminate dims and |
69 | /// symbols, the client needs to specify the number of dims and symbols in |
70 | /// the result. The returned map always has the same number of results. |
71 | IntegerSet replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements, |
72 | ArrayRef<AffineExpr> symReplacements, |
73 | unsigned numResultDims, |
74 | unsigned numResultSyms); |
75 | |
76 | explicit operator bool() { return set; } |
77 | bool operator==(IntegerSet other) const { return set == other.set; } |
78 | bool operator!=(IntegerSet other) const { return set != other.set; } |
79 | |
80 | unsigned getNumDims() const; |
81 | unsigned getNumSymbols() const; |
82 | unsigned getNumInputs() const; |
83 | unsigned getNumConstraints() const; |
84 | unsigned getNumEqualities() const; |
85 | unsigned getNumInequalities() const; |
86 | |
87 | ArrayRef<AffineExpr> getConstraints() const; |
88 | |
89 | AffineExpr getConstraint(unsigned idx) const; |
90 | |
91 | /// Returns the equality bits, which specify whether each of the constraints |
92 | /// is an equality or inequality. |
93 | ArrayRef<bool> getEqFlags() const; |
94 | |
95 | /// Returns true if the idx^th constraint is an equality, false if it is an |
96 | /// inequality. |
97 | bool isEq(unsigned idx) const; |
98 | |
99 | MLIRContext *getContext() const; |
100 | |
101 | /// Walk all of the AffineExpr's in this set's constraints. Each node in an |
102 | /// expression tree is visited in postorder. |
103 | void walkExprs(function_ref<void(AffineExpr)> callback) const; |
104 | |
105 | void print(raw_ostream &os) const; |
106 | void dump() const; |
107 | |
108 | friend ::llvm::hash_code hash_value(IntegerSet arg); |
109 | |
110 | /// Methods supporting C API. |
111 | const void *getAsOpaquePointer() const { |
112 | return static_cast<const void *>(set); |
113 | } |
114 | static IntegerSet getFromOpaquePointer(const void *pointer) { |
115 | return IntegerSet( |
116 | reinterpret_cast<ImplType *>(const_cast<void *>(pointer))); |
117 | } |
118 | |
119 | private: |
120 | ImplType *set{nullptr}; |
121 | }; |
122 | |
123 | // Make AffineExpr hashable. |
124 | inline ::llvm::hash_code hash_value(IntegerSet arg) { |
125 | return ::llvm::hash_value(ptr: arg.set); |
126 | } |
127 | |
128 | } // namespace mlir |
129 | namespace llvm { |
130 | |
131 | // IntegerSet hash just like pointers. |
132 | template <> |
133 | struct DenseMapInfo<mlir::IntegerSet> { |
134 | static mlir::IntegerSet getEmptyKey() { |
135 | auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey(); |
136 | return mlir::IntegerSet(static_cast<mlir::IntegerSet::ImplType *>(pointer)); |
137 | } |
138 | static mlir::IntegerSet getTombstoneKey() { |
139 | auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey(); |
140 | return mlir::IntegerSet(static_cast<mlir::IntegerSet::ImplType *>(pointer)); |
141 | } |
142 | static unsigned getHashValue(mlir::IntegerSet val) { |
143 | return mlir::hash_value(arg: val); |
144 | } |
145 | static bool isEqual(mlir::IntegerSet LHS, mlir::IntegerSet RHS) { |
146 | return LHS == RHS; |
147 | } |
148 | }; |
149 | |
150 | } // namespace llvm |
151 | #endif // MLIR_IR_INTEGERSET_H |
152 | |