1 //===- Matrix.cpp - MLIR Matrix Class -------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Analysis/Presburger/Matrix.h"
10 #include "mlir/Analysis/Presburger/Utils.h"
11 #include "llvm/Support/MathExtras.h"
12 
13 using namespace mlir;
14 using namespace presburger;
15 
16 Matrix::Matrix(unsigned rows, unsigned columns, unsigned reservedRows,
17                unsigned reservedColumns)
18     : nRows(rows), nColumns(columns),
19       nReservedColumns(std::max(nColumns, reservedColumns)),
20       data(nRows * nReservedColumns) {
21   data.reserve(std::max(nRows, reservedRows) * nReservedColumns);
22 }
23 
24 Matrix Matrix::identity(unsigned dimension) {
25   Matrix matrix(dimension, dimension);
26   for (unsigned i = 0; i < dimension; ++i)
27     matrix(i, i) = 1;
28   return matrix;
29 }
30 
31 int64_t &Matrix::at(unsigned row, unsigned column) {
32   assert(row < nRows && "Row outside of range");
33   assert(column < nColumns && "Column outside of range");
34   return data[row * nReservedColumns + column];
35 }
36 
37 int64_t Matrix::at(unsigned row, unsigned column) const {
38   assert(row < nRows && "Row outside of range");
39   assert(column < nColumns && "Column outside of range");
40   return data[row * nReservedColumns + column];
41 }
42 
43 int64_t &Matrix::operator()(unsigned row, unsigned column) {
44   return at(row, column);
45 }
46 
47 int64_t Matrix::operator()(unsigned row, unsigned column) const {
48   return at(row, column);
49 }
50 
51 unsigned Matrix::getNumRows() const { return nRows; }
52 
53 unsigned Matrix::getNumColumns() const { return nColumns; }
54 
55 unsigned Matrix::getNumReservedColumns() const { return nReservedColumns; }
56 
57 unsigned Matrix::getNumReservedRows() const {
58   return data.capacity() / nReservedColumns;
59 }
60 
61 void Matrix::reserveRows(unsigned rows) {
62   data.reserve(rows * nReservedColumns);
63 }
64 
65 unsigned Matrix::appendExtraRow() {
66   resizeVertically(nRows + 1);
67   return nRows - 1;
68 }
69 
70 unsigned Matrix::appendExtraRow(ArrayRef<int64_t> elems) {
71   assert(elems.size() == nColumns && "elems must match row length!");
72   unsigned row = appendExtraRow();
73   for (unsigned col = 0; col < nColumns; ++col)
74     at(row, col) = elems[col];
75   return row;
76 }
77 
78 void Matrix::resizeHorizontally(unsigned newNColumns) {
79   if (newNColumns < nColumns)
80     removeColumns(newNColumns, nColumns - newNColumns);
81   if (newNColumns > nColumns)
82     insertColumns(nColumns, newNColumns - nColumns);
83 }
84 
85 void Matrix::resize(unsigned newNRows, unsigned newNColumns) {
86   resizeHorizontally(newNColumns);
87   resizeVertically(newNRows);
88 }
89 
90 void Matrix::resizeVertically(unsigned newNRows) {
91   nRows = newNRows;
92   data.resize(nRows * nReservedColumns);
93 }
94 
95 void Matrix::swapRows(unsigned row, unsigned otherRow) {
96   assert((row < getNumRows() && otherRow < getNumRows()) &&
97          "Given row out of bounds");
98   if (row == otherRow)
99     return;
100   for (unsigned col = 0; col < nColumns; col++)
101     std::swap(at(row, col), at(otherRow, col));
102 }
103 
104 void Matrix::swapColumns(unsigned column, unsigned otherColumn) {
105   assert((column < getNumColumns() && otherColumn < getNumColumns()) &&
106          "Given column out of bounds");
107   if (column == otherColumn)
108     return;
109   for (unsigned row = 0; row < nRows; row++)
110     std::swap(at(row, column), at(row, otherColumn));
111 }
112 
113 MutableArrayRef<int64_t> Matrix::getRow(unsigned row) {
114   return {&data[row * nReservedColumns], nColumns};
115 }
116 
117 ArrayRef<int64_t> Matrix::getRow(unsigned row) const {
118   return {&data[row * nReservedColumns], nColumns};
119 }
120 
121 void Matrix::insertColumn(unsigned pos) { insertColumns(pos, 1); }
122 void Matrix::insertColumns(unsigned pos, unsigned count) {
123   if (count == 0)
124     return;
125   assert(pos <= nColumns);
126   unsigned oldNReservedColumns = nReservedColumns;
127   if (nColumns + count > nReservedColumns) {
128     nReservedColumns = llvm::NextPowerOf2(nColumns + count);
129     data.resize(nRows * nReservedColumns);
130   }
131   nColumns += count;
132 
133   for (int ri = nRows - 1; ri >= 0; --ri) {
134     for (int ci = nReservedColumns - 1; ci >= 0; --ci) {
135       unsigned r = ri;
136       unsigned c = ci;
137       int64_t &dest = data[r * nReservedColumns + c];
138       if (c >= nColumns) { // NOLINT
139         // Out of bounds columns are zero-initialized. NOLINT because clang-tidy
140         // complains about this branch being the same as the c >= pos one.
141         //
142         // TODO: this case can be skipped if the number of reserved columns
143         // didn't change.
144         dest = 0;
145       } else if (c >= pos + count) {
146         // Shift the data occuring after the inserted columns.
147         dest = data[r * oldNReservedColumns + c - count];
148       } else if (c >= pos) {
149         // The inserted columns are also zero-initialized.
150         dest = 0;
151       } else {
152         // The columns before the inserted columns stay at the same (row, col)
153         // but this corresponds to a different location in the linearized array
154         // if the number of reserved columns changed.
155         if (nReservedColumns == oldNReservedColumns)
156           break;
157         dest = data[r * oldNReservedColumns + c];
158       }
159     }
160   }
161 }
162 
163 void Matrix::removeColumn(unsigned pos) { removeColumns(pos, 1); }
164 void Matrix::removeColumns(unsigned pos, unsigned count) {
165   if (count == 0)
166     return;
167   assert(pos + count - 1 < nColumns);
168   for (unsigned r = 0; r < nRows; ++r) {
169     for (unsigned c = pos; c < nColumns - count; ++c)
170       at(r, c) = at(r, c + count);
171     for (unsigned c = nColumns - count; c < nColumns; ++c)
172       at(r, c) = 0;
173   }
174   nColumns -= count;
175 }
176 
177 void Matrix::insertRow(unsigned pos) { insertRows(pos, 1); }
178 void Matrix::insertRows(unsigned pos, unsigned count) {
179   if (count == 0)
180     return;
181 
182   assert(pos <= nRows);
183   resizeVertically(nRows + count);
184   for (int r = nRows - 1; r >= int(pos + count); --r)
185     copyRow(r - count, r);
186   for (int r = pos + count - 1; r >= int(pos); --r)
187     for (unsigned c = 0; c < nColumns; ++c)
188       at(r, c) = 0;
189 }
190 
191 void Matrix::removeRow(unsigned pos) { removeRows(pos, 1); }
192 void Matrix::removeRows(unsigned pos, unsigned count) {
193   if (count == 0)
194     return;
195   assert(pos + count - 1 <= nRows);
196   for (unsigned r = pos; r + count < nRows; ++r)
197     copyRow(r + count, r);
198   resizeVertically(nRows - count);
199 }
200 
201 void Matrix::copyRow(unsigned sourceRow, unsigned targetRow) {
202   if (sourceRow == targetRow)
203     return;
204   for (unsigned c = 0; c < nColumns; ++c)
205     at(targetRow, c) = at(sourceRow, c);
206 }
207 
208 void Matrix::fillRow(unsigned row, int64_t value) {
209   for (unsigned col = 0; col < nColumns; ++col)
210     at(row, col) = value;
211 }
212 
213 void Matrix::addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale) {
214   if (scale == 0)
215     return;
216   for (unsigned col = 0; col < nColumns; ++col)
217     at(targetRow, col) += scale * at(sourceRow, col);
218 }
219 
220 void Matrix::addToColumn(unsigned sourceColumn, unsigned targetColumn,
221                          int64_t scale) {
222   if (scale == 0)
223     return;
224   for (unsigned row = 0, e = getNumRows(); row < e; ++row)
225     at(row, targetColumn) += scale * at(row, sourceColumn);
226 }
227 
228 void Matrix::negateColumn(unsigned column) {
229   for (unsigned row = 0, e = getNumRows(); row < e; ++row)
230     at(row, column) = -at(row, column);
231 }
232 
233 void Matrix::negateRow(unsigned row) {
234   for (unsigned column = 0, e = getNumColumns(); column < e; ++column)
235     at(row, column) = -at(row, column);
236 }
237 
238 int64_t Matrix::normalizeRow(unsigned row, unsigned cols) {
239   return normalizeRange(getRow(row).slice(0, cols));
240 }
241 
242 int64_t Matrix::normalizeRow(unsigned row) {
243   return normalizeRow(row, getNumColumns());
244 }
245 
246 SmallVector<int64_t, 8>
247 Matrix::preMultiplyWithRow(ArrayRef<int64_t> rowVec) const {
248   assert(rowVec.size() == getNumRows() && "Invalid row vector dimension!");
249 
250   SmallVector<int64_t, 8> result(getNumColumns(), 0);
251   for (unsigned col = 0, e = getNumColumns(); col < e; ++col)
252     for (unsigned i = 0, e = getNumRows(); i < e; ++i)
253       result[col] += rowVec[i] * at(i, col);
254   return result;
255 }
256 
257 SmallVector<int64_t, 8>
258 Matrix::postMultiplyWithColumn(ArrayRef<int64_t> colVec) const {
259   assert(getNumColumns() == colVec.size() &&
260          "Invalid column vector dimension!");
261 
262   SmallVector<int64_t, 8> result(getNumRows(), 0);
263   for (unsigned row = 0, e = getNumRows(); row < e; row++)
264     for (unsigned i = 0, e = getNumColumns(); i < e; i++)
265       result[row] += at(row, i) * colVec[i];
266   return result;
267 }
268 
269 void Matrix::print(raw_ostream &os) const {
270   for (unsigned row = 0; row < nRows; ++row) {
271     for (unsigned column = 0; column < nColumns; ++column)
272       os << at(row, column) << ' ';
273     os << '\n';
274   }
275 }
276 
277 void Matrix::dump() const { print(llvm::errs()); }
278 
279 bool Matrix::hasConsistentState() const {
280   if (data.size() != nRows * nReservedColumns)
281     return false;
282   if (nColumns > nReservedColumns)
283     return false;
284   for (unsigned r = 0; r < nRows; ++r)
285     for (unsigned c = nColumns; c < nReservedColumns; ++c)
286       if (data[r * nReservedColumns + c] != 0)
287         return false;
288   return true;
289 }
290