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