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;
IntegerSet(ImplType * set)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).
getEmptySet(unsigned numDims,unsigned numSymbols,MLIRContext * context)56 static IntegerSet getEmptySet(unsigned numDims, unsigned numSymbols,
57 MLIRContext *context) {
58 auto one = getAffineConstantExpr(1, context);
59 // 1 == 0.
60 return get(numDims, numSymbols, one, 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.
getAsOpaquePointer()111 const void *getAsOpaquePointer() const {
112 return static_cast<const void *>(set);
113 }
getFromOpaquePointer(const void * pointer)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.
hash_value(IntegerSet arg)124 inline ::llvm::hash_code hash_value(IntegerSet arg) {
125 return ::llvm::hash_value(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(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