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