[MLIR] Matrix: support resizing horizontally

[Groverkss]

Added some comments.

mlir/lib/Analysis/Presburger/Matrix.cpp

69–72

Just to maintain the symmetry, this could be moved to a function resizeHorizontally().

[arjunp]

Add support for resizing just horizontally

[Groverkss]

LGTM.

[arjunp]

Add support for resizing just horizontally (repeat to try and resolve some issue with arc land)

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

public:
 
   /// 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.
+  /// the values are truncated; if it is bigger, the new values are initialized
+  /// to zero.
+  ///
+  /// Due to the representation of the matrix, resizing vertically (adding rows)
+  /// is less expensive than increasing the number of columns beyond
+  /// nReservedColumns.
+  void resize(unsigned newNRows, unsigned newNColumns);
+  void resizeHorizontally(unsigned newNColumns);
   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;

mlir/lib/Analysis/Presburger/Matrix.cpp

void Matrix::reserveRows(unsigned rows) {
   data.reserve(rows * nReservedColumns);
 }
 
 unsigned Matrix::appendExtraRow() {
   resizeVertically(nRows + 1);
   return nRows - 1;
 }
 
+void Matrix::resizeHorizontally(unsigned newNColumns) {
+  if (newNColumns < nColumns)
+    removeColumns(newNColumns, nColumns - newNColumns);
+  if (newNColumns > nColumns)
+    insertColumns(nColumns, newNColumns - nColumns);
+}
+
+void Matrix::resize(unsigned newNRows, unsigned newNColumns) {
+  resizeHorizontally(newNColumns);
+  resizeVertically(newNRows);
+}
+
 void Matrix::resizeVertically(unsigned newNRows) {
   nRows = newNRows;
   data.resize(nRows * nReservedColumns);
 }
 
 void Matrix::swapRows(unsigned row, unsigned otherRow) {
   assert((row < getNumRows() && otherRow < getNumRows()) &&
          "Given row out of bounds");

mlir/unittests/Analysis/Presburger/MatrixTest.cpp

TEST(MatrixTest, insertRows) {
 
   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);
 }
 
+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);
+  ASSERT_TRUE(mat.hasConsistentState());
+  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(7, 7);
+  ASSERT_TRUE(mat.hasConsistentState());
+  EXPECT_EQ(mat.getNumRows(), 7u);
+  EXPECT_EQ(mat.getNumColumns(), 7u);
+  for (unsigned row = 0; row < 7; ++row)
+    for (unsigned col = 0; col < 7; ++col)
+      EXPECT_EQ(mat(row, col), row >= 3 || col >= 3 ? 0 : int(10 * row + col));
+}
+
 } // namespace mlir