-
Notifications
You must be signed in to change notification settings - Fork 10.7k
/
PWMAFunction.h
182 lines (153 loc) · 7.66 KB
/
PWMAFunction.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
//===- PWMAFunction.h - MLIR PWMAFunction Class------------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Support for piece-wise multi-affine functions. These are functions that are
// defined on a domain that is a union of IntegerPolyhedrons, and on each domain
// the value of the function is a tuple of integers, with each value in the
// tuple being an affine expression in the vars of the IntegerPolyhedron.
//
//===----------------------------------------------------------------------===//
#ifndef MLIR_ANALYSIS_PRESBURGER_PWMAFUNCTION_H
#define MLIR_ANALYSIS_PRESBURGER_PWMAFUNCTION_H
#include "mlir/Analysis/Presburger/IntegerRelation.h"
#include "mlir/Analysis/Presburger/PresburgerRelation.h"
namespace mlir {
namespace presburger {
/// This class represents a multi-affine function whose domain is given by an
/// IntegerPolyhedron. This can be thought of as an IntegerPolyhedron with a
/// tuple of integer values attached to every point in the polyhedron, with the
/// value of each element of the tuple given by an affine expression in the vars
/// of the polyhedron. For example we could have the domain
///
/// (x, y) : (x >= 5, y >= x)
///
/// and a tuple of three integers defined at every point in the polyhedron:
///
/// (x, y) -> (x + 2, 2*x - 3y + 5, 2*x + y).
///
/// In this way every point in the polyhedron has a tuple of integers associated
/// with it. If the integer polyhedron has local vars, then the output
/// expressions can use them as well. The output expressions are represented as
/// a matrix with one row for every element in the output vector one column for
/// each var, and an extra column at the end for the constant term.
///
/// Checking equality of two such functions is supported, as well as finding the
/// value of the function at a specified point.
class MultiAffineFunction {
public:
MultiAffineFunction(const IntegerPolyhedron &domain, const Matrix &output)
: domainSet(domain), output(output) {}
MultiAffineFunction(const Matrix &output, const PresburgerSpace &space)
: domainSet(space), output(output) {}
unsigned getNumInputs() const { return domainSet.getNumDimAndSymbolVars(); }
unsigned getNumOutputs() const { return output.getNumRows(); }
bool isConsistent() const {
return output.getNumColumns() == domainSet.getNumVars() + 1;
}
const IntegerPolyhedron &getDomain() const { return domainSet; }
const PresburgerSpace &getDomainSpace() const { return domainSet.getSpace(); }
/// Insert `num` variables of the specified kind at position `pos`.
/// Positions are relative to the kind of variable. The coefficient columns
/// corresponding to the added variables are initialized to zero. Return the
/// absolute column position (i.e., not relative to the kind of variable)
/// of the first added variable.
unsigned insertVar(VarKind kind, unsigned pos, unsigned num = 1);
/// Remove the specified range of vars.
void removeVarRange(VarKind kind, unsigned varStart, unsigned varLimit);
/// Given a MAF `other`, merges local variables such that both funcitons
/// have union of local vars, without changing the set of points in domain or
/// the output.
void mergeLocalVars(MultiAffineFunction &other);
/// Return whether the outputs of `this` and `other` agree wherever both
/// functions are defined, i.e., the outputs should be equal for all points in
/// the intersection of the domains.
bool isEqualWhereDomainsOverlap(MultiAffineFunction other) const;
/// Return whether the `this` and `other` are equal. This is the case if
/// they lie in the same space, i.e. have the same dimensions, and their
/// domains are identical and their outputs are equal on their domain.
bool isEqual(const MultiAffineFunction &other) const;
/// Get the value of the function at the specified point. If the point lies
/// outside the domain, an empty optional is returned.
Optional<SmallVector<int64_t, 8>> valueAt(ArrayRef<int64_t> point) const;
/// Truncate the output dimensions to the first `count` dimensions.
///
/// TODO: refactor so that this can be accomplished through removeVarRange.
void truncateOutput(unsigned count);
void print(raw_ostream &os) const;
void dump() const;
private:
/// The IntegerPolyhedron representing the domain over which the function is
/// defined.
IntegerPolyhedron domainSet;
/// The function's output is a tuple of integers, with the ith element of the
/// tuple defined by the affine expression given by the ith row of this output
/// matrix.
Matrix output;
};
/// This class represents a piece-wise MultiAffineFunction. This can be thought
/// of as a list of MultiAffineFunction with disjoint domains, with each having
/// their own affine expressions for their output tuples. For example, we could
/// have a function with two input variables (x, y), defined as
///
/// f(x, y) = (2*x + y, y - 4) if x >= 0, y >= 0
/// = (-2*x + y, y + 4) if x < 0, y < 0
/// = (4, 1) if x < 0, y >= 0
///
/// Note that the domains all have to be *disjoint*. Otherwise, the behaviour of
/// this class is undefined. The domains need not cover all possible points;
/// this represents a partial function and so could be undefined at some points.
///
/// As in PresburgerSets, the input vars are partitioned into dimension vars and
/// symbolic vars.
///
/// Support is provided to compare equality of two such functions as well as
/// finding the value of the function at a point.
class PWMAFunction {
public:
PWMAFunction(const PresburgerSpace &space, unsigned numOutputs)
: space(space), numOutputs(numOutputs) {
assert(space.getNumDomainVars() == 0 &&
"Set type space should have zero domain vars.");
assert(space.getNumLocalVars() == 0 &&
"PWMAFunction cannot have local vars.");
assert(numOutputs >= 1 && "The function must output something!");
}
const PresburgerSpace &getSpace() const { return space; }
void addPiece(const MultiAffineFunction &piece);
void addPiece(const IntegerPolyhedron &domain, const Matrix &output);
const MultiAffineFunction &getPiece(unsigned i) const { return pieces[i]; }
unsigned getNumPieces() const { return pieces.size(); }
unsigned getNumOutputs() const { return numOutputs; }
unsigned getNumInputs() const { return space.getNumVars(); }
MultiAffineFunction &getPiece(unsigned i) { return pieces[i]; }
/// Return the domain of this piece-wise MultiAffineFunction. This is the
/// union of the domains of all the pieces.
PresburgerSet getDomain() const;
/// Return the value at the specified point and an empty optional if the
/// point does not lie in the domain.
Optional<SmallVector<int64_t, 8>> valueAt(ArrayRef<int64_t> point) const;
/// Return whether `this` and `other` are equal as PWMAFunctions, i.e. whether
/// they have the same dimensions, the same domain and they take the same
/// value at every point in the domain.
bool isEqual(const PWMAFunction &other) const;
/// Truncate the output dimensions to the first `count` dimensions.
///
/// TODO: refactor so that this can be accomplished through removeVarRange.
void truncateOutput(unsigned count);
void print(raw_ostream &os) const;
void dump() const;
private:
PresburgerSpace space;
/// The list of pieces in this piece-wise MultiAffineFunction.
SmallVector<MultiAffineFunction, 4> pieces;
/// The number of output vars.
unsigned numOutputs;
};
} // namespace presburger
} // namespace mlir
#endif // MLIR_ANALYSIS_PRESBURGER_PWMAFUNCTION_H