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