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

[arjunp]

Summary

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.

[arjunp]

This patch is not expected to change any functionality, it's just some simplification in preparation for future work.

mlir/lib/Analysis/AffineStructures.cpp

147

Is there a reason the default copy constructor was not used?

[bondhugula]

In D105152#2848501, @arjunp wrote:

This patch is not expected to change any functionality, it's just some simplification in preparation for future work.

Thanks @arjunp Could you please add this line to the commit summary and add an additional line there on this patch itself?

[bondhugula]

This looks good in general. Some minor comments below.

mlir/include/mlir/Analysis/AffineStructures.h

70–76

Thanks for cleaning this up - this looks unrelated to the patch. This should at least be mentioned in the summary. "FlatAffineConstraints cleanup"?

88–90

Thanks for removing duplication here. Please also use the "/*....=*/" form wherever constant integers, etc. are being passed - for readability. I think numReservedRows/Cols was earlier uninitialized?

mlir/include/mlir/Analysis/Presburger/Matrix.h

65–66

It would be good to have a doc comment for these.

131

Doc comment please.

134

nReservedColumns needs a doc comment at least.

mlir/lib/Analysis/Presburger/Matrix.cpp

101

llvm math has this function: https://llvm.org/doxygen/MathExtras_8h_source.html

mlir/unittests/Analysis/Presburger/MatrixTest.cpp

94–115

A bit weird that these tests are under Analysis/Presburger. Is there a better place for them?

[bondhugula]

mlir/include/mlir/Analysis/Presburger/Matrix.h

67–101

All doc comments should be triple ///.

[mehdi_amini]

Nit: the commit title says "Marix" instead of (I think) "Matrix"

[arjunp]

Addressed review comments and fixed some bugs.

[arjunp]

Thanks for the quick review!

mlir/include/mlir/Analysis/AffineStructures.h

70–76

My primary reason to touch this code was to fix the usage of inequalities.reserve and the reference to nReservedColumns in accordance with the change to Matrix. I did slightly simplify the code since I was touching it anyway; I've mentioned this in the summary now.

mlir/unittests/Analysis/Presburger/MatrixTest.cpp

94–115

Not sure where else to put it since Matrix is under Analysis/Presburger too at the moment. I'm happy to have them both moved in another patch if there's a better place for them.

[arjunp]

Fixed another bug, the tests should pass now.

[arjunp]

Added regression tests for the bugs in the patch I fixed.

mlir/include/mlir/Analysis/AffineStructures.h

public:
 
   /// Constructs a constraint system reserving memory for the specified number
   /// of constraints and identifiers..
   FlatAffineConstraints(unsigned numReservedInequalities,
                         unsigned numReservedEqualities,
                         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 >= numDims + numSymbols + numLocals + 1);
+    assert(idArgs.empty() || idArgs.size() == numIds);
     ids.reserve(numReservedCols);
     if (idArgs.empty())
       ids.resize(numIds, None);
     else
       ids.append(idArgs.begin(), idArgs.end());
   }
 
   /// Constructs a constraint system with the specified number of
   /// dimensions and symbols.
   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(0, 0, numDims + numSymbols + numLocals + 1,
+                              numDims, numSymbols, numLocals, idArgs) {}
 
   /// Return a system with no constraints, i.e., one which is satisfied by all
   /// points.
   static FlatAffineConstraints getUniverse(unsigned numDims = 0,
                                            unsigned numSymbols = 0) {
     return FlatAffineConstraints(numDims, numSymbols);
   }
 
   /// Create a flat affine constraint system from an AffineValueMap or a list of
   /// these. The constructed system will only include equalities.
   explicit FlatAffineConstraints(const AffineValueMap &avm);
   explicit FlatAffineConstraints(ArrayRef<const AffineValueMap *> avmRef);
 
   /// Creates an affine constraint system from an IntegerSet.
   explicit FlatAffineConstraints(IntegerSet set);
 
-  FlatAffineConstraints(const FlatAffineConstraints &other);
-
   FlatAffineConstraints(ArrayRef<const AffineValueMap *> avmRef,
                         IntegerSet set);
 
   FlatAffineConstraints(const MutableAffineMap &map);
 
   ~FlatAffineConstraints() {}
 
   // Clears any existing data and reserves memory for the specified constraints.
public:
   /// Returns true if the given point satisfies the constraints, or false
   /// otherwise.
   bool containsPoint(ArrayRef<int64_t> point) const;
 
   // Clones this object.
   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'
   /// operation's Value using IR information stored in its bound maps. The
   /// right identifier is first looked up using forOp's Value. Asserts if the
   /// Value corresponding to the 'affine.for' operation isn't found in the
   /// constraint system. Returns failure for the yet unimplemented/unsupported
   /// cases.  Any new identifiers that are found in the bound operands of the
private:
   /// Normalized each constraints by the GCD of its coefficients.
   void normalizeConstraintsByGCD();
 
   /// Removes identifiers in the column range [idStart, idLimit), and copies any
   /// remaining valid data into place, updates member variables, and resizes
   /// 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;
 
   /// Number of identifiers corresponding to real dimensions.
   unsigned numDims;
 
   /// Number of identifiers corresponding to symbols (unknown but constant for
   /// 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
   /// None.
   SmallVector<Optional<Value>, 8> ids;
 
   /// A parameter that controls detection of an unrealistic number of
   /// constraints. If the number of constraints is this many times the number of

mlir/include/mlir/Analysis/Presburger/Matrix.h

 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <cassert>
 
 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. The reserved columns are actually present in
+/// the array, but are just not used. So the element at (i, j) is stored at
+/// data[i*nReservedColumns + j]. The reserved but unused columns 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);
 
   /// Access the element at the specified row and column.
   int64_t &at(unsigned row, unsigned column);
   int64_t at(unsigned row, unsigned column) const;
   int64_t &operator()(unsigned row, unsigned column);
   int64_t operator()(unsigned row, unsigned column) const;
 
   /// Swap the given columns.
   void swapColumns(unsigned column, unsigned otherColumn);
 
   /// Swap the given rows.
   void swapRows(unsigned row, unsigned otherRow);
 
   unsigned getNumRows() const;
 
   unsigned getNumColumns() const;
 
+  unsigned getNumReservedColumns() const;
+  unsigned getNumReservedRows() 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);
 
   /// Add `scale` multiples of the source column to the target column.
   void addToColumn(unsigned sourceColumn, unsigned targetColumn, int64_t scale);
 
   /// Negate the specified column.
   void negateColumn(unsigned column);
 
   /// Resize the matrix to the specified dimensions. If a dimension is smaller,
   /// the values are truncated; if it is bigger, the new values are default
   /// initialized.
   void resizeVertically(unsigned newNRows);
 
+  /// Clear the array. This sets the size to zero rows and zero columns.
+  void clear();
+
+  // 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;
 
+  bool hasConsistentState() const;
+
 private:
-  unsigned nRows, nColumns;
+  unsigned nRows, nColumns, nReservedColumns;
 
-  /// Stores the data. data.size() is equal to nRows * nColumns.
+  /// Stores the data. data.size() is equal to nRows * nReservedColumns.
+  /// data.capacity() / nReservedColumns is the number of reserved rows.
   SmallVector<int64_t, 64> data;
 };
 
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_MATRIX_H

mlir/lib/Analysis/AffineStructures.cpp

return ::getFlattenedAffineExprs(set.getConstraints(), set.getNumDims(),
                                    set.getNumSymbols(), flattenedExprs,
                                    localVarCst);
 }
 
 //===----------------------------------------------------------------------===//
 // 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));
-  }
-}
+// QUESTION: Is there a reason the default copy constructor was not used?
 
 // Clones this object.
 std::unique_ptr<FlatAffineConstraints> FlatAffineConstraints::clone() const {
   return std::make_unique<FlatAffineConstraints>(*this);
 }
 
 // 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.
   std::vector<SmallVector<int64_t, 8>> flatExprs;
   FlatAffineConstraints localVarCst;
   if (failed(getFlattenedAffineExprs(set, &flatExprs, &localVarCst))) {
     assert(false && "flattening unimplemented for semi-affine integer sets");
     return;

 void FlatAffineConstraints::reset(unsigned numReservedInequalities,
                                   unsigned numReservedEqualities,
                                   unsigned newNumReservedCols,
                                   unsigned newNumDims, unsigned newNumSymbols,
                                   unsigned newNumLocals,
                                   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.resize(numIds, None);
+  else
+    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,
                                   unsigned newNumLocals,
                                   ArrayRef<Value> idArgs) {
   reset(0, 0, newNumDims + newNumSymbols + newNumLocals + 1, newNumDims,
         newNumSymbols, newNumLocals, idArgs);
 }
 
 void FlatAffineConstraints::append(const FlatAffineConstraints &other) {
   assert(other.getNumCols() == getNumCols());
   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));
   }
   for (unsigned r = 0, e = other.getNumEqualities(); r < e; r++) {
     addEquality(other.getEquality(r));
   }
 }

 void FlatAffineConstraints::addId(IdKind kind, unsigned pos, Value id) {
   if (kind == IdKind::Dimension)
     assert(pos <= getNumDimIds());
   else if (kind == IdKind::Symbol)
     assert(pos <= getNumSymbolIds());
   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++;
   } else if (kind == IdKind::Symbol) {
     absolutePos = pos + getNumDimIds();
     numSymbols++;
   } else {
     absolutePos = pos + getNumDimIds() + getNumSymbolIds();
   }
   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)
     ids.insert(ids.begin() + absolutePos, id);
   else
     ids.insert(ids.begin() + absolutePos, None);
   assert(ids.size() == getNumIds());
 }
for (unsigned i = 0, e = getNumEqualities(); i < e; ++i) {
     normalizeConstraintByGCD</*isEq=*/true>(this, i);
   }
   for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
     normalizeConstraintByGCD</*isEq=*/false>(this, i);
   }
 }
 
 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;
 
   // Catches errors where numDims, numSymbols, numIds aren't consistent.
   if (numDims > numIds || numSymbols > numIds || numDims + numSymbols > numIds)
     return false;
 
if (j >= elimColStart && j < pivotCol)
       continue;
     int64_t v = pivotMultiplier * constraints->atEq(pivotRow, j) +
                 rowMultiplier * at(rowIdx, j);
     isEq ? constraints->atEq(rowIdx, j) = v
          : constraints->atIneq(rowIdx, j) = v;
   }
 }
 
-// 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) {
   assert(idLimit < getNumCols() && "invalid id limit");
 
   if (idStart >= idLimit)
     return;
 
   // We are going to be removing one or more identifiers from the range.
   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;
   unsigned numLocalsEliminated = 0;
   unsigned numColsEliminated = idLimit - idStart;
   if (idStart < numDims) {
     numDimsEliminated = std::min(numDims, idLimit) - idStart;
   }
   // Check how many local id's were removed. Note that our identifier order is
   // [dims, symbols, locals]. Local id start at position numDims + numSymbols.
   if (idLimit > numDims + numSymbols) {
     numLocalsEliminated = std::min(
         idLimit - std::max(idStart, numDims + numSymbols), getNumLocalIds());
   }
   unsigned numSymbolsEliminated =
       numColsEliminated - numDimsEliminated - numLocalsEliminated;
 
   numDims -= numDimsEliminated;
   numSymbols -= numSymbolsEliminated;
   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
 /// bounds> times <number of upper bounds> from the specified range of
 /// identifiers [start, end). It is often best to eliminate in the increasing
 /// order of these counts when doing Fourier-Motzkin elimination since FM adds
 /// that many new constraints.
 static unsigned getBestIdToEliminate(const FlatAffineConstraints &cst,
for (unsigned c = 0, e = getNumCols(); c < e; c++) {
       atIneq(dest, c) = atIneq(src, c);
     }
   };
   unsigned pos = 0;
   for (unsigned r = 0, e = getNumInequalities(); r < e; r++) {
     if (!redun[r])
       copyRow(r, pos++);
   }
-  inequalities.resize(numReservedCols * pos);
+  inequalities.resizeVertically(pos);
 }
 
 // A more complex check to eliminate redundant inequalities and equalities. Uses
 // Simplex to check if a constraint is redundant.
 void FlatAffineConstraints::removeRedundantConstraints() {
   // First, we run GCDTightenInequalities. This allows us to catch some
   // constraints which are not redundant when considering rational solutions
   // but are redundant in terms of integer solutions.
void FlatAffineConstraints::removeRedundantConstraints() {
   unsigned pos = 0;
   unsigned numIneqs = getNumInequalities();
   // Scan to get rid of all inequalities marked redundant, in-place. In Simplex,
   // the first constraints added are the inequalities.
   for (unsigned r = 0; r < numIneqs; r++) {
     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.
   // An equality is redundant if both the inequalities in its pair are
   // redundant.
   auto copyEquality = [&](unsigned src, unsigned dest) {
     if (src == dest)
       return;
     for (unsigned c = 0, e = getNumCols(); c < e; c++)
       atEq(dest, c) = atEq(src, c);
   };
   pos = 0;
   for (unsigned r = 0, e = getNumEqualities(); r < e; r++) {
     if (!(simplex.isMarkedRedundant(numIneqs + 2 * r) &&
           simplex.isMarkedRedundant(numIneqs + 2 * r + 1)))
       copyEquality(r, pos++);
   }
-  equalities.resize(numReservedCols * pos);
+  equalities.resizeVertically(pos);
 }
 
 std::pair<AffineMap, AffineMap> FlatAffineConstraints::getLowerAndUpperBound(
     unsigned pos, unsigned offset, unsigned num, unsigned symStartPos,
     ArrayRef<AffineExpr> localExprs, MLIRContext *context) const {
   assert(pos + offset < getNumDimIds() && "invalid dim start pos");
   assert(symStartPos >= (pos + offset) && "invalid sym start pos");
   assert(getNumLocalIds() == localExprs.size() &&
if (lbMap && lbMap.getNumResults() != 0 && ubMap &&
         return failure();
     }
   }
   return success();
 }
 
 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
 /// identifiers, the coefficients of which are provided in 'dividend' and with
 /// respect to a positive constant 'divisor'. Two constraints are added to the
 /// system to capture equivalence with the floordiv.
 ///      q = expr floordiv c    <=>   c*q <= expr <= c*q + c - 1.
 void FlatAffineConstraints::addLocalFloorDiv(ArrayRef<int64_t> dividend,
void FlatAffineConstraints::setDimSymbolSeparation(unsigned newSymbolCount) {
   assert(newSymbolCount <= numDims + numSymbols &&
          "invalid separation position");
   numDims = numDims + numSymbols - newSymbolCount;
   numSymbols = 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
 /// found.
 void FlatAffineConstraints::setIdToConstant(Value id, int64_t val) {
   unsigned pos;
   if (!findId(id, &pos))
     // This is a pre-condition for this method.
     assert(0 && "id not found");
   setIdToConstant(pos, 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.
 /// Returns the position of the equality row. If 'symbolic' is set to true,
 /// symbols are also treated like a constant, i.e., an affine function of the
 /// symbols is also treated like a constant. Returns -1 if such an equality
 /// could not be found.
 static int findEqualityToConstant(const FlatAffineConstraints &cst,
for (unsigned c = 0, e = getNumCols() - 1; c < e; c++) {
         return false;
     }
     return atIneq(r, getNumCols() - 1) >= 0;
   };
 
   // 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;
       continue;
     }
 
     // Among constraints that only differ in the constant term part, mark
     // everything other than the one with the smallest constant term redundant.
if (!ret.second) {
         val = {r, constTerm};
       } else {
         // The one stored makes this one redundant.
         redunIneq[r] = true;
       }
     }
   }
 
-  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.
 }
 
 void FlatAffineConstraints::clearAndCopyFrom(
     const FlatAffineConstraints &other) {
   FlatAffineConstraints copy(other);

mlir/lib/Analysis/Presburger/Matrix.cpp

 //===- Matrix.cpp - MLIR Matrix Class -------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/Matrix.h"
 
 namespace mlir {
 
-Matrix::Matrix(unsigned rows, unsigned columns)
-    : nRows(rows), nColumns(columns), data(nRows * nColumns) {}
+Matrix::Matrix(unsigned rows, unsigned columns, unsigned reservedRows,
+               unsigned reservedColumns)
+    : nRows(rows), nColumns(columns),
+      nReservedColumns(std::max(nColumns, reservedColumns)),
+      data(nRows * nReservedColumns) {
+  // QUESTION: is there no way to reserve in the constructor?
+  data.reserve(std::max(nRows, reservedRows) * nReservedColumns);
+}
 
 Matrix Matrix::identity(unsigned dimension) {
   Matrix matrix(dimension, dimension);
   for (unsigned i = 0; i < dimension; ++i)
     matrix(i, i) = 1;
   return matrix;
 }
 
 int64_t &Matrix::at(unsigned row, unsigned column) {
-  assert(row < getNumRows() && "Row outside of range");
-  assert(column < getNumColumns() && "Column outside of range");
-  return data[row * nColumns + column];
+  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) {
   return at(row, column);
 }
 
 int64_t Matrix::operator()(unsigned row, unsigned column) const {
   return at(row, column);
 }
 
 unsigned Matrix::getNumRows() const { return nRows; }
 
 unsigned Matrix::getNumColumns() const { return nColumns; }
 
+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::clear() {
+  nRows = nColumns = 0;
+  data.clear();
 }
 
 void Matrix::swapRows(unsigned row, unsigned otherRow) {
   assert((row < getNumRows() && otherRow < getNumRows()) &&
          "Given row out of bounds");
   if (row == otherRow)
     return;
   for (unsigned col = 0; col < nColumns; col++)
     std::swap(at(row, col), at(otherRow, col));
 }
 
 void Matrix::swapColumns(unsigned column, unsigned otherColumn) {
   assert((column < getNumColumns() && otherColumn < getNumColumns()) &&
          "Given column out of bounds");
   if (column == otherColumn)
     return;
   for (unsigned row = 0; row < nRows; row++)
     std::swap(at(row, column), at(row, otherColumn));
 }
 
 ArrayRef<int64_t> Matrix::getRow(unsigned row) const {
-  return {&data[row * nColumns], nColumns};
+  return {&data[row * nReservedColumns], nColumns};
+}
+
+namespace {
+unsigned nextPowOfTwo(unsigned x) {
+  unsigned y = 1;
+  while (y <= x)
+    y *= 2;
+  return y;
+}
+} // anonymous namespace
+
+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 = nextPowOfTwo(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 + count < nColumns; ++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); --r)
+    copyRow(r - count, r);
+
+  for (unsigned r = pos; r < pos + count; ++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) {
   if (scale == 0)
     return;
   for (unsigned col = 0; col < nColumns; ++col)
     at(targetRow, col) += scale * at(sourceRow, col);
-  return;
 }
 
 void Matrix::addToColumn(unsigned sourceColumn, unsigned targetColumn,
                          int64_t scale) {
   if (scale == 0)
     return;
   for (unsigned row = 0, e = getNumRows(); row < e; ++row)
     at(row, targetColumn) += scale * at(row, sourceColumn);
for (unsigned row = 0; row < nRows; ++row) {
     for (unsigned column = 0; column < nColumns; ++column)
       os << at(row, column) << ' ';
     os << '\n';
   }
 }
 
 void Matrix::dump() const { print(llvm::errs()); }
 
+bool Matrix::hasConsistentState() const {
+  if (data.size() % nReservedColumns != 0)
+    return false;
+  if (data.size() / nReservedColumns != nRows)
+    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

mlir/unittests/Analysis/Presburger/MatrixTest.cpp

TEST(MatrixTest, resizeVertically) {
   Matrix mat(5, 5);
   EXPECT_EQ(mat.getNumRows(), 5u);
   EXPECT_EQ(mat.getNumColumns(), 5u);
   for (unsigned row = 0; row < 5; ++row)
     for (unsigned col = 0; col < 5; ++col)
       mat(row, col) = 10 * row + col;
 
   mat.resizeVertically(3);
+  ASSERT_TRUE(mat.hasConsistentState());
   EXPECT_EQ(mat.getNumRows(), 3u);
   EXPECT_EQ(mat.getNumColumns(), 5u);
   for (unsigned row = 0; row < 3; ++row)
     for (unsigned col = 0; col < 5; ++col)
       EXPECT_EQ(mat(row, col), int(10 * row + col));
 
   mat.resizeVertically(5);
+  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)
       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, 2);
+  ASSERT_TRUE(mat.hasConsistentState());
+  EXPECT_EQ(mat.getNumRows(), 5u);
+  EXPECT_EQ(mat.getNumColumns(), 7u);
+  for (unsigned row = 0; row < 5; ++row) {
+    for (unsigned col = 0; col < 7; ++col) {
+      if (col < 3)
+        EXPECT_EQ(mat(row, col), int(10 * row + col));
+      else if (col == 3 || col == 4)
+        EXPECT_EQ(mat(row, col), 0);
+      else
+        EXPECT_EQ(mat(row, col), int(10 * row + col - 2));
+    }
+  }
+
+  mat.removeColumns(3, 2);
+  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, 2);
+  ASSERT_TRUE(mat.hasConsistentState());
+  EXPECT_EQ(mat.getNumRows(), 7u);
+  EXPECT_EQ(mat.getNumColumns(), 5u);
+  for (unsigned row = 0; row < 7; ++row) {
+    for (unsigned col = 0; col < 5; ++col) {
+      if (row < 3)
+        EXPECT_EQ(mat(row, col), int(10 * row + col));
+      else if (row == 3 || row == 4)
+        EXPECT_EQ(mat(row, col), 0);
+      else
+        EXPECT_EQ(mat(row, col), int(10 * (row - 2) + col));
+    }
+  }
+
+  mat.removeRows(3, 2);
+  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