[Mlir-commits] [mlir] c605dfc - [MLIR] FlatAffineConstraints: Use Matrix objects to store the constraints

Arjun P llvmlistbot at llvm.org
Thu Jul 1 09:35:37 PDT 2021


Author: Arjun P
Date: 2021-07-01T22:05:31+05:30
New Revision: c605dfcfc0c1c1605d254b00536351e2e8a87e19

URL: https://github.com/llvm/llvm-project/commit/c605dfcfc0c1c1605d254b00536351e2e8a87e19
DIFF: https://github.com/llvm/llvm-project/commit/c605dfcfc0c1c1605d254b00536351e2e8a87e19.diff

LOG: [MLIR] FlatAffineConstraints: Use Matrix objects to store the constraints

This results in significant deduplication of code. This patch is not expected to change any functionality, it's just some simplification in preparation for future work. Also slightly simplified some code that was being touched anyway and added some unit tests for some functions that were touched.

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D105152

Added: 
    

Modified: 
    mlir/include/mlir/Analysis/AffineStructures.h
    mlir/include/mlir/Analysis/Presburger/Matrix.h
    mlir/lib/Analysis/AffineStructures.cpp
    mlir/lib/Analysis/Presburger/Matrix.cpp
    mlir/unittests/Analysis/AffineStructuresTest.cpp
    mlir/unittests/Analysis/Presburger/MatrixTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Analysis/AffineStructures.h b/mlir/include/mlir/Analysis/AffineStructures.h
index 6aa0a38243b97..e96d52c225069 100644
--- a/mlir/include/mlir/Analysis/AffineStructures.h
+++ b/mlir/include/mlir/Analysis/AffineStructures.h
@@ -67,13 +67,12 @@ class FlatAffineConstraints {
                         unsigned numReservedCols, unsigned numDims = 0,
                         unsigned numSymbols = 0, unsigned numLocals = 0,
                         ArrayRef<Optional<Value>> idArgs = {})
-      : numReservedCols(numReservedCols), numDims(numDims),
-        numSymbols(numSymbols) {
-    assert(numReservedCols >= numDims + numSymbols + 1);
-    assert(idArgs.empty() || idArgs.size() == numDims + numSymbols + numLocals);
-    equalities.reserve(numReservedCols * numReservedEqualities);
-    inequalities.reserve(numReservedCols * numReservedInequalities);
-    numIds = numDims + numSymbols + numLocals;
+      : numIds(numDims + numSymbols + numLocals), numDims(numDims),
+        numSymbols(numSymbols),
+        equalities(0, numIds + 1, numReservedEqualities, numReservedCols),
+        inequalities(0, numIds + 1, numReservedInequalities, numReservedCols) {
+    assert(numReservedCols >= numIds + 1);
+    assert(idArgs.empty() || idArgs.size() == numIds);
     ids.reserve(numReservedCols);
     if (idArgs.empty())
       ids.resize(numIds, None);
@@ -86,17 +85,11 @@ class FlatAffineConstraints {
   FlatAffineConstraints(unsigned numDims = 0, unsigned numSymbols = 0,
                         unsigned numLocals = 0,
                         ArrayRef<Optional<Value>> idArgs = {})
-      : numReservedCols(numDims + numSymbols + numLocals + 1), numDims(numDims),
-        numSymbols(numSymbols) {
-    assert(numReservedCols >= numDims + numSymbols + 1);
-    assert(idArgs.empty() || idArgs.size() == numDims + numSymbols + numLocals);
-    numIds = numDims + numSymbols + numLocals;
-    ids.reserve(numIds);
-    if (idArgs.empty())
-      ids.resize(numIds, None);
-    else
-      ids.append(idArgs.begin(), idArgs.end());
-  }
+      : FlatAffineConstraints(/*numReservedInequalities=*/0,
+                              /*numReservedEqualities=*/0,
+                              /*numReservedCols=*/numDims + numSymbols +
+                                  numLocals + 1,
+                              numDims, numSymbols, numLocals, idArgs) {}
 
   /// Return a system with no constraints, i.e., one which is satisfied by all
   /// points.
@@ -113,8 +106,6 @@ class FlatAffineConstraints {
   /// Creates an affine constraint system from an IntegerSet.
   explicit FlatAffineConstraints(IntegerSet set);
 
-  FlatAffineConstraints(const FlatAffineConstraints &other);
-
   FlatAffineConstraints(ArrayRef<const AffineValueMap *> avmRef,
                         IntegerSet set);
 
@@ -173,51 +164,38 @@ class FlatAffineConstraints {
   std::unique_ptr<FlatAffineConstraints> clone() const;
 
   /// Returns the value at the specified equality row and column.
-  inline int64_t atEq(unsigned i, unsigned j) const {
-    return equalities[i * numReservedCols + j];
-  }
-  inline int64_t &atEq(unsigned i, unsigned j) {
-    return equalities[i * numReservedCols + j];
-  }
+  inline int64_t atEq(unsigned i, unsigned j) const { return equalities(i, j); }
+  inline int64_t &atEq(unsigned i, unsigned j) { return equalities(i, j); }
 
   inline int64_t atIneq(unsigned i, unsigned j) const {
-    return inequalities[i * numReservedCols + j];
+    return inequalities(i, j);
   }
 
-  inline int64_t &atIneq(unsigned i, unsigned j) {
-    return inequalities[i * numReservedCols + j];
-  }
+  inline int64_t &atIneq(unsigned i, unsigned j) { return inequalities(i, j); }
 
   /// Returns the number of columns in the constraint system.
   inline unsigned getNumCols() const { return numIds + 1; }
 
-  inline unsigned getNumEqualities() const {
-    assert(equalities.size() % numReservedCols == 0 &&
-           "inconsistent equality buffer size");
-    return equalities.size() / numReservedCols;
-  }
+  inline unsigned getNumEqualities() const { return equalities.getNumRows(); }
 
   inline unsigned getNumInequalities() const {
-    assert(inequalities.size() % numReservedCols == 0 &&
-           "inconsistent inequality buffer size");
-    return inequalities.size() / numReservedCols;
+    return inequalities.getNumRows();
   }
 
   inline unsigned getNumReservedEqualities() const {
-    return equalities.capacity() / numReservedCols;
+    return equalities.getNumReservedRows();
   }
 
   inline unsigned getNumReservedInequalities() const {
-    return inequalities.capacity() / numReservedCols;
+    return inequalities.getNumReservedRows();
   }
 
   inline ArrayRef<int64_t> getEquality(unsigned idx) const {
-    return ArrayRef<int64_t>(&equalities[idx * numReservedCols], getNumCols());
+    return equalities.getRow(idx);
   }
 
   inline ArrayRef<int64_t> getInequality(unsigned idx) const {
-    return ArrayRef<int64_t>(&inequalities[idx * numReservedCols],
-                             getNumCols());
+    return inequalities.getRow(idx);
   }
 
   /// Adds constraints (lower and upper bounds) for the specified 'affine.for'
@@ -649,16 +627,6 @@ class FlatAffineConstraints {
   /// arrays as needed.
   void removeIdRange(unsigned idStart, unsigned idLimit);
 
-  /// Coefficients of affine equalities (in == 0 form).
-  SmallVector<int64_t, 64> equalities;
-
-  /// Coefficients of affine inequalities (in >= 0 form).
-  SmallVector<int64_t, 64> inequalities;
-
-  /// Number of columns reserved. Actual ones in used are returned by
-  /// getNumCols().
-  unsigned numReservedCols;
-
   /// Total number of identifiers.
   unsigned numIds;
 
@@ -669,6 +637,12 @@ class FlatAffineConstraints {
   /// analysis).
   unsigned numSymbols;
 
+  /// Coefficients of affine equalities (in == 0 form).
+  Matrix equalities;
+
+  /// Coefficients of affine inequalities (in >= 0 form).
+  Matrix inequalities;
+
   /// Values corresponding to the (column) identifiers of this constraint
   /// system appearing in the order the identifiers correspond to columns.
   /// Temporary ones or those that aren't associated to any Value are set to

diff  --git a/mlir/include/mlir/Analysis/Presburger/Matrix.h b/mlir/include/mlir/Analysis/Presburger/Matrix.h
index 8ed40bb9c0266..0561aa920d145 100644
--- a/mlir/include/mlir/Analysis/Presburger/Matrix.h
+++ b/mlir/include/mlir/Analysis/Presburger/Matrix.h
@@ -22,16 +22,26 @@
 
 namespace mlir {
 
-/// This is a simple class to represent a resizable matrix.
+/// This is a class to represent a resizable matrix.
 ///
-/// The data is stored in the form of a vector of vectors.
+/// More columns and rows can be reserved than are currently used. The data is
+/// stored as a single 1D array, viewed as a 2D matrix with nRows rows and
+/// nReservedColumns columns, stored in row major form. Thus the element at
+/// (i, j) is stored at data[i*nReservedColumns + j]. The reserved but unused
+/// columns always have all zero values. The reserved rows are just reserved
+/// space in the underlying SmallVector's capacity.
 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);
+  /// The number of reserved rows and columns will be at least the number
+  /// specified, and will always be sufficient to accomodate the number of rows
+  /// and columns specified.
+  ///
+  /// Initially, the entries are initialized to ero.
+  Matrix(unsigned rows, unsigned columns, unsigned reservedRows = 0,
+         unsigned reservedColumns = 0);
 
   /// Return the identity matrix of the specified dimension.
   static Matrix identity(unsigned dimension);
@@ -52,9 +62,50 @@ class Matrix {
 
   unsigned getNumColumns() const;
 
+  /// Return the maximum number of rows/columns that can be added without
+  /// incurring a reallocation.
+  unsigned getNumReservedRows() const;
+  unsigned getNumReservedColumns() const;
+
+  /// Reserve enough space to resize to the specified number of rows without
+  /// reallocations.
+  void reserveRows(unsigned rows);
+
   /// Get an ArrayRef corresponding to the specified row.
   ArrayRef<int64_t> getRow(unsigned row) const;
 
+  /// Insert columns having positions pos, pos + 1, ... pos + count - 1.
+  /// Columns that were at positions 0 to pos - 1 will stay where they are;
+  /// columns that were at positions pos to nColumns - 1 will be pushed to the
+  /// right. pos should be at most nColumns.
+  void insertColumns(unsigned pos, unsigned count);
+  void insertColumn(unsigned pos);
+
+  /// Insert rows having positions pos, pos + 1, ... pos + count - 1.
+  /// Rows that were at positions 0 to pos - 1 will stay where they are;
+  /// rows that were at positions pos to nColumns - 1 will be pushed to the
+  /// right. pos should be at most nRows.
+  void insertRows(unsigned pos, unsigned count);
+  void insertRow(unsigned pos);
+
+  /// Remove the columns having positions pos, pos + 1, ... pos + count - 1.
+  /// Rows that were at positions 0 to pos - 1 will stay where they are;
+  /// columns that were at positions pos + count - 1 or later will be pushed to
+  /// the right. The columns to be deleted must be valid rows: pos + count - 1
+  /// must be at most nColumns - 1.
+  void removeColumns(unsigned pos, unsigned count);
+  void removeColumn(unsigned pos);
+
+  /// Remove the rows having positions pos, pos + 1, ... pos + count - 1.
+  /// Rows that were at positions 0 to pos - 1 will stay where they are;
+  /// rows that were at positions pos + count - 1 or later will be pushed to the
+  /// right. The rows to be deleted must be valid rows: pos + count - 1 must be
+  /// at most nRows - 1.
+  void removeRows(unsigned pos, unsigned count);
+  void removeRow(unsigned pos);
+
+  void copyRow(unsigned sourceRow, unsigned targetRow);
+
   /// Add `scale` multiples of the source row to the target row.
   void addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale);
 
@@ -69,14 +120,26 @@ class Matrix {
   /// initialized.
   void resizeVertically(unsigned newNRows);
 
+  /// Add an extra row at the bottom of the matrix and return its position.
+  unsigned appendExtraRow();
+
   /// Print the matrix.
   void print(raw_ostream &os) const;
   void dump() const;
 
-private:
-  unsigned nRows, nColumns;
+  /// Return whether the Matrix is in a consistent state with all its
+  /// invariants satisfied.
+  bool hasConsistentState() const;
 
-  /// Stores the data. data.size() is equal to nRows * nColumns.
+private:
+  /// The current number of rows, columns, and reserved columns. The underlying
+  /// data vector is viewed as an nRows x nReservedColumns matrix, of which the
+  /// first nColumns columns are currently in use, and the remaining are
+  /// reserved columns filled with zeros.
+  unsigned nRows, nColumns, nReservedColumns;
+
+  /// Stores the data. data.size() is equal to nRows * nReservedColumns.
+  /// data.capacity() / nReservedColumns is the number of reserved rows.
   SmallVector<int64_t, 64> data;
 };
 

diff  --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp
index 8a3b7b6b9a922..7fb82c86b1eed 100644
--- a/mlir/lib/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Analysis/AffineStructures.cpp
@@ -144,32 +144,6 @@ LogicalResult mlir::getFlattenedAffineExprs(
 // FlatAffineConstraints.
 //===----------------------------------------------------------------------===//
 
-// Copy constructor.
-FlatAffineConstraints::FlatAffineConstraints(
-    const FlatAffineConstraints &other) {
-  numReservedCols = other.numReservedCols;
-  numDims = other.getNumDimIds();
-  numSymbols = other.getNumSymbolIds();
-  numIds = other.getNumIds();
-
-  auto otherIds = other.getIds();
-  ids.reserve(numReservedCols);
-  ids.append(otherIds.begin(), otherIds.end());
-
-  unsigned numReservedEqualities = other.getNumReservedEqualities();
-  unsigned numReservedInequalities = other.getNumReservedInequalities();
-
-  equalities.reserve(numReservedEqualities * numReservedCols);
-  inequalities.reserve(numReservedInequalities * numReservedCols);
-
-  for (unsigned r = 0, e = other.getNumInequalities(); r < e; r++) {
-    addInequality(other.getInequality(r));
-  }
-  for (unsigned r = 0, e = other.getNumEqualities(); r < e; r++) {
-    addEquality(other.getEquality(r));
-  }
-}
-
 // Clones this object.
 std::unique_ptr<FlatAffineConstraints> FlatAffineConstraints::clone() const {
   return std::make_unique<FlatAffineConstraints>(*this);
@@ -177,11 +151,10 @@ std::unique_ptr<FlatAffineConstraints> FlatAffineConstraints::clone() const {
 
 // Construct from an IntegerSet.
 FlatAffineConstraints::FlatAffineConstraints(IntegerSet set)
-    : numReservedCols(set.getNumInputs() + 1),
-      numIds(set.getNumDims() + set.getNumSymbols()), numDims(set.getNumDims()),
-      numSymbols(set.getNumSymbols()) {
-  equalities.reserve(set.getNumEqualities() * numReservedCols);
-  inequalities.reserve(set.getNumInequalities() * numReservedCols);
+    : numIds(set.getNumDims() + set.getNumSymbols()), numDims(set.getNumDims()),
+      numSymbols(set.getNumSymbols()),
+      equalities(0, numIds + 1, set.getNumEqualities(), numIds + 1),
+      inequalities(0, numIds + 1, set.getNumInequalities(), numIds + 1) {
   ids.resize(numIds, None);
 
   // Flatten expressions and add them to the constraint system.
@@ -217,22 +190,13 @@ void FlatAffineConstraints::reset(unsigned numReservedInequalities,
                                   ArrayRef<Value> idArgs) {
   assert(newNumReservedCols >= newNumDims + newNumSymbols + newNumLocals + 1 &&
          "minimum 1 column");
-  numReservedCols = newNumReservedCols;
-  numDims = newNumDims;
-  numSymbols = newNumSymbols;
-  numIds = numDims + numSymbols + newNumLocals;
-  assert(idArgs.empty() || idArgs.size() == numIds);
+  SmallVector<Optional<Value>, 8> newIds;
+  if (!idArgs.empty())
+    newIds.assign(idArgs.begin(), idArgs.end());
 
-  clearConstraints();
-  if (numReservedEqualities >= 1)
-    equalities.reserve(newNumReservedCols * numReservedEqualities);
-  if (numReservedInequalities >= 1)
-    inequalities.reserve(newNumReservedCols * numReservedInequalities);
-  if (idArgs.empty()) {
-    ids.resize(numIds, None);
-  } else {
-    ids.assign(idArgs.begin(), idArgs.end());
-  }
+  *this = FlatAffineConstraints(numReservedInequalities, numReservedEqualities,
+                                newNumReservedCols, newNumDims, newNumSymbols,
+                                newNumLocals, newIds);
 }
 
 void FlatAffineConstraints::reset(unsigned newNumDims, unsigned newNumSymbols,
@@ -247,10 +211,9 @@ void FlatAffineConstraints::append(const FlatAffineConstraints &other) {
   assert(other.getNumDimIds() == getNumDimIds());
   assert(other.getNumSymbolIds() == getNumSymbolIds());
 
-  inequalities.reserve(inequalities.size() +
-                       other.getNumInequalities() * numReservedCols);
-  equalities.reserve(equalities.size() +
-                     other.getNumEqualities() * numReservedCols);
+  inequalities.reserveRows(inequalities.getNumRows() +
+                           other.getNumInequalities());
+  equalities.reserveRows(equalities.getNumRows() + other.getNumEqualities());
 
   for (unsigned r = 0, e = other.getNumInequalities(); r < e; r++) {
     addInequality(other.getInequality(r));
@@ -282,17 +245,7 @@ void FlatAffineConstraints::addId(IdKind kind, unsigned pos, Value id) {
   else
     assert(pos <= getNumLocalIds());
 
-  unsigned oldNumReservedCols = numReservedCols;
-
-  // Check if a resize is necessary.
-  if (getNumCols() + 1 > numReservedCols) {
-    equalities.resize(getNumEqualities() * (getNumCols() + 1));
-    inequalities.resize(getNumInequalities() * (getNumCols() + 1));
-    numReservedCols++;
-  }
-
   int absolutePos;
-
   if (kind == IdKind::Dimension) {
     absolutePos = pos;
     numDims++;
@@ -304,35 +257,8 @@ void FlatAffineConstraints::addId(IdKind kind, unsigned pos, Value id) {
   }
   numIds++;
 
-  // Note that getNumCols() now will already return the new size, which will be
-  // at least one.
-  int numInequalities = static_cast<int>(getNumInequalities());
-  int numEqualities = static_cast<int>(getNumEqualities());
-  int numCols = static_cast<int>(getNumCols());
-  for (int r = numInequalities - 1; r >= 0; r--) {
-    for (int c = numCols - 2; c >= 0; c--) {
-      if (c < absolutePos)
-        atIneq(r, c) = inequalities[r * oldNumReservedCols + c];
-      else
-        atIneq(r, c + 1) = inequalities[r * oldNumReservedCols + c];
-    }
-    atIneq(r, absolutePos) = 0;
-  }
-
-  for (int r = numEqualities - 1; r >= 0; r--) {
-    for (int c = numCols - 2; c >= 0; c--) {
-      // All values in column absolutePositions < absolutePos have the same
-      // coordinates in the 2-d view of the coefficient buffer.
-      if (c < absolutePos)
-        atEq(r, c) = equalities[r * oldNumReservedCols + c];
-      else
-        // Those at absolutePosition >= absolutePos, get a shifted
-        // absolutePosition.
-        atEq(r, c + 1) = equalities[r * oldNumReservedCols + c];
-    }
-    // Initialize added dimension to zero.
-    atEq(r, absolutePos) = 0;
-  }
+  inequalities.insertColumn(absolutePos);
+  equalities.insertColumn(absolutePos);
 
   // If an 'id' is provided, insert it; otherwise use None.
   if (id)
@@ -840,9 +766,9 @@ void FlatAffineConstraints::normalizeConstraintsByGCD() {
 }
 
 bool FlatAffineConstraints::hasConsistentState() const {
-  if (inequalities.size() != getNumInequalities() * numReservedCols)
+  if (!inequalities.hasConsistentState())
     return false;
-  if (equalities.size() != getNumEqualities() * numReservedCols)
+  if (!equalities.hasConsistentState())
     return false;
   if (ids.size() != getNumIds())
     return false;
@@ -923,31 +849,6 @@ static void eliminateFromConstraint(FlatAffineConstraints *constraints,
   }
 }
 
-// Remove coefficients in column range [colStart, colLimit) in place.
-// This removes in data in the specified column range, and copies any
-// remaining valid data into place.
-static void shiftColumnsToLeft(FlatAffineConstraints *constraints,
-                               unsigned colStart, unsigned colLimit,
-                               bool isEq) {
-  assert(colLimit <= constraints->getNumIds());
-  if (colLimit <= colStart)
-    return;
-
-  unsigned numCols = constraints->getNumCols();
-  unsigned numRows = isEq ? constraints->getNumEqualities()
-                          : constraints->getNumInequalities();
-  unsigned numToEliminate = colLimit - colStart;
-  for (unsigned r = 0, e = numRows; r < e; ++r) {
-    for (unsigned c = colLimit; c < numCols; ++c) {
-      if (isEq) {
-        constraints->atEq(r, c - numToEliminate) = constraints->atEq(r, c);
-      } else {
-        constraints->atIneq(r, c - numToEliminate) = constraints->atIneq(r, c);
-      }
-    }
-  }
-}
-
 // Removes identifiers in column range [idStart, idLimit), and copies any
 // remaining valid data into place, and updates member variables.
 void FlatAffineConstraints::removeIdRange(unsigned idStart, unsigned idLimit) {
@@ -960,11 +861,9 @@ void FlatAffineConstraints::removeIdRange(unsigned idStart, unsigned idLimit) {
   assert(idStart < numIds && "invalid idStart position");
 
   // TODO: Make 'removeIdRange' a lambda called from here.
-  // Remove eliminated identifiers from equalities.
-  shiftColumnsToLeft(this, idStart, idLimit, /*isEq=*/true);
-
-  // Remove eliminated identifiers from inequalities.
-  shiftColumnsToLeft(this, idStart, idLimit, /*isEq=*/false);
+  // Remove eliminated identifiers from the constraints..
+  equalities.removeColumns(idStart, idLimit - idStart);
+  inequalities.removeColumns(idStart, idLimit - idStart);
 
   // Update members numDims, numSymbols and numIds.
   unsigned numDimsEliminated = 0;
@@ -987,8 +886,6 @@ void FlatAffineConstraints::removeIdRange(unsigned idStart, unsigned idLimit) {
   numIds = numIds - numColsEliminated;
 
   ids.erase(ids.begin() + idStart, ids.begin() + idLimit);
-
-  // No resize necessary. numReservedCols remains the same.
 }
 
 /// Returns the position of the identifier that has the minimum <number of lower
@@ -1737,7 +1634,7 @@ void FlatAffineConstraints::removeRedundantInequalities() {
     if (!redun[r])
       copyRow(r, pos++);
   }
-  inequalities.resize(numReservedCols * pos);
+  inequalities.resizeVertically(pos);
 }
 
 // A more complex check to eliminate redundant inequalities and equalities. Uses
@@ -1764,7 +1661,7 @@ void FlatAffineConstraints::removeRedundantConstraints() {
     if (!simplex.isMarkedRedundant(r))
       copyInequality(r, pos++);
   }
-  inequalities.resize(numReservedCols * pos);
+  inequalities.resizeVertically(pos);
 
   // Scan to get rid of all equalities marked redundant, in-place. In Simplex,
   // after the inequalities, a pair of constraints for each equality is added.
@@ -1782,7 +1679,7 @@ void FlatAffineConstraints::removeRedundantConstraints() {
           simplex.isMarkedRedundant(numIneqs + 2 * r + 1)))
       copyEquality(r, pos++);
   }
-  equalities.resize(numReservedCols * pos);
+  equalities.resizeVertically(pos);
 }
 
 std::pair<AffineMap, AffineMap> FlatAffineConstraints::getLowerAndUpperBound(
@@ -2185,60 +2082,45 @@ LogicalResult FlatAffineConstraints::addSliceBounds(ArrayRef<Value> values,
 
 void FlatAffineConstraints::addEquality(ArrayRef<int64_t> eq) {
   assert(eq.size() == getNumCols());
-  unsigned offset = equalities.size();
-  equalities.resize(equalities.size() + numReservedCols);
-  std::copy(eq.begin(), eq.end(), equalities.begin() + offset);
+  unsigned row = equalities.appendExtraRow();
+  for (unsigned i = 0, e = eq.size(); i < e; ++i)
+    equalities(row, i) = eq[i];
 }
 
 void FlatAffineConstraints::addInequality(ArrayRef<int64_t> inEq) {
   assert(inEq.size() == getNumCols());
-  unsigned offset = inequalities.size();
-  inequalities.resize(inequalities.size() + numReservedCols);
-  std::copy(inEq.begin(), inEq.end(), inequalities.begin() + offset);
+  unsigned row = inequalities.appendExtraRow();
+  for (unsigned i = 0, e = inEq.size(); i < e; ++i)
+    inequalities(row, i) = inEq[i];
 }
 
 void FlatAffineConstraints::addConstantLowerBound(unsigned pos, int64_t lb) {
   assert(pos < getNumCols());
-  unsigned offset = inequalities.size();
-  inequalities.resize(inequalities.size() + numReservedCols);
-  std::fill(inequalities.begin() + offset,
-            inequalities.begin() + offset + getNumCols(), 0);
-  inequalities[offset + pos] = 1;
-  inequalities[offset + getNumCols() - 1] = -lb;
+  unsigned row = inequalities.appendExtraRow();
+  inequalities(row, pos) = 1;
+  inequalities(row, getNumCols() - 1) = -lb;
 }
 
 void FlatAffineConstraints::addConstantUpperBound(unsigned pos, int64_t ub) {
   assert(pos < getNumCols());
-  unsigned offset = inequalities.size();
-  inequalities.resize(inequalities.size() + numReservedCols);
-  std::fill(inequalities.begin() + offset,
-            inequalities.begin() + offset + getNumCols(), 0);
-  inequalities[offset + pos] = -1;
-  inequalities[offset + getNumCols() - 1] = ub;
+  unsigned row = inequalities.appendExtraRow();
+  inequalities(row, pos) = -1;
+  inequalities(row, getNumCols() - 1) = ub;
 }
 
 void FlatAffineConstraints::addConstantLowerBound(ArrayRef<int64_t> expr,
                                                   int64_t lb) {
-  assert(expr.size() == getNumCols());
-  unsigned offset = inequalities.size();
-  inequalities.resize(inequalities.size() + numReservedCols);
-  std::fill(inequalities.begin() + offset,
-            inequalities.begin() + offset + getNumCols(), 0);
-  std::copy(expr.begin(), expr.end(), inequalities.begin() + offset);
-  inequalities[offset + getNumCols() - 1] += -lb;
+  addInequality(expr);
+  inequalities(inequalities.getNumRows() - 1, getNumCols() - 1) += -lb;
 }
 
 void FlatAffineConstraints::addConstantUpperBound(ArrayRef<int64_t> expr,
                                                   int64_t ub) {
   assert(expr.size() == getNumCols());
-  unsigned offset = inequalities.size();
-  inequalities.resize(inequalities.size() + numReservedCols);
-  std::fill(inequalities.begin() + offset,
-            inequalities.begin() + offset + getNumCols(), 0);
-  for (unsigned i = 0, e = getNumCols(); i < e; i++) {
-    inequalities[offset + i] = -expr[i];
-  }
-  inequalities[offset + getNumCols() - 1] += ub;
+  unsigned row = inequalities.appendExtraRow();
+  for (unsigned i = 0, e = expr.size(); i < e; ++i)
+    inequalities(row, i) = -expr[i];
+  inequalities(inequalities.getNumRows() - 1, getNumCols() - 1) += ub;
 }
 
 /// Adds a new local identifier as the floordiv of an affine function of other
@@ -2311,12 +2193,10 @@ void FlatAffineConstraints::setDimSymbolSeparation(unsigned newSymbolCount) {
 
 /// Sets the specified identifier to a constant value.
 void FlatAffineConstraints::setIdToConstant(unsigned pos, int64_t val) {
-  unsigned offset = equalities.size();
-  equalities.resize(equalities.size() + numReservedCols);
-  std::fill(equalities.begin() + offset,
-            equalities.begin() + offset + getNumCols(), 0);
-  equalities[offset + pos] = 1;
-  equalities[offset + getNumCols() - 1] = -val;
+  equalities.resizeVertically(equalities.getNumRows() + 1);
+  unsigned row = equalities.getNumRows() - 1;
+  equalities(row, pos) = 1;
+  equalities(row, getNumCols() - 1) = -val;
 }
 
 /// Sets the specified identifier to a constant value; asserts if the id is not
@@ -2330,29 +2210,11 @@ void FlatAffineConstraints::setIdToConstant(Value id, int64_t val) {
 }
 
 void FlatAffineConstraints::removeEquality(unsigned pos) {
-  unsigned numEqualities = getNumEqualities();
-  assert(pos < numEqualities);
-  unsigned outputIndex = pos * numReservedCols;
-  unsigned inputIndex = (pos + 1) * numReservedCols;
-  unsigned numElemsToCopy = (numEqualities - pos - 1) * numReservedCols;
-  std::copy(equalities.begin() + inputIndex,
-            equalities.begin() + inputIndex + numElemsToCopy,
-            equalities.begin() + outputIndex);
-  assert(equalities.size() >= numReservedCols);
-  equalities.resize(equalities.size() - numReservedCols);
+  equalities.removeRow(pos);
 }
 
 void FlatAffineConstraints::removeInequality(unsigned pos) {
-  unsigned numInequalities = getNumInequalities();
-  assert(pos < numInequalities && "invalid position");
-  unsigned outputIndex = pos * numReservedCols;
-  unsigned inputIndex = (pos + 1) * numReservedCols;
-  unsigned numElemsToCopy = (numInequalities - pos - 1) * numReservedCols;
-  std::copy(inequalities.begin() + inputIndex,
-            inequalities.begin() + inputIndex + numElemsToCopy,
-            inequalities.begin() + outputIndex);
-  assert(inequalities.size() >= numReservedCols);
-  inequalities.resize(inequalities.size() - numReservedCols);
+  inequalities.removeRow(pos);
 }
 
 /// Finds an equality that equates the specified identifier to a constant.
@@ -2716,7 +2578,7 @@ void FlatAffineConstraints::removeTrivialRedundancy() {
   // Detect and mark redundant constraints.
   SmallVector<bool, 256> redunIneq(getNumInequalities(), false);
   for (unsigned r = 0, e = getNumInequalities(); r < e; r++) {
-    int64_t *rowStart = inequalities.data() + numReservedCols * r;
+    int64_t *rowStart = &inequalities(r, 0);
     auto row = ArrayRef<int64_t>(rowStart, getNumCols());
     if (isTriviallyValid(r) || !rowSet.insert(row).second) {
       redunIneq[r] = true;
@@ -2745,21 +2607,13 @@ void FlatAffineConstraints::removeTrivialRedundancy() {
     }
   }
 
-  auto copyRow = [&](unsigned src, unsigned dest) {
-    if (src == dest)
-      return;
-    for (unsigned c = 0, e = getNumCols(); c < e; c++) {
-      atIneq(dest, c) = atIneq(src, c);
-    }
-  };
-
   // Scan to get rid of all rows marked redundant, in-place.
   unsigned pos = 0;
-  for (unsigned r = 0, e = getNumInequalities(); r < e; r++) {
+  for (unsigned r = 0, e = getNumInequalities(); r < e; r++)
     if (!redunIneq[r])
-      copyRow(r, pos++);
-  }
-  inequalities.resize(numReservedCols * pos);
+      inequalities.copyRow(r, pos++);
+
+  inequalities.resizeVertically(pos);
 
   // TODO: consider doing this for equalities as well, but probably not worth
   // the savings.
@@ -3053,8 +2907,8 @@ void FlatAffineConstraints::projectOut(Value id) {
 }
 
 void FlatAffineConstraints::clearConstraints() {
-  equalities.clear();
-  inequalities.clear();
+  equalities.resizeVertically(0);
+  inequalities.resizeVertically(0);
 }
 
 namespace {

diff  --git a/mlir/lib/Analysis/Presburger/Matrix.cpp b/mlir/lib/Analysis/Presburger/Matrix.cpp
index 4a5a53921548c..1b5861012d3cf 100644
--- a/mlir/lib/Analysis/Presburger/Matrix.cpp
+++ b/mlir/lib/Analysis/Presburger/Matrix.cpp
@@ -7,11 +7,17 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/Matrix.h"
+#include "llvm/Support/MathExtras.h"
 
 namespace mlir {
 
-Matrix::Matrix(unsigned rows, unsigned columns)
-    : nRows(rows), nColumns(columns), data(nRows * nColumns) {}
+Matrix::Matrix(unsigned rows, unsigned columns, unsigned reservedRows,
+               unsigned reservedColumns)
+    : nRows(rows), nColumns(columns),
+      nReservedColumns(std::max(nColumns, reservedColumns)),
+      data(nRows * nReservedColumns) {
+  data.reserve(std::max(nRows, reservedRows) * nReservedColumns);
+}
 
 Matrix Matrix::identity(unsigned dimension) {
   Matrix matrix(dimension, dimension);
@@ -21,15 +27,15 @@ Matrix Matrix::identity(unsigned dimension) {
 }
 
 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];
+  assert(row < nRows && "Row outside of range");
+  assert(column < nColumns && "Column outside of range");
+  return data[row * nReservedColumns + 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];
+  assert(row < nRows && "Row outside of range");
+  assert(column < nColumns && "Column outside of range");
+  return data[row * nReservedColumns + column];
 }
 
 int64_t &Matrix::operator()(unsigned row, unsigned column) {
@@ -44,9 +50,24 @@ unsigned Matrix::getNumRows() const { return nRows; }
 
 unsigned Matrix::getNumColumns() const { return nColumns; }
 
+unsigned Matrix::getNumReservedColumns() const { return nReservedColumns; }
+
+unsigned Matrix::getNumReservedRows() const {
+  return data.capacity() / nReservedColumns;
+}
+
+void Matrix::reserveRows(unsigned rows) {
+  data.reserve(rows * nReservedColumns);
+}
+
+unsigned Matrix::appendExtraRow() {
+  resizeVertically(nRows + 1);
+  return nRows - 1;
+}
+
 void Matrix::resizeVertically(unsigned newNRows) {
   nRows = newNRows;
-  data.resize(nRows * nColumns);
+  data.resize(nRows * nReservedColumns);
 }
 
 void Matrix::swapRows(unsigned row, unsigned otherRow) {
@@ -68,7 +89,81 @@ void Matrix::swapColumns(unsigned column, unsigned otherColumn) {
 }
 
 ArrayRef<int64_t> Matrix::getRow(unsigned row) const {
-  return {&data[row * nColumns], nColumns};
+  return {&data[row * nReservedColumns], nColumns};
+}
+
+void Matrix::insertColumn(unsigned pos) { insertColumns(pos, 1); }
+void Matrix::insertColumns(unsigned pos, unsigned count) {
+  if (count == 0)
+    return;
+  assert(pos <= nColumns);
+  unsigned oldNReservedColumns = nReservedColumns;
+  if (nColumns + count > nReservedColumns) {
+    nReservedColumns = llvm::NextPowerOf2(nColumns + count);
+    data.resize(nRows * nReservedColumns);
+  }
+  nColumns += count;
+
+  for (int ri = nRows - 1; ri >= 0; --ri) {
+    for (int ci = nReservedColumns - 1; ci >= 0; --ci) {
+      unsigned r = ri;
+      unsigned c = ci;
+      int64_t &dest = data[r * nReservedColumns + c];
+      if (c >= nColumns)
+        dest = 0;
+      else if (c >= pos + count)
+        dest = data[r * oldNReservedColumns + c - count];
+      else if (c >= pos)
+        dest = 0;
+      else
+        dest = data[r * oldNReservedColumns + c];
+    }
+  }
+}
+
+void Matrix::removeColumn(unsigned pos) { removeColumns(pos, 1); }
+void Matrix::removeColumns(unsigned pos, unsigned count) {
+  if (count == 0)
+    return;
+  assert(pos + count - 1 < nColumns);
+  for (unsigned r = 0; r < nRows; ++r) {
+    for (unsigned c = pos; c < nColumns - count; ++c)
+      at(r, c) = at(r, c + count);
+    for (unsigned c = nColumns - count; c < nColumns; ++c)
+      at(r, c) = 0;
+  }
+  nColumns -= count;
+}
+
+void Matrix::insertRow(unsigned pos) { insertRows(pos, 1); }
+void Matrix::insertRows(unsigned pos, unsigned count) {
+  if (count == 0)
+    return;
+
+  assert(pos <= nRows);
+  resizeVertically(nRows + count);
+  for (int r = nRows - 1; r >= int(pos + count); --r)
+    copyRow(r - count, r);
+  for (int r = pos + count - 1; r >= int(pos); --r)
+    for (unsigned c = 0; c < nColumns; ++c)
+      at(r, c) = 0;
+}
+
+void Matrix::removeRow(unsigned pos) { removeRows(pos, 1); }
+void Matrix::removeRows(unsigned pos, unsigned count) {
+  if (count == 0)
+    return;
+  assert(pos + count - 1 <= nRows);
+  for (unsigned r = pos; r + count < nRows; ++r)
+    copyRow(r + count, r);
+  resizeVertically(nRows - count);
+}
+
+void Matrix::copyRow(unsigned sourceRow, unsigned targetRow) {
+  if (sourceRow == targetRow)
+    return;
+  for (unsigned c = 0; c < nColumns; ++c)
+    at(targetRow, c) = at(sourceRow, c);
 }
 
 void Matrix::addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale) {
@@ -76,7 +171,6 @@ void Matrix::addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale) {
     return;
   for (unsigned col = 0; col < nColumns; ++col)
     at(targetRow, col) += scale * at(sourceRow, col);
-  return;
 }
 
 void Matrix::addToColumn(unsigned sourceColumn, unsigned targetColumn,
@@ -102,4 +196,16 @@ void Matrix::print(raw_ostream &os) const {
 
 void Matrix::dump() const { print(llvm::errs()); }
 
+bool Matrix::hasConsistentState() const {
+  if (data.size() != nRows * nReservedColumns)
+    return false;
+  if (nColumns > nReservedColumns)
+    return false;
+  for (unsigned r = 0; r < nRows; ++r)
+    for (unsigned c = nColumns; c < nReservedColumns; ++c)
+      if (data[r * nReservedColumns + c] != 0)
+        return false;
+  return true;
+}
+
 } // namespace mlir

diff  --git a/mlir/unittests/Analysis/AffineStructuresTest.cpp b/mlir/unittests/Analysis/AffineStructuresTest.cpp
index 2121fec978d88..3ee6f049c9a20 100644
--- a/mlir/unittests/Analysis/AffineStructuresTest.cpp
+++ b/mlir/unittests/Analysis/AffineStructuresTest.cpp
@@ -547,4 +547,44 @@ TEST(FlatAffineConstraintsTest, removeRedundantConstraintsTest) {
   }
 }
 
+TEST(FlatAffineConstraintsTest, addConstantUpperBound) {
+  FlatAffineConstraints fac = makeFACFromConstraints(2, {}, {});
+  fac.addConstantUpperBound(0, 1);
+  EXPECT_EQ(fac.atIneq(0, 0), -1);
+  EXPECT_EQ(fac.atIneq(0, 1), 0);
+  EXPECT_EQ(fac.atIneq(0, 2), 1);
+
+  fac.addConstantUpperBound({1, 2, 3}, 1);
+  EXPECT_EQ(fac.atIneq(1, 0), -1);
+  EXPECT_EQ(fac.atIneq(1, 1), -2);
+  EXPECT_EQ(fac.atIneq(1, 2), -2);
+}
+
+TEST(FlatAffineConstraintsTest, addConstantLowerBound) {
+  FlatAffineConstraints fac = makeFACFromConstraints(2, {}, {});
+  fac.addConstantLowerBound(0, 1);
+  EXPECT_EQ(fac.atIneq(0, 0), 1);
+  EXPECT_EQ(fac.atIneq(0, 1), 0);
+  EXPECT_EQ(fac.atIneq(0, 2), -1);
+
+  fac.addConstantLowerBound({1, 2, 3}, 1);
+  EXPECT_EQ(fac.atIneq(1, 0), 1);
+  EXPECT_EQ(fac.atIneq(1, 1), 2);
+  EXPECT_EQ(fac.atIneq(1, 2), 2);
+}
+
+TEST(FlatAffineConstraintsTest, clearConstraints) {
+  FlatAffineConstraints fac = makeFACFromConstraints(1, {}, {});
+
+  fac.addInequality({1, 0});
+  EXPECT_EQ(fac.atIneq(0, 0), 1);
+  EXPECT_EQ(fac.atIneq(0, 1), 0);
+
+  fac.clearConstraints();
+
+  fac.addInequality({1, 0});
+  EXPECT_EQ(fac.atIneq(0, 0), 1);
+  EXPECT_EQ(fac.atIneq(0, 1), 0);
+}
+
 } // namespace mlir

diff  --git a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
index 4d8801b579a79..9c892f73249e1 100644
--- a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
@@ -75,6 +75,7 @@ TEST(MatrixTest, resizeVertically) {
       mat(row, col) = 10 * row + col;
 
   mat.resizeVertically(3);
+  ASSERT_TRUE(mat.hasConsistentState());
   EXPECT_EQ(mat.getNumRows(), 3u);
   EXPECT_EQ(mat.getNumColumns(), 5u);
   for (unsigned row = 0; row < 3; ++row)
@@ -82,6 +83,7 @@ TEST(MatrixTest, resizeVertically) {
       EXPECT_EQ(mat(row, col), int(10 * row + col));
 
   mat.resizeVertically(5);
+  ASSERT_TRUE(mat.hasConsistentState());
   EXPECT_EQ(mat.getNumRows(), 5u);
   EXPECT_EQ(mat.getNumColumns(), 5u);
   for (unsigned row = 0; row < 5; ++row)
@@ -89,4 +91,79 @@ TEST(MatrixTest, resizeVertically) {
       EXPECT_EQ(mat(row, col), row >= 3 ? 0 : int(10 * row + col));
 }
 
+TEST(MatrixTest, insertColumns) {
+  Matrix mat(5, 5, 5, 10);
+  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.insertColumns(3, 100);
+  ASSERT_TRUE(mat.hasConsistentState());
+  EXPECT_EQ(mat.getNumRows(), 5u);
+  EXPECT_EQ(mat.getNumColumns(), 105u);
+  for (unsigned row = 0; row < 5; ++row) {
+    for (unsigned col = 0; col < 105; ++col) {
+      if (col < 3)
+        EXPECT_EQ(mat(row, col), int(10 * row + col));
+      else if (3 <= col && col <= 102)
+        EXPECT_EQ(mat(row, col), 0);
+      else
+        EXPECT_EQ(mat(row, col), int(10 * row + col - 100));
+    }
+  }
+
+  mat.removeColumns(3, 100);
+  ASSERT_TRUE(mat.hasConsistentState());
+  mat.insertColumns(0, 0);
+  ASSERT_TRUE(mat.hasConsistentState());
+  mat.insertColumn(5);
+  ASSERT_TRUE(mat.hasConsistentState());
+
+  EXPECT_EQ(mat.getNumRows(), 5u);
+  EXPECT_EQ(mat.getNumColumns(), 6u);
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 6; ++col)
+      EXPECT_EQ(mat(row, col), col == 5 ? 0 : 10 * row + col);
+}
+
+TEST(MatrixTest, insertRows) {
+  Matrix mat(5, 5, 5, 10);
+  ASSERT_TRUE(mat.hasConsistentState());
+  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.insertRows(3, 100);
+  ASSERT_TRUE(mat.hasConsistentState());
+  EXPECT_EQ(mat.getNumRows(), 105u);
+  EXPECT_EQ(mat.getNumColumns(), 5u);
+  for (unsigned row = 0; row < 105; ++row) {
+    for (unsigned col = 0; col < 5; ++col) {
+      if (row < 3)
+        EXPECT_EQ(mat(row, col), int(10 * row + col));
+      else if (3 <= row && row <= 102)
+        EXPECT_EQ(mat(row, col), 0);
+      else
+        EXPECT_EQ(mat(row, col), int(10 * (row - 100) + col));
+    }
+  }
+
+  mat.removeRows(3, 100);
+  ASSERT_TRUE(mat.hasConsistentState());
+  mat.insertRows(0, 0);
+  ASSERT_TRUE(mat.hasConsistentState());
+  mat.insertRow(5);
+  ASSERT_TRUE(mat.hasConsistentState());
+
+  EXPECT_EQ(mat.getNumRows(), 6u);
+  EXPECT_EQ(mat.getNumColumns(), 5u);
+  for (unsigned row = 0; row < 6; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      EXPECT_EQ(mat(row, col), row == 5 ? 0 : 10 * row + col);
+}
+
 } // namespace mlir


        


More information about the Mlir-commits mailing list