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,29 @@ /// 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; + /// Returns true if the set of constraints has no integer solution, and false + /// if it does have a solution. + bool isIntegerEmpty() const; + + /// Find a sample point satisfying the constraints. This uses a branch and + /// bound algorithm with general basis reduction. + /// + /// Returns such a point if one exists, or an empty llvm::Optional otherwise. + llvm::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,73 @@ +//===- 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 { + +struct Fraction { + Fraction() : num(0), den(1) {} + Fraction(int64_t oNum, int64_t oDen) : num(oNum), den(oDen) { + if (den < 0) { + num = -num; + den = -den; + } + } + int64_t num, den; +}; + +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,70 @@ +//===- 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 "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); + + static Matrix identity(unsigned dimension); + + /// Access the element at the specified row and column. + 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; + + llvm::ArrayRef getRow(unsigned row) const; + + 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 resize(unsigned newNRows, unsigned newNColumns); + + void print(llvm::raw_ostream &os) const; + void dump() const; + +private: + unsigned nColumns; + llvm::SmallVector, 32> 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,244 @@ +//===- 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 perform 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 "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 the Simplex algorithm. It supports adding affine +/// equalities and inequalities, 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. +/// +/// The implementation involves a matrix called a tableau, having number of rows +/// equal to the number of constraints and number of columns equal to two plus +/// the number of variables. Each row and column except the first two columns +/// correspond to an unknown -- either a constraint or a variable. +/// +/// Every row is a representation of the unknown x corresponding to that row in +/// terms of the unknowns in column positions u_1, u_2, ... u_m. If the row +/// corresponding to the unknown is d, c, a_1, a_2, ... a_m, then x is +/// represented as x = (c + a_1*u_1 + a_2*u_2 + ... + a_m*u_m)/d. +/// +/// An unknown is addressed by its index. If the index i is non-negative, then +/// the ith variable is the unknown being addressed. If the index is negative, +/// then a constraint is being addressed, having index ~i. +/// +/// The unknown corresponding to each row r (resp. column c) has index rowVar[r] +/// (resp. colVar[c]). If at some point it is detected that the +/// constraints are mutually contradictory and have no solution, then the +/// simplex will be marked empty; empty will be set to true. +class Simplex { +public: + enum class Direction { Up, Down }; + enum class UndoOp { Deallocate, UnmarkEmpty }; + + Simplex() = delete; + explicit Simplex(unsigned nVar); + explicit Simplex(const FlatAffineConstraints &constraints); + unsigned getNumRows() const; + unsigned getNumColumns() const; + + /// 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(llvm::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(llvm::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 a null value if no + /// optimum exists, i.e., if the expression is unbounded in this direction. + llvm::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. + llvm::Optional computeOptimum(Direction direction, + llvm::ArrayRef coeffs); + + /// 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 + /// empty llvm::Optional. + llvm::Optional> + getSamplePointIfIntegral() const; + + /// Returns an integral sample point if one exists, or an empty + /// llvm::Optional otherwise. + llvm::Optional> findIntegerSample(); + + // Print the tableau's internal state. + void print(llvm::raw_ostream &os) const; + void dump() const; + +private: + friend class GBRSimplex; + + struct Unknown { + Unknown(bool oOwnsRow, bool oRestricted, unsigned oPos) + : pos(oPos), ownsRow(oOwnsRow), restricted(oRestricted) {} + Unknown() : Unknown(false, false, -1) {} + unsigned pos; + bool ownsRow; + bool restricted; + }; + + // Dump the internal state of the unknown. + void printUnknown(llvm::raw_ostream &os, const Unknown &u) const; + + 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 llvm::Optional if + /// no valid pivot exists. + llvm::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(llvm::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. + bool restoreRow(Unknown &u); + + void undoOp(UndoOp op, llvm::Optional index); + + /// 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 llvm::Optional if the column + /// is unbounded in the specified direction. + llvm::Optional findPivotRow(llvm::Optional skipRow, + Direction direction, + unsigned col) const; + + /// Returns true value is positive and direction is Direction::UP, or if + /// value is negative and direction is Direction::DOWN. Returns false + /// otherwise. + bool signMatchesDirection(int64_t value, Direction direction) const; + + /// Returns Direction::UP if direction is Direction::DOWN and vice versa. + Direction flippedDirection(Direction direction) const; + + /// Searches for an integer sample point recursively using a branch and bound + /// algorithm and general basis reduction. + llvm::Optional> + findIntegerSampleRecursively(Matrix &basis, unsigned level); + + /// 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; + + /// 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. + llvm::SmallVector>, 32> undoLog; + + /// These hold the indexes of the unknown at a given row or column position. + llvm::SmallVector rowVar, colVar; + + /// These hold information about each unknown. + llvm::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,20 @@ return false; } +// First, try the GCD test heuristic. If that doesn't find the set empty, +// use the complete emptiness check provided by findIntegerSample, which gives +// a definitive answer. +bool FlatAffineConstraints::isIntegerEmpty() const { + if (isEmptyByGCDTest()) + return true; + return !findIntegerSample().hasValue(); +} + +llvm::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,83 @@ +//===- Matrix.cpp - 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Analysis/Presburger/Matrix.h" + +namespace mlir { + +Matrix::Matrix(unsigned rows, unsigned columns) + : nColumns(columns), data(rows, llvm::SmallVector(columns)) {} + +Matrix Matrix::identity(unsigned dimension) { + Matrix matrix(dimension, dimension); + for (size_t i = 0; i < dimension; ++i) + matrix(i, i) = 1; + return matrix; +} + +int64_t &Matrix::operator()(unsigned row, unsigned column) { + assert(row < getNumRows() && "Row outside of range"); + assert(column < getNumColumns() && "Column outside of range"); + return data[row][column]; +} + +int64_t Matrix::operator()(unsigned row, unsigned column) const { + assert(row < getNumRows() && "Row outside of range"); + assert(column < getNumColumns() && "Column outside of range"); + return data[row][column]; +} + +unsigned Matrix::getNumRows() const { return data.size(); } + +unsigned Matrix::getNumColumns() const { return nColumns; } + +void Matrix::resize(unsigned newNRows, unsigned newNColumns) { + nColumns = newNColumns; + for (auto &row : data) + row.resize(nColumns); + data.resize(newNRows, llvm::SmallVector(nColumns)); +} + +void Matrix::swapRows(unsigned row, unsigned otherRow) { + assert((row < getNumRows() && otherRow < getNumRows()) && + "Given row out of bounds"); + if (row == otherRow) + return; + std::swap(data[row], data[otherRow]); +} + +void Matrix::swapColumns(unsigned column, unsigned otherColumn) { + assert((column < getNumColumns() && otherColumn < getNumColumns()) && + "Given column out of bounds"); + if (column == otherColumn) + return; + for (auto &row : data) + std::swap(row[column], row[otherColumn]); +} + +llvm::ArrayRef Matrix::getRow(unsigned row) const { return data[row]; } + +void Matrix::addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale) { + if (scale == 0) + return; + for (unsigned col = 0; col < nColumns; ++col) + data[targetRow][col] += scale * data[sourceRow][col]; + return; +} + +void Matrix::print(llvm::raw_ostream &os) const { + for (const auto &row : data) { + for (int64_t elem : row) + os << elem << " "; + 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,934 @@ +//===- Simplex.cpp - 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Analysis/Presburger/Simplex.h" +#include "mlir/Analysis/Presburger/Matrix.h" +#include "mlir/Support/MathExtras.h" + +#include +#include + +namespace mlir { +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) { + colVar.push_back(nullIndex); + colVar.push_back(nullIndex); + for (unsigned i = 0; i < nVar; ++i) { + var.emplace_back( + /*ownsRow=*/false, /*restricted=*/false, /*pos=*/nCol); + colVar.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(colVar[col]); +} + +const Simplex::Unknown &Simplex::unknownFromRow(unsigned row) const { + assert(row < nRow && "Invalid row"); + return unknownFromIndex(rowVar[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(colVar[col]); +} + +Simplex::Unknown &Simplex::unknownFromRow(unsigned row) { + assert(row < nRow && "Invalid row"); + return unknownFromIndex(rowVar[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. +/// +/// If the tableau is not big enough to accomodate the extra row, we extend it. +/// +/// We then process each given variable coefficient. +/// 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. +/// +/// 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. +unsigned Simplex::addRow(llvm::ArrayRef coeffs) { + assert(coeffs.size() == 1 + var.size() && + "Incorrect number of coefficients!"); + + ++nRow; + if (nRow >= tableau.getNumRows()) + tableau.resize(nRow, tableau.getNumColumns()); + rowVar.push_back(~con.size()); + con.emplace_back(true, 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; + + for (unsigned i = 0; i < var.size(); ++i) { + unsigned pos = var[i].pos; + if (coeffs[i] == 0) + continue; + + if (!var[i].ownsRow) { + tableau(nRow - 1, pos) += coeffs[i] * tableau(nRow - 1, 0); + continue; + } + + 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 stack along with the index of the new constraint. + undoLog.emplace_back(UndoOp::Deallocate, ~(con.size() - 1)); + return con.size() - 1; +} + +unsigned Simplex::getNumColumns() const { return nCol; } + +unsigned Simplex::getNumRows() const { return nRow; } + +/// 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; +} + +bool Simplex::signMatchesDirection(int64_t elem, Direction direction) const { + assert(elem != 0 && "elem should not be 0"); + return direction == Direction::Up ? elem > 0 : elem < 0; +} + +Simplex::Direction Simplex::flippedDirection(Direction direction) const { + return direction == Direction::Up ? Direction::Down : Direction::Up; +} + +/// 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. +llvm::Optional Simplex::findPivot(int row, + Direction direction) const { + llvm::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 || colVar[j] < colVar[*col]) + col = j; + } + + if (!col) + return {}; + + Direction newDirection = + tableau(row, *col) < 0 ? flippedDirection(direction) : direction; + auto opt = findPivotRow(row, newDirection, *col); + return Pivot{opt.getValueOr(row), *col}; +} + +void Simplex::swapRowWithCol(unsigned row, unsigned col) { + std::swap(rowVar[row], colVar[col]); + Unknown *uCol = &unknownFromColumn(col); + Unknown *uRow = &unknownFromRow(row); + uCol->ownsRow = false; + uRow->ownsRow = true; + 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. +bool Simplex::restoreRow(Unknown &u) { + assert(u.ownsRow && "unknown should be in row position"); + + while (tableau(u.pos, 1) < 0) { + auto p = findPivot(u.pos, Direction::Up); + if (!p) + break; + + pivot(*p); + if (!u.ownsRow) + return 1; // the unknown is unbounded above. + } + return 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. +llvm::Optional Simplex::findPivotRow(llvm::Optional skipRow, + Direction direction, + unsigned col) const { + llvm::Optional retRow; + int64_t retElem, retConst; + for (unsigned row = 0; row < nRow; ++row) { + if (skipRow && row == *skipRow) + continue; + auto elem = tableau(row, col); + if (elem == 0) + continue; + if (!unknownFromRow(row).restricted) + continue; + if (signMatchesDirection(elem, direction)) + continue; + auto constTerm = tableau(row, 1); + + if (!retRow) { + retRow = row; + retElem = elem; + retConst = constTerm; + continue; + } + + auto diff = retConst * elem - constTerm * retElem; + if ((diff == 0 && rowVar[row] < rowVar[*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(rowVar[i], rowVar[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.emplace_back(UndoOp::UnmarkEmpty, llvm::Optional()); + 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(llvm::ArrayRef coeffs) { + unsigned conIndex = addRow(coeffs); + Unknown &u = con[conIndex]; + u.restricted = true; + bool success = restoreRow(u); + if (!success) + 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(llvm::ArrayRef coeffs) { + addInequality(coeffs); + SmallVector negatedCoeffs; + for (auto 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::undoOp(UndoOp op, llvm::Optional index) { + if (op == UndoOp::Deallocate) { + assert(index.hasValue() && + "DEALLOCATE undo entry must be accompanied by an index"); + + assert(index < 0 && "Unknown to be deallocated must be a constraint"); + Unknown &unknown = unknownFromIndex(*index); + + if (!unknown.ownsRow) { + unsigned column = unknown.pos; + llvm::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 column 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 (auto maybeRow = findPivotRow({}, Direction::Up, column)) + row = *maybeRow; + else if (auto 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(unknown.pos, nRow - 1); + // It is not strictly necessary to resize tableau, but for now we maintain + // the invariant that the actual size of the tableau is (nRow, nCol). + tableau.resize(nRow - 1, nCol); + nRow--; + rowVar.pop_back(); + con.pop_back(); + } else if (op == UndoOp::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) { + auto entry = undoLog.back(); + undoLog.pop_back(); + undoOp(entry.first, entry.second); + } +} + +llvm::Optional Simplex::computeRowOptimum(Direction direction, + unsigned row) { + while (auto maybePivot = findPivot(row, direction)) { + if (maybePivot->row == row) + return {}; + pivot(*maybePivot); + } + return Fraction(tableau(row, 1), tableau(row, 0)); +} + +llvm::Optional +Simplex::computeOptimum(Direction direction, llvm::ArrayRef coeffs) { + assert(!empty && "Tableau should not be empty"); + + auto snap = getSnapshot(); + unsigned conIndex = addRow(coeffs); + unsigned row = con[conIndex].pos; + auto optimum = computeRowOptimum(direction, row); + rollback(snap); + return optimum; +} + +/// 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.resize(numCon, 2 + numVar); + result.empty = a.empty || b.empty; + + auto concat = [](llvm::ArrayRef v, llvm::ArrayRef w) { + llvm::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.colVar.assign(2, nullIndex); + for (unsigned i = 2; i < a.nCol; i++) { + result.colVar.push_back(a.colVar[i]); + result.unknownFromIndex(result.colVar.back()).pos = + result.colVar.size() - 1; + } + for (unsigned i = 2; i < b.nCol; i++) { + result.colVar.push_back(indexFromBIndex(b.colVar[i])); + result.unknownFromIndex(result.colVar.back()).pos = + result.colVar.size() - 1; + } + + auto appendRowFromA = [&](unsigned row) { + for (unsigned col = 0; col < a.nCol; col++) + result.tableau(result.nRow, col) = a.tableau(row, col); + result.rowVar.push_back(a.rowVar[row]); + result.unknownFromIndex(result.rowVar.back()).pos = + result.rowVar.size() - 1; + result.nRow++; + }; + + // Also fixes the corresponding entry in rowVar and var/con. + auto appendRowFromB = [&](unsigned row) { + result.tableau(result.nRow, 0) = a.tableau(row, 0); + result.tableau(result.nRow, 1) = a.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.rowVar.push_back(indexFromBIndex(b.rowVar[row])); + result.unknownFromIndex(result.rowVar.back()).pos = + result.rowVar.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; +} + +llvm::Optional> +Simplex::getSamplePointIfIntegral() const { + if (empty) + return {}; + + llvm::SmallVector sample; + for (const Unknown &u : var) { + if (!u.ownsRow) + 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 general basis reduction. In what follows, +/// denotes the dot product of the vectors x and y. +class GBRSimplex { +public: + GBRSimplex(const Simplex &originalSimplex) + : simplex(Simplex::makeProduct(originalSimplex, originalSimplex)), + simplexConstraintOffset(simplex.numConstraints()) {} + + // Add an equality = 0. + void addEqualityForDirection(llvm::ArrayRef dir) { + snapshotStack.push_back(simplex.getSnapshot()); + simplex.addEquality(getCoeffsForDirection(dir)); + } + + // Compute max() and save the dual variables for only the + // direction constraints to `dual`. + Fraction computeWidthAndDuals(llvm::ArrayRef dir, + llvm::SmallVector &dual, + int64_t &dualDenom) { + unsigned snap = simplex.getSnapshot(); + unsigned conIndex = simplex.addRow(getCoeffsForDirection(dir)); + unsigned row = simplex.con[conIndex].pos; + auto width = simplex.computeRowOptimum(Simplex::Direction::Up, row); + assert(width && "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.) + // + // Otherwise, if the constraint is in row position, the dual is zero. + if (!simplex.con[i].ownsRow) + dual.push_back(-simplex.tableau(row, simplex.con[i].pos)); + else if (!simplex.con[i + 1].ownsRow) + dual.push_back(simplex.tableau(row, simplex.con[i + 1].pos)); + else + dual.push_back(0); + } + simplex.rollback(snap); + return *width; + } + + // Remove the last equality that was added through addEqualityForDirection. + void removeLastEquality() { + assert(!snapshotStack.empty() && "Snapshot stack is empty!"); + simplex.rollback(snapshotStack.back()); + snapshotStack.pop_back(); + } + +private: + // Returns coefficients for the expression . + llvm::SmallVector + getCoeffsForDirection(llvm::ArrayRef dir) { + assert(2 * dir.size() == simplex.numVariables() && + "Direction vector has wrong dimensionality"); + llvm::SmallVector coeffs; + for (unsigned i = 0; i < dir.size(); i++) + coeffs.emplace_back(dir[i]); + for (unsigned i = 0; i < dir.size(); i++) + coeffs.emplace_back(-dir[i]); + 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. + llvm::SmallVector snapshotStack; +}; + +/// Let b_{level}, b_{level + 1}, ... b_n be the current basis. +/// Let f_i(v) = max where x and y are points in the original +/// polytope and = 0 is satisfied for all j < i. (here +/// denotes the inner product) +/// +/// In every iteration, we first replace b_{i+1} with b_{i+1} + u*b_i, where u +/// is the integer such that f_i(b_{i+1} + u*b_i) minimized. Let alpha be the +/// dual variable associated with the constraint = 0 when computing +/// F_{i+1}(b_{i+1}). alpha must be the minimizing value of u, if it were +/// allowed to be rational. Due to convexity, the minimizing integer value is +/// either floor(alpha) or ceil(alpha), so we just need to check which of these +/// gives a lower F_{i+1} value. If alpha turned out to be an integer, then u = +/// alpha. +/// +/// Now if f_i(b_{i+1}) < eps * f_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 use eps = 0.75. +/// +/// In an iteration we need to compute: +/// +/// Some of the required F and alpha values may already be known. We cache the +/// known values and reuse them if possible. In particular: +/// +/// When we set b_{i+1} to b_{i+1} + u*b_i, no F values are changed since we +/// only added a multiple of b_i to b_{i+1}. In particular F_{i+1}(b_{i+1}) = +/// min f_i(b_{i+1} + alpha * b_i) so adding u*b_i to b_{i+1} does not change +/// this. Also = 0 and = 0 already imply = 0, so the constraints are unchanged. +/// +/// When we decrement i, we swap b_i and b_{i+1}. In the following paragraphs we +/// always refer to the final vector b_{i+1} after updating). But note that when +/// we come to the end of an iteration we always know f_i(b_{i+1}), so we just +/// need to update the cached value to reflect this. However the value of +/// F_{i+1}(b_i) is not already known, so if there was a stale value of F[i+1] +/// in the cache we remove this. Moreover, the iteration after decrementing will +/// want the dual variables from this computation so we cache this when we +/// compute the minimizing value of u. +/// +/// If alpha turns out to be an integer, then we never compute f_i(b_{i+1}) in +/// this iteration. But in this case, F_{i+1}(b_{i+1}) = F_{i+1}(b'_{i+1}) where +/// b'_{i+1} is the vector before updating. Therefore, we can update the cache +/// with this value. Furthermore, we can just inherit the dual variables from +/// this computation. +/// +/// When incrementing i we do not make any changes to the basis so no +/// invalidation occurs. +void Simplex::reduceBasis(Matrix &basis, unsigned level) { + const Fraction epsilon(3, 4); + + if (level == basis.getNumRows() - 1) + return; + + GBRSimplex gbrSimplex(*this); + llvm::SmallVector f; + llvm::SmallVector alpha; + int64_t alphaDenom; + auto findUAndGetFCandidate = [&](unsigned i) -> Fraction { + assert(i < level + alpha.size() && "alpha_i is not known!"); + + int64_t u = floorDiv(alpha[i - level], alphaDenom); + basis.addToRow(i, i + 1, u); + if (alpha[i - level] % alphaDenom != 0) { + llvm::SmallVector uAlpha[2]; + int64_t uAlphaDenom[2]; + Fraction fAtI[2]; + + // Initially u is floor(alpha) and basis reflects this. + fAtI[0] = gbrSimplex.computeWidthAndDuals(basis.getRow(i + 1), uAlpha[0], + uAlphaDenom[0]); + + // Now try ceil(alpha), i.e. floor(alpha) + 1. + ++u; + basis.addToRow(i, i + 1, 1); + fAtI[1] = gbrSimplex.computeWidthAndDuals(basis.getRow(i + 1), uAlpha[1], + uAlphaDenom[1]); + + int j = fAtI[0] < fAtI[1] ? 0 : 1; + if (j == 0) + // Subtract 1 to go from u = ceil(alpha) back to floor(alpha). + basis.addToRow(i, i + 1, -1); + + alpha = std::move(uAlpha[j]); + alphaDenom = uAlphaDenom[j]; + return fAtI[j]; + } + assert(i + 1 - level < f.size() && "F_{i+1} wasn't saved"); + // When alpha minimizes f_i(b_{i+1} + alpha*b_i), this is equal to + // F_{i+1}(b_{i+1}). + return f[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) { + Fraction fAtICandidate; // f_i(b_{i+1} + u*b_i) + if (i >= level + f.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 f already contains + // values up to and including i. + + assert((i == 0 || i - 1 < level + f.size()) && + "We are at level i but we don't know the value of F_{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"); + f.push_back( + gbrSimplex.computeWidthAndDuals(basis.getRow(i), alpha, alphaDenom)); + } + + if (i >= level + alpha.size()) { + assert(i + 1 >= level + f.size() && "We don't know alpha_i but we know " + "F_{i+1}, this should never happen"); + // We don't know alpha for our level, so let's find it. + gbrSimplex.addEqualityForDirection(basis.getRow(i)); + f.push_back(gbrSimplex.computeWidthAndDuals(basis.getRow(i + 1), alpha, + alphaDenom)); + gbrSimplex.removeLastEquality(); + } + + fAtICandidate = findUAndGetFCandidate(i); + + if (fAtICandidate < epsilon * f[i - level]) { + basis.swapRows(i, i + 1); + f[i - level] = fAtICandidate; + // The values of f_{i+1}(b_{i+1}) and higher may change after the swap, + // so we remove the cached values here. + f.resize(i - level + 1); + if (i == level) { + alpha.clear(); + continue; + } + + gbrSimplex.removeLastEquality(); + i--; + continue; + } + + alpha.clear(); + gbrSimplex.addEqualityForDirection(basis.getRow(i)); + i++; + } +} + +/// Try to find an integer sample point in the polytope. +/// +/// If such a point exists, this function returns it. Otherwise, it returns and +/// empty llvm::Optional. +llvm::Optional> Simplex::findIntegerSample() { + if (empty) + return {}; + + unsigned nDims = var.size(); + Matrix basis = Matrix::identity(nDims); + return findIntegerSampleRecursively(basis, 0); +} + +/// 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 called with 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 general basis reduction via +/// reduceBasis and then compute the bounds again. Now we can't do any better +/// than this, so we just recurse on every integer value in this range. +/// +/// 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 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, +/// or an empty llvm::Optional otherwise. +llvm::Optional> +Simplex::findIntegerSampleRecursively(Matrix &basis, unsigned level) { + if (level == basis.getNumRows()) + return getSamplePointIfIntegral(); + + auto getBounds = + [&](llvm::ArrayRef dir) -> std::pair { + llvm::SmallVector basisCoeffVector(dir.begin(), dir.end()); + basisCoeffVector.emplace_back(0); // constant term + int64_t minRoundedUp; + if (auto opt = computeOptimum(Direction::Down, basisCoeffVector)) + minRoundedUp = ceil(*opt); + else + llvm_unreachable("Tableau should not be unbounded"); + + int64_t maxRoundedDown; + if (auto opt = computeOptimum(Direction::Up, basisCoeffVector)) + maxRoundedDown = floor(*opt); + else + llvm_unreachable("Tableau should not be unbounded"); + + return {minRoundedUp, maxRoundedDown}; + }; + + int64_t minRoundedUp, maxRoundedDown; + std::tie(minRoundedUp, maxRoundedDown) = getBounds(basis.getRow(level)); + + // Heuristic: if the sample point is integral at this point, just return it. + if (auto opt = getSamplePointIfIntegral()) + return *opt; + + if (minRoundedUp < maxRoundedDown) { + reduceBasis(basis, level); + std::tie(minRoundedUp, maxRoundedDown) = getBounds(basis.getRow(level)); + } + + llvm::SmallVector basisCoeffVector(basis.getRow(level).begin(), + basis.getRow(level).end()); + basisCoeffVector.push_back(0); + for (int64_t i = minRoundedUp; i <= maxRoundedDown; ++i) { + auto snapshot = getSnapshot(); + basisCoeffVector.back() = -i; + // Add the constraint `basisCoeffVector = i`. + addEquality(basisCoeffVector); + if (auto opt = findIntegerSampleRecursively(basis, level + 1)) + return *opt; + rollback(snapshot); + } + + return {}; +} + +void Simplex::printUnknown(llvm::raw_ostream &os, const Unknown &u) const { + os << (u.ownsRow ? "r" : "c"); + os << u.pos; + if (u.restricted) + os << " [>=0]"; +} + +void Simplex::print(llvm::raw_ostream &os) const { + os << "Dumping Simplex, 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 << ", "; + printUnknown(os, var[i]); + } + os << "\ncon: "; + for (unsigned i = 0; i < con.size(); ++i) { + if (i > 0) + os << ", "; + printUnknown(os, con[i]); + } + os << '\n'; + for (unsigned row = 0; row < nRow; ++row) { + if (row > 0) + os << ", "; + os << "r" << row << ": " << rowVar[row]; + } + os << '\n'; + os << "c0: denom, c1: const"; + for (unsigned col = 2; col < nCol; ++col) + os << ", c" << col << ": " << colVar[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,132 @@ +//===- 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 + +namespace mlir { + +int64_t valueAt(ArrayRef expr, + const llvm::SmallVector &point) { + int64_t value = expr.back(); + assert(expr.size() == 1 + point.size()); + for (unsigned i = 0; i < point.size(); ++i) + value += expr[i] * point[i]; + return value; +} + +void checkSample(bool hasValue, const FlatAffineConstraints &fac) { + auto maybeSample = fac.findIntegerSample(); + if (!hasValue) + EXPECT_FALSE(maybeSample.hasValue()); + 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); + } +} + +FlatAffineConstraints makeFACFromConstraints( + unsigned dims, llvm::SmallVector, 8> ineqs, + llvm::SmallVector, 8> 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; +} + +TEST(FlatAffineConstraintsTest, FindSampleTest) { + // bounded sets, 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}}, + {})); + + // We run the following case with the variables in different orders, which + // causes the algorithm to run differently. At least one of these should get + // past our heuristic and actually test the basis reduction. + // + // This is a tetrahedron with vertices at + // (1/3, 0, 0), (2/3, 0, 0), (2/3, 0, 10000), and (10000, 10000, 10000). + SmallVector, 8> ineqs{ + {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 + }; + + SmallVector perm = {2, 0, 1}; + do { + SmallVector, 8> permutedIneqs; + for (const auto &ineq : ineqs) { + SmallVector permutedIneq; + for (auto id : perm) + permutedIneq.push_back(ineq[id]); + permutedIneq.push_back(ineq.back()); + permutedIneqs.emplace_back(permutedIneq); + } + checkSample(true, makeFACFromConstraints(3, permutedIneqs, {})); + } while (std::next_permutation(perm.begin(), perm.end())); +} + +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()); +} + +} // 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,90 @@ +//===- MatrixTest.cpp - Tests for Matrix ------------------------*- 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/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); + + 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(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, Resize) { + 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.resize(3, 3); + EXPECT_EQ(mat.getNumRows(), 3u); + EXPECT_EQ(mat.getNumColumns(), 3u); + for (unsigned row = 0; row < 3; ++row) + for (unsigned col = 0; col < 3; ++col) + EXPECT_EQ(mat(row, col), int(10 * row + col)); + + mat.resize(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) + EXPECT_EQ(mat(row, col), row >= 3 || col >= 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,78 @@ +//===- SimplexTest.cpp - Tests for Simplex ----------------------*- 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/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); + } +} + +} // 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)