Skip to content

[MLIR][Presburger] Template Matrix to allow MPInt and Fraction and separate out IntMatrix #66897

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 19 commits into from
Sep 20, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
Original file line number Diff line number Diff line change
Expand Up @@ -366,7 +366,7 @@ class IntegerRelation {
/// bounded. The span of the returned vectors is guaranteed to contain all
/// such vectors. The returned vectors are NOT guaranteed to be linearly
/// independent. This function should not be called on empty sets.
Matrix getBoundedDirections() const;
IntMatrix getBoundedDirections() const;

/// Find an integer sample point satisfying the constraints using a
/// branch and bound algorithm with generalized basis reduction, with some
Expand Down Expand Up @@ -792,10 +792,10 @@ class IntegerRelation {
PresburgerSpace space;

/// Coefficients of affine equalities (in == 0 form).
Matrix equalities;
IntMatrix equalities;

/// Coefficients of affine inequalities (in >= 0 form).
Matrix inequalities;
IntMatrix inequalities;
};

/// An IntegerPolyhedron represents the set of points from a PresburgerSpace
Expand Down
12 changes: 8 additions & 4 deletions mlir/include/mlir/Analysis/Presburger/LinearTransform.h
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@ namespace presburger {

class LinearTransform {
public:
explicit LinearTransform(Matrix &&oMatrix);
explicit LinearTransform(const Matrix &oMatrix);
explicit LinearTransform(IntMatrix &&oMatrix);
explicit LinearTransform(const IntMatrix &oMatrix);

// Returns a linear transform T such that MT is M in column echelon form.
// Also returns the number of non-zero columns in MT.
Expand All @@ -32,7 +32,7 @@ class LinearTransform {
// strictly below that of the previous column, and all columns which have only
// zeros are at the end.
static std::pair<unsigned, LinearTransform>
makeTransformToColumnEchelon(const Matrix &m);
makeTransformToColumnEchelon(const IntMatrix &m);

// Returns an IntegerRelation having a constraint vector vT for every
// constraint vector v in rel, where T is this transform.
Expand All @@ -50,8 +50,12 @@ class LinearTransform {
return matrix.postMultiplyWithColumn(colVec);
}

// Compute the determinant of the transform by converting it to row echelon
// form and then taking the product of the diagonal.
MPInt determinant();

private:
Matrix matrix;
IntMatrix matrix;
};

} // namespace presburger
Expand Down
110 changes: 70 additions & 40 deletions mlir/include/mlir/Analysis/Presburger/Matrix.h
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,17 @@
//===----------------------------------------------------------------------===//
//
// This is a simple 2D matrix class that supports reading, writing, resizing,
// swapping rows, and swapping columns.
// swapping rows, and swapping columns. It can hold integers (MPInt) or rational
// numbers (Fraction).
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_ANALYSIS_PRESBURGER_MATRIX_H
#define MLIR_ANALYSIS_PRESBURGER_MATRIX_H

#include "mlir/Analysis/Presburger/MPInt.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Analysis/Presburger/Fraction.h"
#include "mlir/Analysis/Presburger/Matrix.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/Support/raw_ostream.h"

Expand All @@ -32,7 +34,12 @@ namespace presburger {
/// (i, j) is stored at data[i*nReservedColumns + j]. The reserved but unused
/// columns always have all zero values. The reserved rows are just reserved
/// space in the underlying SmallVector's capacity.
/// This class only works for the types MPInt and Fraction, since the method
/// implementations are in the Matrix.cpp file. Only these two types have
/// been explicitly instantiated there.
template<typename T>
class Matrix {
static_assert(std::is_same_v<T,MPInt> || std::is_same_v<T,Fraction>, "T must be MPInt or Fraction.");
public:
Matrix() = delete;

Expand All @@ -49,21 +56,21 @@ class Matrix {
static Matrix identity(unsigned dimension);

/// Access the element at the specified row and column.
MPInt &at(unsigned row, unsigned column) {
T &at(unsigned row, unsigned column) {
assert(row < nRows && "Row outside of range");
assert(column < nColumns && "Column outside of range");
return data[row * nReservedColumns + column];
}

MPInt at(unsigned row, unsigned column) const {
T at(unsigned row, unsigned column) const {
assert(row < nRows && "Row outside of range");
assert(column < nColumns && "Column outside of range");
return data[row * nReservedColumns + column];
}

MPInt &operator()(unsigned row, unsigned column) { return at(row, column); }
T &operator()(unsigned row, unsigned column) { return at(row, column); }

MPInt operator()(unsigned row, unsigned column) const {
T operator()(unsigned row, unsigned column) const {
return at(row, column);
}

Expand All @@ -87,11 +94,11 @@ class Matrix {
void reserveRows(unsigned rows);

/// Get a [Mutable]ArrayRef corresponding to the specified row.
MutableArrayRef<MPInt> getRow(unsigned row);
ArrayRef<MPInt> getRow(unsigned row) const;
MutableArrayRef<T> getRow(unsigned row);
ArrayRef<T> getRow(unsigned row) const;

/// Set the specified row to `elems`.
void setRow(unsigned row, ArrayRef<MPInt> elems);
void setRow(unsigned row, ArrayRef<T> elems);

/// Insert columns having positions pos, pos + 1, ... pos + count - 1.
/// Columns that were at positions 0 to pos - 1 will stay where they are;
Expand Down Expand Up @@ -125,23 +132,23 @@ class Matrix {

void copyRow(unsigned sourceRow, unsigned targetRow);

void fillRow(unsigned row, const MPInt &value);
void fillRow(unsigned row, int64_t value) { fillRow(row, MPInt(value)); }
void fillRow(unsigned row, const T &value);
void fillRow(unsigned row, int64_t value) { fillRow(row, T(value)); }

/// Add `scale` multiples of the source row to the target row.
void addToRow(unsigned sourceRow, unsigned targetRow, const MPInt &scale);
void addToRow(unsigned sourceRow, unsigned targetRow, const T &scale);
void addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale) {
addToRow(sourceRow, targetRow, MPInt(scale));
addToRow(sourceRow, targetRow, T(scale));
}
/// Add `scale` multiples of the rowVec row to the specified row.
void addToRow(unsigned row, ArrayRef<MPInt> rowVec, const MPInt &scale);
void addToRow(unsigned row, ArrayRef<T> rowVec, const T &scale);

/// Add `scale` multiples of the source column to the target column.
void addToColumn(unsigned sourceColumn, unsigned targetColumn,
const MPInt &scale);
const T &scale);
void addToColumn(unsigned sourceColumn, unsigned targetColumn,
int64_t scale) {
addToColumn(sourceColumn, targetColumn, MPInt(scale));
addToColumn(sourceColumn, targetColumn, T(scale));
}

/// Negate the specified column.
Expand All @@ -150,32 +157,13 @@ class Matrix {
/// Negate the specified row.
void negateRow(unsigned row);

/// Divide the first `nCols` of the specified row by their GCD.
/// Returns the GCD of the first `nCols` of the specified row.
MPInt normalizeRow(unsigned row, unsigned nCols);
/// Divide the columns of the specified row by their GCD.
/// Returns the GCD of the columns of the specified row.
MPInt normalizeRow(unsigned row);

/// The given vector is interpreted as a row vector v. Post-multiply v with
/// this matrix, say M, and return vM.
SmallVector<MPInt, 8> preMultiplyWithRow(ArrayRef<MPInt> rowVec) const;
SmallVector<T, 8> preMultiplyWithRow(ArrayRef<T> rowVec) const;

/// The given vector is interpreted as a column vector v. Pre-multiply v with
/// this matrix, say M, and return Mv.
SmallVector<MPInt, 8> postMultiplyWithColumn(ArrayRef<MPInt> colVec) const;

/// Given the current matrix M, returns the matrices H, U such that H is the
/// column hermite normal form of M, i.e. H = M * U, where U is unimodular and
/// the matrix H has the following restrictions:
/// - H is lower triangular.
/// - The leading coefficient (the first non-zero entry from the top, called
/// the pivot) of a non-zero column is always strictly below of the leading
/// coefficient of the column before it; moreover, it is positive.
/// - The elements to the right of the pivots are zero and the elements to
/// the left of the pivots are nonnegative and strictly smaller than the
/// pivot.
std::pair<Matrix, Matrix> computeHermiteNormalForm() const;
SmallVector<T, 8> postMultiplyWithColumn(ArrayRef<T> colVec) const;

/// Resize the matrix to the specified dimensions. If a dimension is smaller,
/// the values are truncated; if it is bigger, the new values are initialized
Expand All @@ -192,7 +180,7 @@ class Matrix {
unsigned appendExtraRow();
/// Same as above, but copy the given elements into the row. The length of
/// `elems` must be equal to the number of columns.
unsigned appendExtraRow(ArrayRef<MPInt> elems);
unsigned appendExtraRow(ArrayRef<T> elems);

/// Print the matrix.
void print(raw_ostream &os) const;
Expand All @@ -211,10 +199,52 @@ class Matrix {

/// Stores the data. data.size() is equal to nRows * nReservedColumns.
/// data.capacity() / nReservedColumns is the number of reserved rows.
SmallVector<MPInt, 16> data;
SmallVector<T, 16> data;
};

// An inherited class for integer matrices, with no new data attributes.
// This is only used for the matrix-related methods which apply only
// to integers (hermite normal form computation and row normalisation).
class IntMatrix : public Matrix<MPInt>
{
public:
IntMatrix(unsigned rows, unsigned columns, unsigned reservedRows = 0,
unsigned reservedColumns = 0) :
Matrix<MPInt>(rows, columns, reservedRows, reservedColumns) {};

IntMatrix(Matrix<MPInt> m) :
Matrix<MPInt>(m.getNumRows(), m.getNumColumns(), m.getNumReservedRows(), m.getNumReservedColumns())
{
for (unsigned i = 0; i < m.getNumRows(); i++)
for (unsigned j = 0; j < m.getNumColumns(); j++)
at(i, j) = m(i, j);
};

/// Return the identity matrix of the specified dimension.
static IntMatrix identity(unsigned dimension);

/// Given the current matrix M, returns the matrices H, U such that H is the
/// column hermite normal form of M, i.e. H = M * U, where U is unimodular and
/// the matrix H has the following restrictions:
/// - H is lower triangular.
/// - The leading coefficient (the first non-zero entry from the top, called
/// the pivot) of a non-zero column is always strictly below of the leading
/// coefficient of the column before it; moreover, it is positive.
/// - The elements to the right of the pivots are zero and the elements to
/// the left of the pivots are nonnegative and strictly smaller than the
/// pivot.
std::pair<IntMatrix, IntMatrix> computeHermiteNormalForm() const;

/// Divide the first `nCols` of the specified row by their GCD.
/// Returns the GCD of the first `nCols` of the specified row.
MPInt normalizeRow(unsigned row, unsigned nCols);
/// Divide the columns of the specified row by their GCD.
/// Returns the GCD of the columns of the specified row.
MPInt normalizeRow(unsigned row);

};

} // namespace presburger
} // namespace mlir

#endif // MLIR_ANALYSIS_PRESBURGER_MATRIX_H
#endif // MLIR_ANALYSIS_PRESBURGER_MATRIX_H
8 changes: 4 additions & 4 deletions mlir/include/mlir/Analysis/Presburger/PWMAFunction.h
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,13 @@ enum class OrderingKind { EQ, NE, LT, LE, GT, GE };
/// value of the function at a specified point.
class MultiAffineFunction {
public:
MultiAffineFunction(const PresburgerSpace &space, const Matrix &output)
MultiAffineFunction(const PresburgerSpace &space, const IntMatrix &output)
: space(space), output(output),
divs(space.getNumVars() - space.getNumRangeVars()) {
assertIsConsistent();
}

MultiAffineFunction(const PresburgerSpace &space, const Matrix &output,
MultiAffineFunction(const PresburgerSpace &space, const IntMatrix &output,
const DivisionRepr &divs)
: space(space), output(output), divs(divs) {
assertIsConsistent();
Expand All @@ -65,7 +65,7 @@ class MultiAffineFunction {
PresburgerSpace getOutputSpace() const { return space.getRangeSpace(); }

/// Get a matrix with each row representing row^th output expression.
const Matrix &getOutputMatrix() const { return output; }
const IntMatrix &getOutputMatrix() const { return output; }
/// Get the `i^th` output expression.
ArrayRef<MPInt> getOutputExpr(unsigned i) const { return output.getRow(i); }

Expand Down Expand Up @@ -124,7 +124,7 @@ class MultiAffineFunction {
/// 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;
IntMatrix output;

/// Storage for division representation for each local variable in space.
DivisionRepr divs;
Expand Down
4 changes: 2 additions & 2 deletions mlir/include/mlir/Analysis/Presburger/Simplex.h
Original file line number Diff line number Diff line change
Expand Up @@ -338,7 +338,7 @@ class SimplexBase {
unsigned nSymbol;

/// The matrix representing the tableau.
Matrix tableau;
IntMatrix tableau;

/// This is true if the tableau has been detected to be empty, false
/// otherwise.
Expand Down Expand Up @@ -861,7 +861,7 @@ class Simplex : public SimplexBase {

/// Reduce the given basis, starting at the specified level, using general
/// basis reduction.
void reduceBasis(Matrix &basis, unsigned level);
void reduceBasis(IntMatrix &basis, unsigned level);
};

/// Takes a snapshot of the simplex state on construction and rolls back to the
Expand Down
2 changes: 1 addition & 1 deletion mlir/include/mlir/Analysis/Presburger/Utils.h
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ class DivisionRepr {
/// Each row of the Matrix represents a single division dividend. The
/// `i^th` row represents the dividend of the variable at `divOffset + i`
/// in the constraint system (and the `i^th` local variable).
Matrix dividends;
IntMatrix dividends;

/// Denominators of each division. If a denominator of a division is `0`, the
/// division variable is considered to not have a division representation.
Expand Down
2 changes: 1 addition & 1 deletion mlir/lib/Analysis/FlatLinearValueConstraints.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1292,7 +1292,7 @@ mlir::getMultiAffineFunctionFromMap(AffineMap map,
"AffineMap cannot produce divs without local representation");

// TODO: We shouldn't have to do this conversion.
Matrix mat(map.getNumResults(), map.getNumInputs() + divs.getNumDivs() + 1);
Matrix<MPInt> mat(map.getNumResults(), map.getNumInputs() + divs.getNumDivs() + 1);
for (unsigned i = 0, e = flattenedExprs.size(); i < e; ++i)
for (unsigned j = 0, f = flattenedExprs[i].size(); j < f; ++j)
mat(i, j) = flattenedExprs[i][j];
Expand Down
8 changes: 4 additions & 4 deletions mlir/lib/Analysis/Presburger/IntegerRelation.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -304,7 +304,7 @@ SymbolicLexOpt IntegerRelation::findSymbolicIntegerLexMax() const {
// Get lexmax by flipping range sign in the PWMA constraints.
for (auto &flippedPiece :
flippedSymbolicIntegerLexMax.lexopt.getAllPieces()) {
Matrix mat = flippedPiece.output.getOutputMatrix();
IntMatrix mat = flippedPiece.output.getOutputMatrix();
for (unsigned i = 0, e = mat.getNumRows(); i < e; i++)
mat.negateRow(i);
MultiAffineFunction maf(flippedPiece.output.getSpace(), mat);
Expand Down Expand Up @@ -738,7 +738,7 @@ bool IntegerRelation::isEmptyByGCDTest() const {
//
// It is sufficient to check the perpendiculars of the constraints, as the set
// of perpendiculars which are bounded must span all bounded directions.
Matrix IntegerRelation::getBoundedDirections() const {
IntMatrix IntegerRelation::getBoundedDirections() const {
// Note that it is necessary to add the equalities too (which the constructor
// does) even though we don't need to check if they are bounded; whether an
// inequality is bounded or not depends on what other constraints, including
Expand All @@ -759,7 +759,7 @@ Matrix IntegerRelation::getBoundedDirections() const {
// The direction vector is given by the coefficients and does not include the
// constant term, so the matrix has one fewer column.
unsigned dirsNumCols = getNumCols() - 1;
Matrix dirs(boundedIneqs.size() + getNumEqualities(), dirsNumCols);
IntMatrix dirs(boundedIneqs.size() + getNumEqualities(), dirsNumCols);

// Copy the bounded inequalities.
unsigned row = 0;
Expand Down Expand Up @@ -845,7 +845,7 @@ IntegerRelation::findIntegerSample() const {
// m is a matrix containing, in each row, a vector in which S is
// bounded, such that the linear span of all these dimensions contains all
// bounded dimensions in S.
Matrix m = getBoundedDirections();
IntMatrix m = getBoundedDirections();
// In column echelon form, each row of m occupies only the first rank(m)
// columns and has zeros on the other columns. The transform T that brings S
// to column echelon form is unimodular as well, so this is a suitable
Expand Down
6 changes: 3 additions & 3 deletions mlir/lib/Analysis/Presburger/LinearTransform.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,11 @@
using namespace mlir;
using namespace presburger;

LinearTransform::LinearTransform(Matrix &&oMatrix) : matrix(oMatrix) {}
LinearTransform::LinearTransform(const Matrix &oMatrix) : matrix(oMatrix) {}
LinearTransform::LinearTransform(IntMatrix &&oMatrix) : matrix(oMatrix) {}
LinearTransform::LinearTransform(const IntMatrix &oMatrix) : matrix(oMatrix) {}

std::pair<unsigned, LinearTransform>
LinearTransform::makeTransformToColumnEchelon(const Matrix &m) {
LinearTransform::makeTransformToColumnEchelon(const IntMatrix &m) {
// Compute the hermite normal form of m. This, is by definition, is in column
// echelon form.
auto [h, u] = m.computeHermiteNormalForm();
Expand Down
Loading