diff --git a/mlir/include/mlir/Analysis/AffineStructures.h b/mlir/include/mlir/Analysis/AffineStructures.h --- a/mlir/include/mlir/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Analysis/AffineStructures.h @@ -126,19 +126,33 @@ /// intersection with no simplification of any sort attempted. void append(const FlatAffineConstraints &other); - // Checks for emptiness by performing variable elimination on all identifiers, - // running the GCD test on each equality constraint, and checking for invalid - // constraints. - // Returns true if the GCD test fails for any equality, or if any invalid - // constraints are discovered on any row. Returns false otherwise. + /// Checks for emptiness by performing variable elimination on all + /// identifiers, running the GCD test on each equality constraint, and + /// checking for invalid constraints. Returns true if the GCD test fails for + /// any equality, or if any invalid constraints are discovered on any row. + /// Returns false otherwise. bool isEmpty() const; - // Runs the GCD test on all equality constraints. Returns 'true' if this test - // fails on any equality. Returns 'false' otherwise. - // This test can be used to disprove the existence of a solution. If it - // returns true, no integer solution to the equality constraints can exist. + /// Runs the GCD test on all equality constraints. Returns 'true' if this test + /// fails on any equality. Returns 'false' otherwise. + /// This test can be used to disprove the existence of a solution. If it + /// returns true, no integer solution to the equality constraints can exist. bool isEmptyByGCDTest() const; + /// Runs the GCD test heuristic. If it proves inconclusive, falls back to + /// generalized basis reduction if the set is bounded. + /// + /// Returns true if the set of constraints is found to have no solution, + /// false if a solution exists or all tests were inconclusive. + bool isIntegerEmpty() const; + + /// Find a sample point satisfying the constraints. This uses a branch and + /// bound algorithm with generalized basis reduction, which always works if + /// the set is bounded. This should not be called for unbounded sets. + /// + /// Returns such a point if one exists, or an empty Optional otherwise. + Optional> findIntegerSample() const; + // Clones this object. std::unique_ptr clone() const; diff --git a/mlir/include/mlir/Analysis/Presburger/Fraction.h b/mlir/include/mlir/Analysis/Presburger/Fraction.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Analysis/Presburger/Fraction.h @@ -0,0 +1,86 @@ +//===- Fraction.h - MLIR Fraction 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 +// +//===----------------------------------------------------------------------===// +// +// This is a simple class to represent fractions. It supports multiplication, +// comparison, floor, and ceiling operations. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_ANALYSIS_PRESBURGER_FRACTION_H +#define MLIR_ANALYSIS_PRESBURGER_FRACTION_H + +#include + +namespace mlir { + +/// A class to represent fractions. The sign of the fraction is represented +/// in the sign of the numerator; the denominator is always positive. +/// +/// Note that overflows may occur if the numerator or denominator are not +/// representable by 64-bit integers. +struct Fraction { + /// Default constructor initializes the represented rational number to zero. + Fraction() : num(0), den(1) {} + + /// Construct a Fraction from a numerator and denominator. + Fraction(int64_t oNum, int64_t oDen) : num(oNum), den(oDen) { + if (den < 0) { + num = -num; + den = -den; + } + } + + // The numerator and denominator, respectively. The denominator is + int64_t num, den; +}; + +/// Three-way comparison between two fractions. +/// Returns +1, 0, and -1 if the first fraction is greater than, equal to, or +/// less than the second fraction, respectively. +inline int compare(Fraction x, Fraction y) { + int64_t diff = x.num * y.den - y.num * x.den; + if (diff > 0) + return +1; + if (diff < 0) + return -1; + return 0; +} + +/// Division rounds towards zero, so if the result is already non-negative +/// then num/den is the floor. If the result is negative and the division +/// leaves a remainder then we need to subtract one to get the floor. +inline int64_t floor(Fraction f) { + return f.num / f.den - (f.num < 0 && f.num % f.den != 0); +} + +/// Division rounds towards zero, so if the result is already non-positive +/// then num/den is the ceiling. If the result is positive and the division +/// leaves a remainder then we need to add one to get the ceiling. +inline int64_t ceil(Fraction f) { + return f.num / f.den + (f.num > 0 && f.num % f.den != 0); +} + +inline Fraction operator-(Fraction x) { return Fraction(-x.num, x.den); } + +inline bool operator<(Fraction x, Fraction y) { return compare(x, y) < 0; } + +inline bool operator<=(Fraction x, Fraction y) { return compare(x, y) <= 0; } + +inline bool operator==(Fraction x, Fraction y) { return compare(x, y) == 0; } + +inline bool operator>(Fraction x, Fraction y) { return compare(x, y) > 0; } + +inline bool operator>=(Fraction x, Fraction y) { return compare(x, y) >= 0; } + +inline Fraction operator*(Fraction x, Fraction y) { + return Fraction(x.num * y.num, x.den * y.den); +} + +} // namespace mlir + +#endif // MLIR_ANALYSIS_PRESBURGER_FRACTION_H diff --git a/mlir/include/mlir/Analysis/Presburger/Matrix.h b/mlir/include/mlir/Analysis/Presburger/Matrix.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Analysis/Presburger/Matrix.h @@ -0,0 +1,79 @@ +//===- Matrix.h - MLIR Matrix 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 +// +//===----------------------------------------------------------------------===// +// +// This is a simple 2D matrix class that supports reading, writing, resizing, +// swapping rows, and swapping columns. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_ANALYSIS_PRESBURGER_MATRIX_H +#define MLIR_ANALYSIS_PRESBURGER_MATRIX_H + +#include "mlir/Support/LLVM.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/Support/raw_ostream.h" + +#include + +namespace mlir { + +/// This is a simple class to represent a resizable matrix. +/// +/// The data is stored in the form of a vector of vectors. +class Matrix { +public: + Matrix() = delete; + + /// Construct a matrix with the specified number of rows and columns. + /// Initially, the values are default initialized. + Matrix(unsigned rows, unsigned columns); + + /// Return the identity matrix of the specified dimension. + static Matrix identity(unsigned dimension); + + /// Access the element at the specified row and column. + int64_t &at(unsigned row, unsigned column); + int64_t at(unsigned row, unsigned column) const; + int64_t &operator()(unsigned row, unsigned column); + int64_t operator()(unsigned row, unsigned column) const; + + /// Swap the given columns. + void swapColumns(unsigned column, unsigned otherColumn); + + /// Swap the given rows. + void swapRows(unsigned row, unsigned otherRow); + + unsigned getNumRows() const; + + unsigned getNumColumns() const; + + /// Get an ArrayRef corresponding to the specified row. + ArrayRef getRow(unsigned row) const; + + /// Add `scale` multiples of the source row to the target row. + void addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale); + + /// Resize the matrix to the specified dimensions. If a dimension is smaller, + /// the values are truncated; if it is bigger, the new values are default + /// initialized. + void resizeVertically(unsigned newNRows); + + /// Print the matrix. + void print(raw_ostream &os) const; + void dump() const; + +private: + unsigned nRows, nColumns; + + /// Stores the data. data.size() is equal to nRows * nColumns. + SmallVector data; +}; + +} // namespace mlir + +#endif // MLIR_ANALYSIS_PRESBURGER_MATRIX_H diff --git a/mlir/include/mlir/Analysis/Presburger/Simplex.h b/mlir/include/mlir/Analysis/Presburger/Simplex.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Analysis/Presburger/Simplex.h @@ -0,0 +1,327 @@ +//===- Simplex.h - MLIR Simplex 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 +// +//===----------------------------------------------------------------------===// +// +// Functionality to perform analysis on FlatAffineConstraints. In particular, +// support for performing emptiness checks. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H +#define MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H + +#include "mlir/Analysis/AffineStructures.h" +#include "mlir/Analysis/Presburger/Fraction.h" +#include "mlir/Analysis/Presburger/Matrix.h" +#include "mlir/Support/LogicalResult.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/Optional.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/raw_ostream.h" + +namespace mlir { + +class GBRSimplex; + +/// This class implements a version of the Simplex and Generalized Basis +/// Reduction algorithms, which can perform analysis of integer sets with affine +/// inequalities and equalities. A Simplex can be constructed +/// by specifying the dimensionality of the set. It supports adding affine +/// inequalities and equalities, and can perform emptiness checks, i.e., it can +/// find a solution to the set of constraints if one exists, or say that the +/// set is empty if no solution exists. Currently, this only works for bounded +/// sets. SImplex can also be constructed from a FlatAffineConstraints object. +/// +/// The implementation of this Simplex class, other than the functionality +/// for sampling, is based on the paper +/// "Simplify: A Theorem Prover for Program Checking" +/// by D. Detlefs, G. Nelson, J. B. Saxe. +/// +/// We define variables, constraints, and unknowns. Consider the example of a +/// two-dimensional set defined by 1 + 2x + 3y >= 0 and 2x - 3z >= 0. Here, +/// x, y, are variables while 1 + 2x + 3y >= 0, 2x - 3y >= 0 are +/// constraints. Unknowns are either variables or constraints, i.e., x, y, +/// 1 + 2x + 3y >= 0, 2x - 3y >= 0 are all unknowns. +/// +/// The implementation involves a matrix called a tableau, which can be thought +/// of as a 2D matrix of rational numbers having number of rows equal to the +/// number of constraints and number of columns equal to one plus the number of +/// variables. In our implementation, instead of storing rational numbers, we +/// store a common denominator for each row, so it is in fact a matrix of +/// integers with number of rows equal to number of constraints and number of +/// columns equal to _two_ plus the number of variables. For example, instead of +/// storing a row of three rationals [1/2, 2/3, 3], we would store [6, 3, 4, 18] +/// since 3/6 = 1/2, 4/6 = 2/3, and 18/6 = 3. +/// +/// Every row and column except the first and second columns is associated with +/// an unknown and every unknown is associated with a row or column. The second +/// column represents the constant, explained in more detail below. An unknown +/// associated with a row or column is said to be in row or column position +/// respectively. +/// +/// The vectors var and con store information about the variables and +/// constraints respectively, namely, whether they are in row or column +/// position, which row or column they are associated with, and whether they +/// correspond to a variable or a constraint. +/// +/// An unknown is addressed by its index. If the index i is non-negative, then +/// the variable var[i] is being addressed. If the index i is negative, then +/// the constraint con[~i] is being addressed. Effectively this maps +/// 0 -> var[0], 1 -> var[1], -1 -> con[0], -2 -> con[1], etc. rowUnknown[r] and +/// colUnknown[c] are the indexes of the unknowns associated with row r and +/// column c, respectively. +/// +/// The unknowns in column position are together called the basis. Initially the +/// basis is the set of variables -- in our example above, the initial basis is +/// x, y. +/// +/// The unknowns in row position are represented in terms of the basis unknowns. +/// If the basis unknowns are u_1, u_2, ... u_m, and a row in the tableau is +/// d, c, a_1, a_2, ... a_m, this representats the unknown for that row as +/// (c + a_1*u_1 + a_2*u_2 + ... + a_m*u_m)/d. In our running example, if the +/// basis is the initial basis of x, y, then the constraint 1 + 2x + 3y >= 0 +/// would be represented by the row [1, 1, 2, 3]. +/// +/// The association of unknowns to rows and columns can be changed by a process +/// called pivoting, where a row unknown and a column unknown exchange places +/// and the remaining row variables' representation is changed accordingly +/// by eliminating the old column unknown in favour of the new column unknown. +/// If we had pivoted the column for x with the row for 2x - 3y >= 0, +/// the new row for x would be [2, 1, 3] since x = (1*(2x - 3y) + 3*y)/2. +/// See the documentation for the pivot member function for details. +/// +/// The association of unknowns to rows and columns is called the _tableau +/// configuration_. The _sample value_ of an unknown in a particular tableau +/// configuration is its value if all the column unknowns were set to zero. +/// Concretely, for unknowns in column position the sample value is zero and +/// for unknowns in row position the sample value is the constant term divided +/// by the common denominator. +/// +/// The tableau configuration is called _consistent_ if the sample value of all +/// restricted unknowns is non-negative. Initially there are no constraints, and +/// the tableau is consistent. When a new constraint is added, its sample value +/// in the current tableau configuration may be negative. In that case, we try +/// to find a series of pivots to bring us to a consistent tableau +/// configuration, i.e. we try to make the new constraint's sample value +/// non-negative without making that of any other constraints negative. (See +/// findPivot and findPivotRow for details.) If this is not possible, then the +/// set of constraints is mutually contradictory and the tableau is marked +/// _empty_, which means the set of constraints has no solution. +/// +/// This Simplex class also supports taking snapshots of the current state +/// and rolling back to prior snapshots. This works by maintaing an undo log +/// of operations. Snapshots are just pointers to a particular location in the +/// log, and rolling back to a snapshot is done by reverting each log entry's +/// operation from the end until we reach the snapshot's location. +/// +/// Finding an integer sample is done with the Generalized Basis Reduction +/// algorithm. See the documentation for findIntegerSample and reduceBasis. +class Simplex { +public: + enum class Direction { Up, Down }; + + Simplex() = delete; + explicit Simplex(unsigned nVar); + explicit Simplex(const FlatAffineConstraints &constraints); + + /// Returns true if the tableau is empty (has conflicting constraints), + /// false otherwise. + bool isEmpty() const; + + /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n + /// is the current number of variables, then the corresponding inequality is + /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0. + void addInequality(ArrayRef coeffs); + + /// Returns the number of variables in the tableau. + unsigned numVariables() const; + + /// Returns the number of constraints in the tableau. + unsigned numConstraints() const; + + /// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n + /// is the current number of variables, then the corresponding equality is + /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0. + void addEquality(ArrayRef coeffs); + + /// Mark the tableau as being empty. + void markEmpty(); + + /// Get a snapshot of the current state. This is used for rolling back. + unsigned getSnapshot() const; + + /// Rollback to a snapshot. This invalidates all later snapshots. + void rollback(unsigned snapshot); + + /// Compute the maximum or minimum value of the given row, depending on + /// direction. + /// + /// Returns a (num, den) pair denoting the optimum, or None if no + /// optimum exists, i.e., if the expression is unbounded in this direction. + Optional computeRowOptimum(Direction direction, unsigned row); + + /// Compute the maximum or minimum value of the given expression, depending on + /// direction. + /// + /// Returns a (num, den) pair denoting the optimum, or a null value if no + /// optimum exists, i.e., if the expression is unbounded in this direction. + Optional computeOptimum(Direction direction, + ArrayRef coeffs); + + /// Returns a (min, max) pair denoting the minimum and maximum integer values + /// of the given expression. + std::pair computeIntegerBounds(ArrayRef coeffs); + + /// Returns true if the polytope is unbounded, i.e., extends to infinity in + /// some direction. Otherwise, returns false. + bool isUnbounded(); + + /// Make a tableau to represent a pair of points in the given tableaus, one in + /// tableau A and one in B. + static Simplex makeProduct(const Simplex &a, const Simplex &b); + + /// Returns the current sample point if it is integral. Otherwise, returns an + /// None. + Optional> getSamplePointIfIntegral() const; + + /// Returns an integer sample point if one exists, or None + /// otherwise. This should only be called for bounded sets. + Optional> findIntegerSample(); + + /// Print the tableau's internal state. + void print(raw_ostream &os) const; + void dump() const; + +private: + friend class GBRSimplex; + + enum class Orientation { Row, Column }; + + /// An Unknown is either a variable or a constraint. It is always associated + /// with either a row or column. Whether it's a row or a column is specified + /// by the orientation and pos identifies the specific row or column it is + /// associated with. If the unknown is restricted, then it has a + /// non-negativity constraint associated with it, i.e., its sample value must + /// always be non-negative and if it cannot be made non-negative without + /// violating other constraints, the tableau is empty. + struct Unknown { + Unknown(Orientation oOrientation, bool oRestricted, unsigned oPos) + : pos(oPos), orientation(oOrientation), restricted(oRestricted) {} + unsigned pos; + Orientation orientation; + bool restricted : 1; + + void print(raw_ostream &os) const { + os << (orientation == Orientation::Row ? "r" : "c"); + os << pos; + if (restricted) + os << " [>=0]"; + } + }; + + struct Pivot { + unsigned row, column; + }; + + /// Find a pivot to change the sample value of row in the specified + /// direction. The returned pivot row will be row if and only + /// if the unknown is unbounded in the specified direction. + /// + /// Returns a (row, col) pair denoting a pivot, or an empty Optional if + /// no valid pivot exists. + Optional findPivot(int row, Direction direction) const; + + /// Swap the row with the column in the tableau's data structures but not the + /// tableau itself. This is used by pivot. + void swapRowWithCol(unsigned row, unsigned col); + + /// Pivot the row with the column. + void pivot(unsigned row, unsigned col); + void pivot(Pivot pair); + + /// Returns the unknown associated with index. + const Unknown &unknownFromIndex(int index) const; + /// Returns the unknown associated with col. + const Unknown &unknownFromColumn(unsigned col) const; + /// Returns the unknown associated with row. + const Unknown &unknownFromRow(unsigned row) const; + /// Returns the unknown associated with index. + Unknown &unknownFromIndex(int index); + /// Returns the unknown associated with col. + Unknown &unknownFromColumn(unsigned col); + /// Returns the unknown associated with row. + Unknown &unknownFromRow(unsigned row); + + /// Add a new row to the tableau and the associated data structures. + unsigned addRow(ArrayRef coeffs); + + /// Normalize the given row by removing common factors between the numerator + /// and the denominator. + void normalizeRow(unsigned row); + + /// Swap the two rows in the tableau and associated data structures. + void swapRows(unsigned i, unsigned j); + + /// Restore the unknown to a non-negative sample value. + /// + /// Returns true if the unknown was successfully restored to a non-negative + /// sample value, false otherwise. + LogicalResult restoreRow(Unknown &u); + + enum class UndoLogEntry { RemoveLastConstraint, UnmarkEmpty }; + + /// Undo the operation represented by the log entry. + void undo(UndoLogEntry entry); + + /// Find a row that can be used to pivot the column in the specified + /// direction. If skipRow is not null, then this row is excluded + /// from consideration. The returned pivot will maintain all constraints + /// except the column itself and skipRow, if it is set. (if these unknowns + /// are restricted). + /// + /// Returns the row to pivot to, or an empty Optional if the column + /// is unbounded in the specified direction. + Optional findPivotRow(Optional skipRow, + Direction direction, unsigned col) const; + + /// Reduce the given basis, starting at the specified level, using general + /// basis reduction. + void reduceBasis(Matrix &basis, unsigned level); + + /// The number of rows in the tableau. + unsigned nRow; + + /// The number of columns in the tableau, including the common denominator + /// and the constant column. + unsigned nCol; + + /// The matrix representing the tableau. + Matrix tableau; + + /// This is true if the tableau has been detected to be empty, false + /// otherwise. + bool empty; + + /// Holds a log of operations, used for rolling back to a previous state. + SmallVector undoLog; + + /// These hold the indexes of the unknown at a given row or column position. + /// We keep these as signed integers since that makes it convenient to check + /// if an index corresponds to a variable or a constraint by checking the + /// sign. + /// + /// colUnknown is padded with two null indexes at the front since the first + /// two columns don't correspond to any unknowns. + SmallVector rowUnknown, colUnknown; + + /// These hold information about each unknown. + SmallVector con, var; +}; + +} // namespace mlir + +#endif // MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp --- a/mlir/lib/Analysis/AffineStructures.cpp +++ b/mlir/lib/Analysis/AffineStructures.cpp @@ -11,6 +11,7 @@ //===----------------------------------------------------------------------===// #include "mlir/Analysis/AffineStructures.h" +#include "mlir/Analysis/Presburger/Simplex.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineValueMap.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" @@ -1035,6 +1036,28 @@ return false; } +// First, try the GCD test heuristic. +// +// If that doesn't find the set empty, check if the set is unbounded. If it is, +// we cannot use the GBR algorithm and we conservatively return false. +// +// If the set is bounded, we use the complete emptiness check for this case +// provided by Simplex::findIntegerSample(), which gives a definitive answer. +bool FlatAffineConstraints::isIntegerEmpty() const { + if (isEmptyByGCDTest()) + return true; + + Simplex simplex(*this); + if (simplex.isUnbounded()) + return false; + return !simplex.findIntegerSample().hasValue(); +} + +Optional> +FlatAffineConstraints::findIntegerSample() const { + return Simplex(*this).findIntegerSample(); +} + /// Tightens inequalities given that we are dealing with integer spaces. This is /// analogous to the GCD test but applied to inequalities. The constant term can /// be reduced to the preceding multiple of the GCD of the coefficients, i.e., diff --git a/mlir/lib/Analysis/CMakeLists.txt b/mlir/lib/Analysis/CMakeLists.txt --- a/mlir/lib/Analysis/CMakeLists.txt +++ b/mlir/lib/Analysis/CMakeLists.txt @@ -25,6 +25,7 @@ MLIRCallInterfaces MLIRControlFlowInterfaces MLIRInferTypeOpInterface + MLIRPresburger MLIRSCF ) @@ -46,5 +47,8 @@ MLIRCallInterfaces MLIRControlFlowInterfaces MLIRInferTypeOpInterface + MLIRPresburger MLIRSCF ) + +add_subdirectory(Presburger) diff --git a/mlir/lib/Analysis/Presburger/CMakeLists.txt b/mlir/lib/Analysis/Presburger/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Analysis/Presburger/CMakeLists.txt @@ -0,0 +1,4 @@ +add_mlir_library(MLIRPresburger + Simplex.cpp + Matrix.cpp + ) diff --git a/mlir/lib/Analysis/Presburger/Matrix.cpp b/mlir/lib/Analysis/Presburger/Matrix.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Analysis/Presburger/Matrix.cpp @@ -0,0 +1,92 @@ +//===- Matrix.cpp - MLIR Matrix Class -------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Analysis/Presburger/Matrix.h" + +namespace mlir { + +Matrix::Matrix(unsigned rows, unsigned columns) + : nRows(rows), nColumns(columns), data(nRows * nColumns) {} + +Matrix Matrix::identity(unsigned dimension) { + Matrix matrix(dimension, dimension); + for (unsigned i = 0; i < dimension; ++i) + matrix(i, i) = 1; + return matrix; +} + +int64_t &Matrix::at(unsigned row, unsigned column) { + assert(row < getNumRows() && "Row outside of range"); + assert(column < getNumColumns() && "Column outside of range"); + return data[row * nColumns + column]; +} + +int64_t Matrix::at(unsigned row, unsigned column) const { + assert(row < getNumRows() && "Row outside of range"); + assert(column < getNumColumns() && "Column outside of range"); + return data[row * nColumns + column]; +} + +int64_t &Matrix::operator()(unsigned row, unsigned column) { + return at(row, column); +} + +int64_t Matrix::operator()(unsigned row, unsigned column) const { + return at(row, column); +} + +unsigned Matrix::getNumRows() const { return nRows; } + +unsigned Matrix::getNumColumns() const { return nColumns; } + +void Matrix::resizeVertically(unsigned newNRows) { + nRows = newNRows; + data.resize(nRows * nColumns); +} + +void Matrix::swapRows(unsigned row, unsigned otherRow) { + assert((row < getNumRows() && otherRow < getNumRows()) && + "Given row out of bounds"); + if (row == otherRow) + return; + for (unsigned col = 0; col < nColumns; col++) + std::swap(at(row, col), at(otherRow, col)); +} + +void Matrix::swapColumns(unsigned column, unsigned otherColumn) { + assert((column < getNumColumns() && otherColumn < getNumColumns()) && + "Given column out of bounds"); + if (column == otherColumn) + return; + for (unsigned row = 0; row < nRows; row++) + std::swap(at(row, column), at(row, otherColumn)); +} + +ArrayRef Matrix::getRow(unsigned row) const { + return {&data[row * nColumns], nColumns}; +} + +void Matrix::addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale) { + if (scale == 0) + return; + for (unsigned col = 0; col < nColumns; ++col) + at(targetRow, col) += scale * at(sourceRow, col); + return; +} + +void Matrix::print(raw_ostream &os) const { + for (unsigned row = 0; row < nRows; ++row) { + for (unsigned column = 0; column < nColumns; ++column) + os << at(row, column) << ' '; + os << '\n'; + } +} + +void Matrix::dump() const { print(llvm::errs()); } + +} // namespace mlir diff --git a/mlir/lib/Analysis/Presburger/Simplex.cpp b/mlir/lib/Analysis/Presburger/Simplex.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Analysis/Presburger/Simplex.cpp @@ -0,0 +1,1062 @@ +//===- Simplex.cpp - MLIR Simplex Class -----------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Analysis/Presburger/Simplex.h" +#include "mlir/Analysis/Presburger/Matrix.h" +#include "mlir/Support/MathExtras.h" + +namespace mlir { +using Direction = Simplex::Direction; + +const int nullIndex = std::numeric_limits::max(); + +/// Construct a Simplex object with `nVar` variables. +Simplex::Simplex(unsigned nVar) + : nRow(0), nCol(2), tableau(0, 2 + nVar), empty(false) { + colUnknown.push_back(nullIndex); + colUnknown.push_back(nullIndex); + for (unsigned i = 0; i < nVar; ++i) { + var.emplace_back(Orientation::Column, /*restricted=*/false, /*pos=*/nCol); + colUnknown.push_back(i); + nCol++; + } +} + +Simplex::Simplex(const FlatAffineConstraints &constraints) + : Simplex(constraints.getNumIds()) { + for (unsigned i = 0, numIneqs = constraints.getNumInequalities(); + i < numIneqs; ++i) + addInequality(constraints.getInequality(i)); + for (unsigned i = 0, numEqs = constraints.getNumEqualities(); i < numEqs; ++i) + addEquality(constraints.getEquality(i)); +} + +const Simplex::Unknown &Simplex::unknownFromIndex(int index) const { + assert(index != nullIndex && "nullIndex passed to unknownFromIndex"); + return index >= 0 ? var[index] : con[~index]; +} + +const Simplex::Unknown &Simplex::unknownFromColumn(unsigned col) const { + assert(col < nCol && "Invalid column"); + return unknownFromIndex(colUnknown[col]); +} + +const Simplex::Unknown &Simplex::unknownFromRow(unsigned row) const { + assert(row < nRow && "Invalid row"); + return unknownFromIndex(rowUnknown[row]); +} + +Simplex::Unknown &Simplex::unknownFromIndex(int index) { + assert(index != nullIndex && "nullIndex passed to unknownFromIndex"); + return index >= 0 ? var[index] : con[~index]; +} + +Simplex::Unknown &Simplex::unknownFromColumn(unsigned col) { + assert(col < nCol && "Invalid column"); + return unknownFromIndex(colUnknown[col]); +} + +Simplex::Unknown &Simplex::unknownFromRow(unsigned row) { + assert(row < nRow && "Invalid row"); + return unknownFromIndex(rowUnknown[row]); +} + +/// Add a new row to the tableau corresponding to the given constant term and +/// list of coefficients. The coefficients are specified as a vector of +/// (variable index, coefficient) pairs. +unsigned Simplex::addRow(ArrayRef coeffs) { + assert(coeffs.size() == 1 + var.size() && + "Incorrect number of coefficients!"); + + ++nRow; + // If the tableau is not big enough to accomodate the extra row, we extend it. + if (nRow >= tableau.getNumRows()) + tableau.resizeVertically(nRow); + rowUnknown.push_back(~con.size()); + con.emplace_back(Orientation::Row, false, nRow - 1); + + tableau(nRow - 1, 0) = 1; + tableau(nRow - 1, 1) = coeffs.back(); + for (unsigned col = 2; col < nCol; ++col) + tableau(nRow - 1, col) = 0; + + // Process each given variable coefficient. + for (unsigned i = 0; i < var.size(); ++i) { + unsigned pos = var[i].pos; + if (coeffs[i] == 0) + continue; + + if (var[i].orientation == Orientation::Column) { + // If a variable is in column position at column col, then we just add the + // coefficient for that variable (scaled by the common row denominator) to + // the corresponding entry in the new row. + tableau(nRow - 1, pos) += coeffs[i] * tableau(nRow - 1, 0); + continue; + } + + // If the variable is in row position, we need to add that row to the new + // row, scaled by the coefficient for the variable, accounting for the two + // rows potentially having different denominators. The new denominator is + // the lcm of the two. + int64_t lcm = mlir::lcm(tableau(nRow - 1, 0), tableau(pos, 0)); + int64_t nRowCoeff = lcm / tableau(nRow - 1, 0); + int64_t idxRowCoeff = coeffs[i] * (lcm / tableau(pos, 0)); + tableau(nRow - 1, 0) = lcm; + for (unsigned col = 1; col < nCol; ++col) + tableau(nRow - 1, col) = + nRowCoeff * tableau(nRow - 1, col) + idxRowCoeff * tableau(pos, col); + } + + normalizeRow(nRow - 1); + // Push to undo log along with the index of the new constraint. + undoLog.push_back(UndoLogEntry::RemoveLastConstraint); + return con.size() - 1; +} + +/// Normalize the row by removing factors that are common between the +/// denominator and all the numerator coefficients. +void Simplex::normalizeRow(unsigned row) { + int64_t gcd = 0; + for (unsigned col = 0; col < nCol; ++col) { + if (gcd == 1) + break; + gcd = llvm::greatestCommonDivisor(gcd, std::abs(tableau(row, col))); + } + for (unsigned col = 0; col < nCol; ++col) + tableau(row, col) /= gcd; +} + +namespace { +bool signMatchesDirection(int64_t elem, Direction direction) { + assert(elem != 0 && "elem should not be 0"); + return direction == Direction::Up ? elem > 0 : elem < 0; +} + +Direction flippedDirection(Direction direction) { + return direction == Direction::Up ? Direction::Down : Simplex::Direction::Up; +} +} // anonymous namespace + +/// Find a pivot to change the sample value of the row in the specified +/// direction. The returned pivot row will involve `row` if and only if the +/// unknown is unbounded in the specified direction. +/// +/// To increase (resp. decrease) the value of a row, we need to find a live +/// column with a non-zero coefficient. If the coefficient is positive, we need +/// to increase (decrease) the value of the column, and if the coefficient is +/// negative, we need to decrease (increase) the value of the column. Also, +/// we cannot decrease the sample value of restricted columns. +/// +/// If multiple columns are valid, we break ties by considering a lexicographic +/// ordering where we prefer unknowns with lower index. +Optional Simplex::findPivot(int row, + Direction direction) const { + Optional col; + for (unsigned j = 2; j < nCol; ++j) { + int64_t elem = tableau(row, j); + if (elem == 0) + continue; + + if (unknownFromColumn(j).restricted && + !signMatchesDirection(elem, direction)) + continue; + if (!col || colUnknown[j] < colUnknown[*col]) + col = j; + } + + if (!col) + return {}; + + Direction newDirection = + tableau(row, *col) < 0 ? flippedDirection(direction) : direction; + Optional maybePivotRow = findPivotRow(row, newDirection, *col); + return Pivot{maybePivotRow.getValueOr(row), *col}; +} + +/// Swap the associated unknowns for the row and the column. +/// +/// First we swap the index associated with the row and column. Then we update +/// the unknowns to reflect their new position and orientation. +void Simplex::swapRowWithCol(unsigned row, unsigned col) { + std::swap(rowUnknown[row], colUnknown[col]); + Unknown &uCol = unknownFromColumn(col); + Unknown &uRow = unknownFromRow(row); + uCol.orientation = Orientation::Column; + uRow.orientation = Orientation::Row; + uCol.pos = col; + uRow.pos = row; +} + +void Simplex::pivot(Pivot pair) { pivot(pair.row, pair.column); } + +/// Pivot pivotRow and pivotCol. +/// +/// Let R be the pivot row unknown and let C be the pivot col unknown. +/// Since initially R = a*C + sum b_i * X_i +/// (where the sum is over the other column's unknowns, x_i) +/// C = (R - (sum b_i * X_i))/a +/// +/// Let u be some other row unknown. +/// u = c*C + sum d_i * X_i +/// So u = c*(R - sum b_i * X_i)/a + sum d_i * X_i +/// +/// This results in the following transform: +/// pivot col other col pivot col other col +/// pivot row a b -> pivot row 1/a -b/a +/// other row c d other row c/a d - bc/a +/// +/// Taking into account the common denominators p and q: +/// +/// pivot col other col pivot col other col +/// pivot row a/p b/p -> pivot row p/a -b/a +/// other row c/q d/q other row cp/aq (da - bc)/aq +/// +/// The pivot row transform is accomplished be swapping a with the pivot row's +/// common denominator and negating the pivot row except for the pivot column +/// element. +void Simplex::pivot(unsigned pivotRow, unsigned pivotCol) { + assert(pivotCol >= 2 && "Refusing to pivot invalid column"); + + swapRowWithCol(pivotRow, pivotCol); + std::swap(tableau(pivotRow, 0), tableau(pivotRow, pivotCol)); + // We need to negate the whole pivot row except for the pivot column. + if (tableau(pivotRow, 0) < 0) { + // If the denominator is negative, we negate the row by simply negating the + // denominator. + tableau(pivotRow, 0) = -tableau(pivotRow, 0); + tableau(pivotRow, pivotCol) = -tableau(pivotRow, pivotCol); + } else { + for (unsigned col = 1; col < nCol; ++col) { + if (col == pivotCol) + continue; + tableau(pivotRow, col) = -tableau(pivotRow, col); + } + } + normalizeRow(pivotRow); + + for (unsigned row = 0; row < nRow; ++row) { + if (row == pivotRow) + continue; + if (tableau(row, pivotCol) == 0) // Nothing to do. + continue; + tableau(row, 0) *= tableau(pivotRow, 0); + for (unsigned j = 1; j < nCol; ++j) { + if (j == pivotCol) + continue; + // Add rather than subtract because the pivot row has been negated. + tableau(row, j) = tableau(row, j) * tableau(pivotRow, 0) + + tableau(row, pivotCol) * tableau(pivotRow, j); + } + tableau(row, pivotCol) *= tableau(pivotRow, pivotCol); + normalizeRow(row); + } +} + +/// Perform pivots until the unknown has a non-negative sample value or until +/// no more upward pivots can be performed. Return the sign of the final sample +/// value. +LogicalResult Simplex::restoreRow(Unknown &u) { + assert(u.orientation == Orientation::Row && + "unknown should be in row position"); + + while (tableau(u.pos, 1) < 0) { + Optional maybePivot = findPivot(u.pos, Direction::Up); + if (!maybePivot) + break; + + pivot(*maybePivot); + if (u.orientation == Orientation::Column) + return LogicalResult::Success; // the unknown is unbounded above. + } + return success(tableau(u.pos, 1) >= 0); +} + +/// Find a row that can be used to pivot the column in the specified direction. +/// This returns an empty optional if and only if the column is unbounded in the +/// specified direction (ignoring skipRow, if skipRow is set). +/// +/// If skipRow is set, this row is not considered, and (if it is restricted) its +/// restriction may be violated by the returned pivot. Usually, skipRow is set +/// because we don't want to move it to column position unless it is unbounded, +/// and we are either trying to increase the value of skipRow or explicitly +/// trying to make skipRow negative, so we are not concerned about this. +/// +/// If the direction is up (resp. down) and a restricted row has a negative +/// (positive) coefficient for the column, then this row imposes a bound on how +/// much the sample value of the column can change. Such a row with constant +/// term c and coefficient f for the column imposes a bound of c/|f| on the +/// change in sample value (in the specified direction). (note that c is +/// non-negative here since the row is restricted and the tableau is consistent) +/// +/// We iterate through the rows and pick the row which imposes the most +/// stringent bound, since pivoting with a row changes the row's sample value to +/// 0 and hence saturates the bound it imposes. We break ties between rows that +/// impose the same bound by considering a lexicographic ordering where we +/// prefer unknowns with lower index value. +Optional Simplex::findPivotRow(Optional skipRow, + Direction direction, + unsigned col) const { + Optional retRow; + int64_t retElem, retConst; + for (unsigned row = 0; row < nRow; ++row) { + if (skipRow && row == *skipRow) + continue; + int64_t elem = tableau(row, col); + if (elem == 0) + continue; + if (!unknownFromRow(row).restricted) + continue; + if (signMatchesDirection(elem, direction)) + continue; + int64_t constTerm = tableau(row, 1); + + if (!retRow) { + retRow = row; + retElem = elem; + retConst = constTerm; + continue; + } + + int64_t diff = retConst * elem - constTerm * retElem; + if ((diff == 0 && rowUnknown[row] < rowUnknown[*retRow]) || + (diff != 0 && !signMatchesDirection(diff, direction))) { + retRow = row; + retElem = elem; + retConst = constTerm; + } + } + return retRow; +} + +bool Simplex::isEmpty() const { return empty; } + +void Simplex::swapRows(unsigned i, unsigned j) { + if (i == j) + return; + tableau.swapRows(i, j); + std::swap(rowUnknown[i], rowUnknown[j]); + unknownFromRow(i).pos = i; + unknownFromRow(j).pos = j; +} + +/// Mark this tableau empty and push an entry to the undo stack. +void Simplex::markEmpty() { + undoLog.push_back(UndoLogEntry::UnmarkEmpty); + empty = true; +} + +/// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n +/// is the curent number of variables, then the corresponding inequality is +/// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0. +/// +/// We add the inequality and mark it as restricted. We then try to make its +/// sample value non-negative. If this is not possible, the tableau has become +/// empty and we mark it as such. +void Simplex::addInequality(ArrayRef coeffs) { + unsigned conIndex = addRow(coeffs); + Unknown &u = con[conIndex]; + u.restricted = true; + LogicalResult result = restoreRow(u); + if (failed(result)) + markEmpty(); +} + +/// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n +/// is the curent number of variables, then the corresponding equality is +/// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0. +/// +/// We simply add two opposing inequalities, which force the expression to +/// be zero. +void Simplex::addEquality(ArrayRef coeffs) { + addInequality(coeffs); + SmallVector negatedCoeffs; + for (int64_t coeff : coeffs) + negatedCoeffs.emplace_back(-coeff); + addInequality(negatedCoeffs); +} + +unsigned Simplex::numVariables() const { return var.size(); } +unsigned Simplex::numConstraints() const { return con.size(); } + +/// Return a snapshot of the curent state. This is just the current size of the +/// undo log. +unsigned Simplex::getSnapshot() const { return undoLog.size(); } + +void Simplex::undo(UndoLogEntry entry) { + if (entry == UndoLogEntry::RemoveLastConstraint) { + Unknown &constraint = con.back(); + if (constraint.orientation == Orientation::Column) { + unsigned column = constraint.pos; + Optional row; + + // Try to find any pivot row for this column that preserves tableau + // consistency (except possibly the column itself, which is going to be + // deallocated anyway). + // + // If no pivot row is found in either direction, then the unknown is + // unbounded in both directions and we are free to + // perform any pivot at all. To do this, we just need to find any row with + // a non-zero coefficient for the column. + if (Optional maybeRow = + findPivotRow({}, Direction::Up, column)) { + row = *maybeRow; + } else if (Optional maybeRow = + findPivotRow({}, Direction::Down, column)) { + row = *maybeRow; + } else { + // The loop doesn't find a pivot row only if the column has zero + // coefficients for every row. But the unknown is a constraint, + // so it was added initially as a row. Such a row could never have been + // pivoted to a column. So a pivot row will always be found. + for (unsigned i = 0; i < nRow; ++i) { + if (tableau(i, column) != 0) { + row = i; + break; + } + } + } + assert(row.hasValue() && "No pivot row found!"); + pivot(*row, column); + } + + // Move this unknown to the last row and remove the last row from the + // tableau. + swapRows(constraint.pos, nRow - 1); + // It is not strictly necessary to shrink the tableau, but for now we + // maintain the invariant that the tableau has exactly nRow rows. + tableau.resizeVertically(nRow - 1); + nRow--; + rowUnknown.pop_back(); + con.pop_back(); + } else if (entry == UndoLogEntry::UnmarkEmpty) { + empty = false; + } +} + +/// Rollback to the specified snapshot. +/// +/// We undo all the log entries until the log size when the snapshot was taken +/// is reached. +void Simplex::rollback(unsigned snapshot) { + while (undoLog.size() > snapshot) { + undo(undoLog.back()); + undoLog.pop_back(); + } +} + +Optional Simplex::computeRowOptimum(Direction direction, + unsigned row) { + while (Optional maybePivot = findPivot(row, direction)) { + if (maybePivot->row == row) + return {}; + pivot(*maybePivot); + } + return Fraction(tableau(row, 1), tableau(row, 0)); +} + +Optional Simplex::computeOptimum(Direction direction, + ArrayRef coeffs) { + assert(!empty && "Tableau should not be empty"); + + unsigned snapshot = getSnapshot(); + unsigned conIndex = addRow(coeffs); + unsigned row = con[conIndex].pos; + Optional optimum = computeRowOptimum(direction, row); + rollback(snapshot); + return optimum; +} + +bool Simplex::isUnbounded() { + if (empty) + return false; + + SmallVector dir(var.size() + 1); + for (unsigned i = 0; i < var.size(); ++i) { + dir[i] = 1; + + Optional maybeMax = computeOptimum(Direction::Up, dir); + if (!maybeMax) + return true; + + Optional maybeMin = computeOptimum(Direction::Down, dir); + if (!maybeMin) + return true; + + dir[i] = 0; + } + return false; +} + +/// Make a tableau to represent a pair of points in the original tableau. +/// +/// The product constraints and variables are stored as: first A's, then B's. +/// +/// The product tableau has row layout: +/// A's rows, B's rows. +/// +/// It has column layout: +/// denominator, constant, A's columns, B's columns. +Simplex Simplex::makeProduct(const Simplex &a, const Simplex &b) { + unsigned numVar = a.numVariables() + b.numVariables(); + unsigned numCon = a.numConstraints() + b.numConstraints(); + Simplex result(numVar); + + result.tableau.resizeVertically(numCon); + result.empty = a.empty || b.empty; + + auto concat = [](ArrayRef v, ArrayRef w) { + SmallVector result; + result.reserve(v.size() + w.size()); + result.insert(result.end(), v.begin(), v.end()); + result.insert(result.end(), w.begin(), w.end()); + return result; + }; + result.con = concat(a.con, b.con); + result.var = concat(a.var, b.var); + + auto indexFromBIndex = [&](int index) { + return index >= 0 ? a.numVariables() + index + : ~(a.numConstraints() + ~index); + }; + + result.colUnknown.assign(2, nullIndex); + for (unsigned i = 2; i < a.nCol; ++i) { + result.colUnknown.push_back(a.colUnknown[i]); + result.unknownFromIndex(result.colUnknown.back()).pos = + result.colUnknown.size() - 1; + } + for (unsigned i = 2; i < b.nCol; ++i) { + result.colUnknown.push_back(indexFromBIndex(b.colUnknown[i])); + result.unknownFromIndex(result.colUnknown.back()).pos = + result.colUnknown.size() - 1; + } + + auto appendRowFromA = [&](unsigned row) { + for (unsigned col = 0; col < a.nCol; ++col) + result.tableau(result.nRow, col) = a.tableau(row, col); + result.rowUnknown.push_back(a.rowUnknown[row]); + result.unknownFromIndex(result.rowUnknown.back()).pos = + result.rowUnknown.size() - 1; + result.nRow++; + }; + + // Also fixes the corresponding entry in rowUnknown and var/con (as the case + // may be). + auto appendRowFromB = [&](unsigned row) { + result.tableau(result.nRow, 0) = b.tableau(row, 0); + result.tableau(result.nRow, 1) = b.tableau(row, 1); + + unsigned offset = a.nCol - 2; + for (unsigned col = 2; col < b.nCol; ++col) + result.tableau(result.nRow, offset + col) = b.tableau(row, col); + result.rowUnknown.push_back(indexFromBIndex(b.rowUnknown[row])); + result.unknownFromIndex(result.rowUnknown.back()).pos = + result.rowUnknown.size() - 1; + result.nRow++; + }; + + for (unsigned row = 0; row < a.nRow; ++row) + appendRowFromA(row); + for (unsigned row = 0; row < b.nRow; ++row) + appendRowFromB(row); + + return result; +} + +Optional> Simplex::getSamplePointIfIntegral() const { + if (empty) + return {}; + + SmallVector sample; + for (const Unknown &u : var) { + if (u.orientation == Orientation::Column) { + sample.push_back(0); + } else { + if (tableau(u.pos, 1) % tableau(u.pos, 0) != 0) + return {}; + sample.push_back(tableau(u.pos, 1) / tableau(u.pos, 0)); + } + } + return sample; +} + +/// Given a simplex for a polytope, construct a new simplex whose variables are +/// identified with a pair of points (x, y) in the original polytope. Supports +/// some operations needed for generalized basis reduction. In what follows, +/// denotes the dot product of the vectors x and y. +/// +/// This supports adding equality constraints == 0. It also +/// supports rolling back this addition, by maintaining a snapshot stack that +/// contains a snapshot of the Simplex's state for each equality, just before +/// that equality was added. +class GBRSimplex { + using Orientation = Simplex::Orientation; + +public: + GBRSimplex(const Simplex &originalSimplex) + : simplex(Simplex::makeProduct(originalSimplex, originalSimplex)), + simplexConstraintOffset(simplex.numConstraints()) {} + + /// Add an equality == 0. + /// First pushes a snapshot for the current simplex state to the stack so + /// that this can be rolled back later. + void addEqualityForDirection(ArrayRef dir) { + assert( + std::any_of(dir.begin(), dir.end(), [](int64_t x) { return x != 0; }) && + "Direction passed is the zero vector!"); + snapshotStack.push_back(simplex.getSnapshot()); + simplex.addEquality(getCoeffsForDirection(dir)); + } + + /// Compute max() and save the dual variables for only the + /// direction equalities to `dual`. + Fraction computeWidthAndDuals(ArrayRef dir, + SmallVectorImpl &dual, + int64_t &dualDenom) { + unsigned snap = simplex.getSnapshot(); + unsigned conIndex = simplex.addRow(getCoeffsForDirection(dir)); + unsigned row = simplex.con[conIndex].pos; + Optional maybeWidth = + simplex.computeRowOptimum(Simplex::Direction::Up, row); + assert(maybeWidth.hasValue() && "Width should not be unbounded!"); + dualDenom = simplex.tableau(row, 0); + dual.clear(); + // The increment is i += 2 because equalities are added as two inequalities, + // one positive and one negative. Each iteration processes one equality. + for (unsigned i = simplexConstraintOffset; i < conIndex; i += 2) { + // The dual variable is the negative of the coefficient of the new row + // in the column of the constraint, if the constraint is in a column. + // Note that the second inequality for the equality is negated. + // + // We want the dual for the original equality. If the positive inequality + // is in column position, the negative of its row coefficient is the + // desired dual. If the negative inequality is in column position, its row + // coefficient is the desired dual. (its coefficients are already the + // negated coefficients of the original equality, so we don't need to + // negate it now.) + // + // If neither are in column position, we move the negated inequality to + // column position. Since the inequality must have sample value zero + // (since it corresponds to an equality), we are free to pivot with + // any column. Since both the unknowns have sample value before and after + // pivoting, no other sample values will change and the tableau will + // remain consistent. To pivot, we just need to find a column that has a + // non-zero coefficient in this row. There must be one since otherwise the + // equality would be 0 == 0, which should never be passed to + // addEqualityForDirection. + // + // After finding a column, we pivot with the column, after which we can + // get the dual from the inequality in column position as explained above. + if (simplex.con[i].orientation == Orientation::Column) { + dual.push_back(-simplex.tableau(row, simplex.con[i].pos)); + } else { + if (simplex.con[i + 1].orientation == Orientation::Row) { + unsigned ineqRow = simplex.con[i + 1].pos; + // Since it is an equality, the the sample value must be zero. + assert(simplex.tableau(ineqRow, 1) == 0 && + "Equality's sample value must be zero."); + for (unsigned col = 2; col < simplex.nCol; ++col) { + if (simplex.tableau(ineqRow, col) != 0) { + simplex.pivot(ineqRow, col); + break; + } + } + assert(simplex.con[i + 1].orientation == Orientation::Column && + "No pivot found. Equality has all-zeros row in tableau!"); + } + dual.push_back(simplex.tableau(row, simplex.con[i + 1].pos)); + } + } + simplex.rollback(snap); + return *maybeWidth; + } + + /// Remove the last equality that was added through addEqualityForDirection. + /// + /// We do this by rolling back to the snapshot at the top of the stack, which + /// should be a snapshot taken just before the last equality was added. + void removeLastEquality() { + assert(!snapshotStack.empty() && "Snapshot stack is empty!"); + simplex.rollback(snapshotStack.back()); + snapshotStack.pop_back(); + } + +private: + // Returns coefficients for the expression . + SmallVector getCoeffsForDirection(ArrayRef dir) { + assert(2 * dir.size() == simplex.numVariables() && + "Direction vector has wrong dimensionality"); + SmallVector coeffs(dir.begin(), dir.end()); + for (int64_t coeff : dir) + coeffs.emplace_back(-coeff); + coeffs.emplace_back(0); // constant term + return coeffs; + } + + Simplex simplex; + // The first index of the equality constraints, the index immediately after + // the last constraint in the initial product simplex. + unsigned simplexConstraintOffset; + // A stack of snapshots, used for rolling back. + SmallVector snapshotStack; +}; + +/// Reduce the basis to try and find a direction in which the polytope is +/// "thin". This only works for bounded polytopes. +/// +/// This is an implementation of the algorithm described in the paper +/// "An Implementation of Generalized Basis Reduction for Integer Programming" +/// by W. Cook, T. Rutherford, H. E. Scarf, D. Shallcross. +/// +/// Let b_{level}, b_{level + 1}, ... b_n be the current basis. +/// Let width_i(v) = max where x and y are points in the original +/// polytope such that = 0 is satisfied for all level <= j < i. +/// +/// In every iteration, we first replace b_{i+1} with b_{i+1} + u*b_i, where u +/// is the integer such that width_i(b_{i+1} + u*b_i) is minimized. Let dual_i +/// be the dual variable associated with the constraint = 0 when +/// computing width_{i+1}(b_{i+1}). It can be shown that dual_i is the +/// minimizing value of u, if it were allowed to be fractional. Due to +/// convexity, the minimizing integer value is either floor(dual_i) or +/// ceil(dual_i), so we just need to check which of these gives a lower +/// width_{i+1} value. If dual_i turned out to be an integer, then u = dual_i. +/// +/// Now if width_i(b_{i+1}) < 0.75 * width_i(b_i), we swap b_i and (the new) +/// b_{i + 1} and decrement i (unless i = level, in which case we stay at the +/// same i). Otherwise, we increment i. +/// +/// We keep f values and duals cached and invalidate them when necessary. +/// Whenever possible, we use them instead of recomputing them. We implement the +/// algorithm as follows. +/// +/// In an iteration at i we need to compute: +/// a) width_i(b_{i + 1}) +/// b) width_i(b_i) +/// c) the integer u that minimizes width_i(b_{i + 1} + u*b_i) +/// +/// If width_i(b_i) is not already cached, we compute it. +/// +/// If the duals are not already cached, we compute width_{i+1}(b_{i+1}) and +/// store the duals from this computation. +/// +/// We call updateBasisWithUAndGetFCandidate, which finds the minimizing value +/// of u as explained before, caches the duals from this computation, sets +/// b_{i+1} to b_{i+1} + u*b_i, and returns the new value of width_i(b_{i+1}). +/// +/// Now if width_i(b_{i+1}) < 0.75 * width_i(b_i), we swap b_i and b_{i+1} and +/// decrement i, resulting in the basis +/// ... b_{i - 1}, b_{i + 1} + u*b_i, b_i, b_{i+2}, ... +/// with corresponding f values +/// ... width_{i-1}(b_{i-1}), width_i(b_{i+1} + u*b_i), width_{i+1}(b_i), ... +/// The values up to i - 1 remain unchanged. We have just gotten the middle +/// value from updateBasisWithUAndGetFCandidate, so we can update that in the +/// cache. The value at width_{i+1}(b_i) is unknown, so we evict this value from +/// the cache. The iteration after decrementing needs exactly the duals from the +/// computation of width_i(b_{i + 1} + u*b_i), so we keep these in the cache. +/// +/// When incrementing i, no cached f values get invalidated. However, the cached +/// duals do get invalidated as the duals for the higher levels are different. +void Simplex::reduceBasis(Matrix &basis, unsigned level) { + const Fraction epsilon(3, 4); + + if (level == basis.getNumRows() - 1) + return; + + GBRSimplex gbrSimplex(*this); + SmallVector width; + SmallVector dual; + int64_t dualDenom; + + // Finds the value of u that minimizes width_i(b_{i+1} + u*b_i), caches the + // duals from this computation, sets b_{i+1} to b_{i+1} + u*b_i, and returns + // the new value of width_i(b_{i+1}). + // + // If dual_i is not an integer, the minimizing value must be either + // floor(dual_i) or ceil(dual_i). We compute the expression for both and + // choose the minimizing value. + // + // If dual_i is an integer, we don't need to perform these computations. We + // know that in this case, + // a) u = dual_i. + // b) one can show that dual_j for j < i are the same duals we would have + // gotten from computing width_i(b_{i + 1} + u*b_i), so the correct duals + // are the ones already in the cache. + // c) width_i(b_{i+1} + u*b_i) = min_{alpha} width_i(b_{i+1} + alpha * b_i), + // which + // one can show is equal to width_{i+1}(b_{i+1}). The latter value must + // be in the cache, so we get it from there and return it. + auto updateBasisWithUAndGetFCandidate = [&](unsigned i) -> Fraction { + assert(i < level + dual.size() && "dual_i is not known!"); + + int64_t u = floorDiv(dual[i - level], dualDenom); + basis.addToRow(i, i + 1, u); + if (dual[i - level] % dualDenom != 0) { + SmallVector candidateDual[2]; + int64_t candidateDualDenom[2]; + Fraction widthI[2]; + + // Initially u is floor(dual) and basis reflects this. + widthI[0] = gbrSimplex.computeWidthAndDuals( + basis.getRow(i + 1), candidateDual[0], candidateDualDenom[0]); + + // Now try ceil(dual), i.e. floor(dual) + 1. + ++u; + basis.addToRow(i, i + 1, 1); + widthI[1] = gbrSimplex.computeWidthAndDuals( + basis.getRow(i + 1), candidateDual[1], candidateDualDenom[1]); + + unsigned j = widthI[0] < widthI[1] ? 0 : 1; + if (j == 0) + // Subtract 1 to go from u = ceil(dual) back to floor(dual). + basis.addToRow(i, i + 1, -1); + dual = std::move(candidateDual[j]); + dualDenom = candidateDualDenom[j]; + return widthI[j]; + } + assert(i + 1 - level < width.size() && "width_{i+1} wasn't saved"); + // When dual minimizes f_i(b_{i+1} + dual*b_i), this is equal to + // width_{i+1}(b_{i+1}). + return width[i + 1 - level]; + }; + + // In the ith iteration of the loop, gbrSimplex has constraints for directions + // from `level` to i - 1. + unsigned i = level; + while (i < basis.getNumRows() - 1) { + if (i >= level + width.size()) { + // We don't even know the value of f_i(b_i), so let's find that first. + // We have to do this first since later we assume that width already + // contains values up to and including i. + + assert((i == 0 || i - 1 < level + width.size()) && + "We are at level i but we don't know the value of width_{i-1}"); + + // We don't actually use these duals at all, but it doesn't matter + // because this case should only occur when i is level, and there are no + // duals in that case anyway. + assert(i == level && "This case should only occur when i == level"); + width.push_back( + gbrSimplex.computeWidthAndDuals(basis.getRow(i), dual, dualDenom)); + } + + if (i >= level + dual.size()) { + assert(i + 1 >= level + width.size() && + "We don't know dual_i but we know " + "width_{i+1}"); + // We don't know dual for our level, so let's find it. + gbrSimplex.addEqualityForDirection(basis.getRow(i)); + width.push_back(gbrSimplex.computeWidthAndDuals(basis.getRow(i + 1), dual, + dualDenom)); + gbrSimplex.removeLastEquality(); + } + + // width_i(b_{i+1} + u*b_i) + Fraction widthICandidate = updateBasisWithUAndGetFCandidate(i); + if (widthICandidate < epsilon * width[i - level]) { + basis.swapRows(i, i + 1); + width[i - level] = widthICandidate; + // The values of width_{i+1}(b_{i+1}) and higher may change after the + // swap, so we remove the cached values here. + width.resize(i - level + 1); + if (i == level) { + dual.clear(); + continue; + } + + gbrSimplex.removeLastEquality(); + i--; + continue; + } + + // Invalidate duals since the higher level needs to recompute its own duals. + dual.clear(); + gbrSimplex.addEqualityForDirection(basis.getRow(i)); + i++; + } +} + +/// Search for an integer sample point using a branch and bound algorithm. +/// +/// Each row in the basis matrix is a vector, and the set of basis vectors +/// should span the space. Initially this is the identity matrix, +/// i.e., the basis vectors are just the variables. +/// +/// In every level, a value is assigned to the level-th basis vector, as +/// follows. Compute the minimum and maximum rational values of this direction. +/// If only one integer point lies in this range, constrain the variable to +/// have this value and recurse to the next variable. +/// +/// If the range has multiple values, perform generalized basis reduction via +/// reduceBasis and then compute the bounds again. Now we try constraining +/// this direction in the first value in this range and "recurse" to the next +/// level. If we fail to find a sample, we try assigning the direction the next +/// value in this range, and so on. +/// +/// If no integer sample is found from any of the assignments, or if the range +/// contains no integer value, then of course the polytope is empty for the +/// current assignment of the values in previous levels, so we return to +/// the previous level. +/// +/// If we reach the last level where all the variables have been assigned values +/// already, then we simply return the current sample point if it is integral, +/// and go back to the previous level otherwise. +/// +/// To avoid potentially arbitrarily large recursion depths leading to stack +/// overflows, this algorithm is implemented iteratively. +Optional> Simplex::findIntegerSample() { + if (empty) + return {}; + + unsigned nDims = var.size(); + Matrix basis = Matrix::identity(nDims); + + unsigned level = 0; + // The snapshot just before constraining a direction to a value at each level. + SmallVector snapshotStack; + // The maximum value in the range of the direction for each level. + SmallVector upperBoundStack; + // The next value to try constraining the basis vector to at each level. + SmallVector nextValueStack; + + snapshotStack.reserve(basis.getNumRows()); + upperBoundStack.reserve(basis.getNumRows()); + nextValueStack.reserve(basis.getNumRows()); + while (level != -1u) { + if (level == basis.getNumRows()) { + // We've assigned values to all variables. Return if we have a sample, + // or go back up to the previous level otherwise. + if (auto maybeSample = getSamplePointIfIntegral()) + return maybeSample; + level--; + continue; + } + + if (level >= upperBoundStack.size()) { + // We haven't populated the stack values for this level yet, so we have + // just come down a level ("recursed"). Find the lower and upper bounds. + // If there is more than one integer point in the range, perform + // generalized basis reduction. + SmallVector basisCoeffs(basis.getRow(level).begin(), + basis.getRow(level).end()); + basisCoeffs.push_back(0); + + int64_t minRoundedUp, maxRoundedDown; + std::tie(minRoundedUp, maxRoundedDown) = + computeIntegerBounds(basisCoeffs); + + // Heuristic: if the sample point is integral at this point, just return + // it. + if (auto maybeSample = getSamplePointIfIntegral()) + return *maybeSample; + + if (minRoundedUp < maxRoundedDown) { + reduceBasis(basis, level); + basisCoeffs = SmallVector(basis.getRow(level).begin(), + basis.getRow(level).end()); + basisCoeffs.push_back(0); + std::tie(minRoundedUp, maxRoundedDown) = + computeIntegerBounds(basisCoeffs); + } + + snapshotStack.push_back(getSnapshot()); + // The smallest value in the range is the next value to try. + nextValueStack.push_back(minRoundedUp); + upperBoundStack.push_back(maxRoundedDown); + } + + assert((snapshotStack.size() - 1 == level && + nextValueStack.size() - 1 == level && + upperBoundStack.size() - 1 == level) && + "Mismatched variable stack sizes!"); + + // Whether we "recursed" or "returned" from a lower level, we rollback + // to the snapshot of the starting state at this level. (in the "recursed" + // case this has no effect) + rollback(snapshotStack.back()); + int64_t nextValue = nextValueStack.back(); + nextValueStack.back()++; + if (nextValue > upperBoundStack.back()) { + // We have exhausted the range and found no solution. Pop the stack and + // return up a level. + snapshotStack.pop_back(); + nextValueStack.pop_back(); + upperBoundStack.pop_back(); + level--; + continue; + } + + // Try the next value in the range and "recurse" into the next level. + SmallVector basisCoeffs(basis.getRow(level).begin(), + basis.getRow(level).end()); + basisCoeffs.push_back(-nextValue); + addEquality(basisCoeffs); + level++; + } + + return {}; +} + +/// Compute the minimum and maximum integer values the expression can take. We +/// compute each separately. +std::pair +Simplex::computeIntegerBounds(ArrayRef coeffs) { + int64_t minRoundedUp; + if (Optional maybeMin = + computeOptimum(Simplex::Direction::Down, coeffs)) + minRoundedUp = ceil(*maybeMin); + else + llvm_unreachable("Tableau should not be unbounded"); + + int64_t maxRoundedDown; + if (Optional maybeMax = + computeOptimum(Simplex::Direction::Up, coeffs)) + maxRoundedDown = floor(*maybeMax); + else + llvm_unreachable("Tableau should not be unbounded"); + + return {minRoundedUp, maxRoundedDown}; +} + +void Simplex::print(raw_ostream &os) const { + os << "rows = " << nRow << ", columns = " << nCol << "\n"; + if (empty) + os << "Simplex marked empty!\n"; + os << "var: "; + for (unsigned i = 0; i < var.size(); ++i) { + if (i > 0) + os << ", "; + var[i].print(os); + } + os << "\ncon: "; + for (unsigned i = 0; i < con.size(); ++i) { + if (i > 0) + os << ", "; + con[i].print(os); + } + os << '\n'; + for (unsigned row = 0; row < nRow; ++row) { + if (row > 0) + os << ", "; + os << "r" << row << ": " << rowUnknown[row]; + } + os << '\n'; + os << "c0: denom, c1: const"; + for (unsigned col = 2; col < nCol; ++col) + os << ", c" << col << ": " << colUnknown[col]; + os << '\n'; + for (unsigned row = 0; row < nRow; ++row) { + for (unsigned col = 0; col < nCol; ++col) + os << tableau(row, col) << '\t'; + os << '\n'; + } + os << '\n'; +} + +void Simplex::dump() const { print(llvm::errs()); } + +} // namespace mlir diff --git a/mlir/unittests/Analysis/AffineStructuresTest.cpp b/mlir/unittests/Analysis/AffineStructuresTest.cpp new file mode 100644 --- /dev/null +++ b/mlir/unittests/Analysis/AffineStructuresTest.cpp @@ -0,0 +1,200 @@ +//===- AffineStructuresTest.cpp - Tests for AffineStructures ----*- 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Analysis/AffineStructures.h" + +#include +#include + +#include + +namespace mlir { + +/// Evaluate the value of the given affine expression at the specified point. +/// The expression is a list of coefficients for the dimensions followed by the +/// constant term. +int64_t valueAt(ArrayRef expr, ArrayRef point) { + assert(expr.size() == 1 + point.size()); + int64_t value = expr.back(); + for (unsigned i = 0; i < point.size(); ++i) + value += expr[i] * point[i]; + return value; +} + +/// If hasValue is true, check that findIntegerSample returns a valid sample +/// for the FlatAffineConstraints fac. +/// +/// If hasValue is false, check that findIntegerSample does not return None. +void checkSample(bool hasValue, const FlatAffineConstraints &fac) { + Optional> maybeSample = fac.findIntegerSample(); + if (!hasValue) { + EXPECT_FALSE(maybeSample.hasValue()); + if (maybeSample.hasValue()) { + for (auto x : *maybeSample) + llvm::errs() << x << ' '; + llvm::errs() << '\n'; + } + } else { + ASSERT_TRUE(maybeSample.hasValue()); + for (unsigned i = 0; i < fac.getNumEqualities(); ++i) + EXPECT_EQ(valueAt(fac.getEquality(i), *maybeSample), 0); + for (unsigned i = 0; i < fac.getNumInequalities(); ++i) + EXPECT_GE(valueAt(fac.getInequality(i), *maybeSample), 0); + } +} + +/// Construct a FlatAffineConstraints from a set of inequality and +/// equality constraints. +FlatAffineConstraints +makeFACFromConstraints(unsigned dims, + SmallVector, 4> ineqs, + SmallVector, 4> eqs) { + FlatAffineConstraints fac(ineqs.size(), eqs.size(), dims + 1, dims); + for (const auto &eq : eqs) + fac.addEquality(eq); + for (const auto &ineq : ineqs) + fac.addInequality(ineq); + return fac; +} + +/// Check sampling for all the permutations of the dimensions for the given +/// constraint set. Since the GBR algorithm progresses dimension-wise, different +/// orderings may cause the algorithm to proceed differently. At lesast some of +///.these permutations should make it past the heuristics and test the +/// implementation of the GBR algorithm itself. +void checkPermutationsSample(bool hasValue, unsigned nDim, + SmallVector, 4> ineqs, + SmallVector, 4> eqs) { + SmallVector perm(nDim); + std::iota(perm.begin(), perm.end(), 0); + auto permute = [&perm](ArrayRef coeffs) { + SmallVector permuted; + for (unsigned id : perm) + permuted.push_back(coeffs[id]); + permuted.push_back(coeffs.back()); + return permuted; + }; + do { + SmallVector, 4> permutedIneqs, permutedEqs; + for (const auto &ineq : ineqs) + permutedIneqs.push_back(permute(ineq)); + for (const auto &eq : eqs) + permutedEqs.push_back(permute(eq)); + checkSample(hasValue, + makeFACFromConstraints(nDim, permutedIneqs, permutedEqs)); + } while (std::next_permutation(perm.begin(), perm.end())); +} + +TEST(FlatAffineConstraintsTest, FindSampleTest) { + // Bounded sets with only inequalities. + // 1 <= 5x and 5x <= 4 (no solution). + checkSample(false, makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {})); + + // 1 <= 5x and 5x <= 9 (solution: x = 1). + checkSample(true, makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {})); + + // Bounded sets with equalities. + // x >= 8 and 40 >= y and x = y. + checkSample( + true, makeFACFromConstraints(2, {{1, 0, -8}, {0, -1, 40}}, {{1, -1, 0}})); + + // x <= 10 and y <= 10 and 10 <= z and x + 2y = 3z. + // solution: x = y = z = 10. + checkSample(true, makeFACFromConstraints( + 3, {{-1, 0, 0, 10}, {0, -1, 0, 10}, {0, 0, 1, -10}}, + {{1, 2, -3, 0}})); + + // x <= 10 and y <= 10 and 11 <= z and x + 2y = 3z. + // This implies x + 2y >= 33 and x + 2y <= 30, which has no solution. + checkSample(false, makeFACFromConstraints( + 3, {{-1, 0, 0, 10}, {0, -1, 0, 10}, {0, 0, 1, -11}}, + {{1, 2, -3, 0}})); + + // 0 <= r and r <= 3 and 4q + r = 7. + // Solution: q = 1, r = 3. + checkSample(true, + makeFACFromConstraints(2, {{0, 1, 0}, {0, -1, 3}}, {{4, 1, -7}})); + + // 4q + r = 7 and r = 0. + // Solution: q = 1, r = 3. + checkSample(false, makeFACFromConstraints(2, {}, {{4, 1, -7}, {0, 1, 0}})); + + // The next two sets are large sets that should take a long time to sample + // with a naive branch and bound algorithm but can be sampled efficiently with + // the GBR algroithm. + // + // This is a triangle with vertices at (1/3, 0), (2/3, 0) and (10000, 10000). + checkSample( + true, + makeFACFromConstraints( + 2, {{0, 1, 0}, {300000, -299999, -100000}, {-300000, 299998, 200000}}, + {})); + + // This is a tetrahedron with vertices at + // (1/3, 0, 0), (2/3, 0, 0), (2/3, 0, 10000), and (10000, 10000, 10000). + // The first three points form a triangular base on the xz plane with the + // apex at the fourth point, which is the only integer point. + checkPermutationsSample( + true, 3, + { + {0, 1, 0, 0}, // y >= 0 + {0, -1, 1, 0}, // z >= y + {300000, -299998, -1, + -100000}, // -300000x + 299998y + 100000 + z <= 0. + {-150000, 149999, 0, 100000}, // -150000x + 149999y + 100000 >= 0. + }, + {}); + + // Same thing with some spurious extra dimensions equated to constants. + checkSample(true, + makeFACFromConstraints( + 5, + { + {0, 1, 0, 1, -1, 0}, + {0, -1, 1, -1, 1, 0}, + {300000, -299998, -1, -9, 21, -112000}, + {-150000, 149999, 0, -15, 47, 68000}, + }, + {{0, 0, 0, 1, -1, 0}, // p = q. + {0, 0, 0, 1, 1, -2000}})); // p + q = 20000 => p = q = 10000. + + // This is a tetrahedron with vertices at + // (1/3, 0, 0), (2/3, 0, 0), (2/3, 0, 100), (100, 100 - 1/3, 100). + checkPermutationsSample(false, 3, + { + {0, 1, 0, 0}, + {0, -300, 299, 0}, + {300 * 299, -89400, -299, -100 * 299}, + {-897, 894, 0, 598}, + }, + {}); +} + +TEST(FlatAffineConstraintsTest, IsIntegerEmptyTest) { + // 1 <= 5x and 5x <= 4 (no solution). + EXPECT_TRUE( + makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}).isIntegerEmpty()); + // 1 <= 5x and 5x <= 9 (solution: x = 1). + EXPECT_FALSE( + makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}).isIntegerEmpty()); + + // An unbounded set, which isIntegerEmpty should detect as unbounded and + // return without calling findIntegerSample. + EXPECT_FALSE(makeFACFromConstraints(3, + { + {2, 0, 0, -1}, + {-2, 0, 0, 1}, + {0, 2, 0, -1}, + {0, -2, 0, 1}, + {0, 0, 2, -1}, + }, + {}) + .isIntegerEmpty()); +} + +} // namespace mlir diff --git a/mlir/unittests/Analysis/CMakeLists.txt b/mlir/unittests/Analysis/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/unittests/Analysis/CMakeLists.txt @@ -0,0 +1,8 @@ +add_mlir_unittest(MLIRAnalysisTests + AffineStructuresTest.cpp +) + +target_link_libraries(MLIRAnalysisTests + PRIVATE MLIRLoopAnalysis) + +add_subdirectory(Presburger) diff --git a/mlir/unittests/Analysis/Presburger/CMakeLists.txt b/mlir/unittests/Analysis/Presburger/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/unittests/Analysis/Presburger/CMakeLists.txt @@ -0,0 +1,7 @@ +add_mlir_unittest(MLIRPresburgerTests + MatrixTest.cpp + SimplexTest.cpp +) + +target_link_libraries(MLIRPresburgerTests + PRIVATE MLIRPresburger) diff --git a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp new file mode 100644 --- /dev/null +++ b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp @@ -0,0 +1,92 @@ +//===- MatrixTest.cpp - Tests for Matrix ----------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Analysis/Presburger/Matrix.h" +#include +#include + +namespace mlir { + +TEST(MatrixTest, ReadWrite) { + Matrix mat(5, 5); + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + mat(row, col) = 10 * row + col; + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + EXPECT_EQ(mat(row, col), int(10 * row + col)); +} + +TEST(MatrixTest, SwapColumns) { + Matrix mat(5, 5); + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + mat(row, col) = col == 3 ? 1 : 0; + mat.swapColumns(3, 1); + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + EXPECT_EQ(mat(row, col), col == 1 ? 1 : 0); + + // swap around all the other columns, swap (1, 3) twice for no effect. + mat.swapColumns(3, 1); + mat.swapColumns(2, 4); + mat.swapColumns(1, 3); + mat.swapColumns(0, 4); + mat.swapColumns(2, 2); + + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + EXPECT_EQ(mat(row, col), col == 1 ? 1 : 0); +} + +TEST(MatrixTest, SwapRows) { + Matrix mat(5, 5); + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + mat(row, col) = row == 2 ? 1 : 0; + mat.swapRows(2, 0); + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + EXPECT_EQ(mat(row, col), row == 0 ? 1 : 0); + + // swap around all the other rows, swap (2, 0) twice for no effect. + mat.swapRows(3, 4); + mat.swapRows(1, 4); + mat.swapRows(2, 0); + mat.swapRows(1, 1); + mat.swapRows(0, 2); + + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + EXPECT_EQ(mat(row, col), row == 0 ? 1 : 0); +} + +TEST(MatrixTest, resizeVertically) { + Matrix mat(5, 5); + EXPECT_EQ(mat.getNumRows(), 5u); + EXPECT_EQ(mat.getNumColumns(), 5u); + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + mat(row, col) = 10 * row + col; + + mat.resizeVertically(3); + EXPECT_EQ(mat.getNumRows(), 3u); + EXPECT_EQ(mat.getNumColumns(), 5u); + for (unsigned row = 0; row < 3; ++row) + for (unsigned col = 0; col < 5; ++col) + EXPECT_EQ(mat(row, col), int(10 * row + col)); + + mat.resizeVertically(5); + EXPECT_EQ(mat.getNumRows(), 5u); + EXPECT_EQ(mat.getNumColumns(), 5u); + for (unsigned row = 0; row < 5; ++row) + for (unsigned col = 0; col < 5; ++col) + EXPECT_EQ(mat(row, col), row >= 3 ? 0 : int(10 * row + col)); +} + +} // namespace mlir diff --git a/mlir/unittests/Analysis/Presburger/SimplexTest.cpp b/mlir/unittests/Analysis/Presburger/SimplexTest.cpp new file mode 100644 --- /dev/null +++ b/mlir/unittests/Analysis/Presburger/SimplexTest.cpp @@ -0,0 +1,175 @@ +//===- SimplexTest.cpp - Tests for Simplex --------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Analysis/Presburger/Simplex.h" + +#include +#include + +namespace mlir { + +TEST(SimplexTest, emptyRollback) { + Simplex tab(2); + // (u - v) >= 0 + tab.addInequality({1, -1, 0}); + EXPECT_FALSE(tab.isEmpty()); + + auto snap = tab.getSnapshot(); + // (u - v) <= -1 + tab.addInequality({-1, 1, -1}); + EXPECT_TRUE(tab.isEmpty()); + tab.rollback(snap); + EXPECT_FALSE(tab.isEmpty()); +} + +TEST(SimplexTest, addEquality_separate) { + Simplex tab(1); + tab.addInequality({1, -1}); // x >= 1 + ASSERT_FALSE(tab.isEmpty()); + tab.addEquality({1, 0}); // x == 0 + EXPECT_TRUE(tab.isEmpty()); +} + +void expectInequalityMakesTabEmpty(Simplex &tab, ArrayRef coeffs, + bool expect) { + ASSERT_FALSE(tab.isEmpty()); + auto snapshot = tab.getSnapshot(); + tab.addInequality(coeffs); + EXPECT_EQ(tab.isEmpty(), expect); + tab.rollback(snapshot); +} + +TEST(SimplexTest, addInequality_rollback) { + Simplex tab(3); + + std::vector coeffs[]{{1, 0, 0, 0}, // u >= 0 + {-1, 0, 0, 0}, // u <= 0 + {1, -1, 1, 0}, // u - v + w >= 0 + {1, 1, -1, 0}}; // u + v - w >= 0 + // The above constraints force u = 0 and v = w. + // The constraints below violate v = w. + std::vector checkCoeffs[]{{0, 1, -1, -1}, // v - w >= 1 + {0, -1, 1, -1}}; // v - w <= -1 + for (int run = 0; run < 4; run++) { + auto snapshot = tab.getSnapshot(); + + expectInequalityMakesTabEmpty(tab, checkCoeffs[0], false); + expectInequalityMakesTabEmpty(tab, checkCoeffs[1], false); + + for (int i = 0; i < 4; i++) + tab.addInequality(coeffs[(run + i) % 4]); + + expectInequalityMakesTabEmpty(tab, checkCoeffs[0], true); + expectInequalityMakesTabEmpty(tab, checkCoeffs[1], true); + + tab.rollback(snapshot); + EXPECT_EQ(tab.numConstraints(), 0u); + + expectInequalityMakesTabEmpty(tab, checkCoeffs[0], false); + expectInequalityMakesTabEmpty(tab, checkCoeffs[1], false); + } +} + +Simplex simplexFromInequalities(unsigned nDim, + SmallVector, 8> ineqs) { + Simplex simplex(nDim); + for (const auto &ineq : ineqs) + simplex.addInequality(ineq); + return simplex; +} + +TEST(SimplexTest, isUnbounded) { + EXPECT_FALSE(simplexFromInequalities( + 2, {{1, 1, 0}, {-1, -1, 0}, {1, -1, 5}, {-1, 1, -5}}) + .isUnbounded()); + + EXPECT_TRUE(simplexFromInequalities(2, {{1, 1, 0}, {1, -1, 5}, {-1, 1, -5}}) + .isUnbounded()); + + EXPECT_TRUE(simplexFromInequalities(2, {{-1, -1, 0}, {1, -1, 5}, {-1, 1, -5}}) + .isUnbounded()); + + EXPECT_TRUE(simplexFromInequalities(2, {}).isUnbounded()); + + EXPECT_FALSE(simplexFromInequalities(3, + { + {2, 0, 0, -1}, + {-2, 0, 0, 1}, + {0, 2, 0, -1}, + {0, -2, 0, 1}, + {0, 0, 2, -1}, + {0, 0, -2, 1}, + }) + .isUnbounded()); + + EXPECT_TRUE(simplexFromInequalities(3, + { + {2, 0, 0, -1}, + {-2, 0, 0, 1}, + {0, 2, 0, -1}, + {0, -2, 0, 1}, + {0, 0, -2, 1}, + }) + .isUnbounded()); + + EXPECT_TRUE(simplexFromInequalities(3, + { + {2, 0, 0, -1}, + {-2, 0, 0, 1}, + {0, 2, 0, -1}, + {0, -2, 0, 1}, + {0, 0, 2, -1}, + }) + .isUnbounded()); + + // Empty set. + EXPECT_FALSE(simplexFromInequalities(3, + { + {2, 0, 0, -1}, + {-2, 0, 0, 1}, + {0, 2, 2, -1}, + {0, -2, -2, 1}, + {3, 3, 3, -4}, + }) + .isUnbounded()); +} + +TEST(SimplexTest, getSamplePointIfIntegral) { + // Empty set. + EXPECT_FALSE(simplexFromInequalities(3, + { + {2, 0, 0, -1}, + {-2, 0, 0, 1}, + {0, 2, 2, -1}, + {0, -2, -2, 1}, + {3, 3, 3, -4}, + }) + .getSamplePointIfIntegral() + .hasValue()); + + auto maybeSample = simplexFromInequalities(2, + {// x = y - 2 + {1, -1, 2}, + {-1, 1, -2}, + // x + y = 2 + {1, 1, -2}, + {-1, -1, 2}}) + .getSamplePointIfIntegral(); + + EXPECT_TRUE(maybeSample.hasValue()); + EXPECT_THAT(*maybeSample, testing::ElementsAre(0, 2)); + + EXPECT_FALSE(simplexFromInequalities(1, + {// 2x = 1 (no integer solutions) + {2, -1}, + {-2, +1}}) + .getSamplePointIfIntegral() + .hasValue()); +} + +} // namespace mlir diff --git a/mlir/unittests/CMakeLists.txt b/mlir/unittests/CMakeLists.txt --- a/mlir/unittests/CMakeLists.txt +++ b/mlir/unittests/CMakeLists.txt @@ -5,6 +5,7 @@ add_unittest(MLIRUnitTests ${test_dirname} ${ARGN}) endfunction() +add_subdirectory(Analysis) add_subdirectory(Dialect) add_subdirectory(IR) add_subdirectory(Pass)