[MLIR] Matrix: support matrix-vector multiplication

[arjunp]

Summary

This just moves in the implementation from LinearTransform.

[Groverkss]

LGTM.

[bondhugula]

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

120–125

Triple ///.

[arjunp]

Address Uday's comment

mlir/include/mlir/Analysis/Presburger/LinearTransform.h

public:
   makeTransformToColumnEchelon(Matrix m);
 
   // Returns an IntegerPolyhedron having a constraint vector vT for every
   // constraint vector v in poly, where T is this transform.
   IntegerPolyhedron applyTo(const IntegerPolyhedron &poly) const;
 
   // The given vector is interpreted as a row vector v. Post-multiply v with
   // this transform, say T, and return vT.
-  SmallVector<int64_t, 8> preMultiplyWithRow(ArrayRef<int64_t> rowVec) const;
+  SmallVector<int64_t, 8> preMultiplyWithRow(ArrayRef<int64_t> rowVec) const {
+    return matrix.preMultiplyWithRow(rowVec);
+  }
 
   // The given vector is interpreted as a column vector v. Pre-multiply v with
   // this transform, say T, and return Tv.
   SmallVector<int64_t, 8>
-  postMultiplyWithColumn(ArrayRef<int64_t> colVec) const;
+  postMultiplyWithColumn(ArrayRef<int64_t> colVec) const {
+    return matrix.postMultiplyWithColumn(colVec);
+  }
 
 private:
   Matrix matrix;
 };
 
 } // namespace mlir
 #endif // MLIR_ANALYSIS_PRESBURGER_LINEARTRANSFORM_H

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

public:
   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);
 
+  /// The given vector is interpreted as a row vector v. Post-multiply v with
+  /// this matrix, say M, and return vM.
+  SmallVector<int64_t, 8> preMultiplyWithRow(ArrayRef<int64_t> rowVec) const;
+
+  /// The given vector is interpreted as a column vector v. Pre-multiply v with
+  /// this matrix, say M, and return Mv.
+  SmallVector<int64_t, 8>
+  postMultiplyWithColumn(ArrayRef<int64_t> colVec) const;
+
   /// Resize the matrix to the specified dimensions. If a dimension is smaller,
   /// 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);

mlir/lib/Analysis/Presburger/LinearTransform.cpp

for (unsigned row = 0; row < m.getNumRows(); ++row) {
     }
 
     ++echelonCol;
   }
 
   return {echelonCol, LinearTransform(std::move(resultMatrix))};
 }
 
-SmallVector<int64_t, 8>
-LinearTransform::preMultiplyWithRow(ArrayRef<int64_t> rowVec) const {
-  assert(rowVec.size() == matrix.getNumRows() &&
-         "row vector dimension should match transform output dimension");
-
-  SmallVector<int64_t, 8> result(matrix.getNumColumns(), 0);
-  for (unsigned col = 0, e = matrix.getNumColumns(); col < e; ++col)
-    for (unsigned i = 0, e = matrix.getNumRows(); i < e; ++i)
-      result[col] += rowVec[i] * matrix(i, col);
-  return result;
-}
-
-SmallVector<int64_t, 8>
-LinearTransform::postMultiplyWithColumn(ArrayRef<int64_t> colVec) const {
-  assert(matrix.getNumColumns() == colVec.size() &&
-         "column vector dimension should match transform input dimension");
-
-  SmallVector<int64_t, 8> result(matrix.getNumRows(), 0);
-  for (unsigned row = 0, e = matrix.getNumRows(); row < e; row++)
-    for (unsigned i = 0, e = matrix.getNumColumns(); i < e; i++)
-      result[row] += matrix(row, i) * colVec[i];
-  return result;
-}
-
 IntegerPolyhedron
 LinearTransform::applyTo(const IntegerPolyhedron &poly) const {
   IntegerPolyhedron result(poly.getNumIds());
 
   for (unsigned i = 0, e = poly.getNumEqualities(); i < e; ++i) {
     ArrayRef<int64_t> eq = poly.getEquality(i);
 
     int64_t c = eq.back();

mlir/lib/Analysis/Presburger/Matrix.cpp

for (unsigned row = 0, e = getNumRows(); row < e; ++row)
     at(row, targetColumn) += scale * at(row, sourceColumn);
 }
 
 void Matrix::negateColumn(unsigned column) {
   for (unsigned row = 0, e = getNumRows(); row < e; ++row)
     at(row, column) = -at(row, column);
 }
 
+SmallVector<int64_t, 8>
+Matrix::preMultiplyWithRow(ArrayRef<int64_t> rowVec) const {
+  assert(rowVec.size() == getNumRows() && "Invalid row vector dimension!");
+
+  SmallVector<int64_t, 8> result(getNumColumns(), 0);
+  for (unsigned col = 0, e = getNumColumns(); col < e; ++col)
+    for (unsigned i = 0, e = getNumRows(); i < e; ++i)
+      result[col] += rowVec[i] * at(i, col);
+  return result;
+}
+
+SmallVector<int64_t, 8>
+Matrix::postMultiplyWithColumn(ArrayRef<int64_t> colVec) const {
+  assert(getNumColumns() == colVec.size() &&
+         "Invalid column vector dimension!");
+
+  SmallVector<int64_t, 8> result(getNumRows(), 0);
+  for (unsigned row = 0, e = getNumRows(); row < e; row++)
+    for (unsigned i = 0, e = getNumColumns(); i < e; i++)
+      result[row] += at(row, i) * colVec[i];
+  return result;
+}
+
 void Matrix::print(raw_ostream &os) const {
   for (unsigned row = 0; row < nRows; ++row) {
     for (unsigned column = 0; column < nColumns; ++column)
       os << at(row, column) << ' ';
     os << '\n';
   }
 }