| 14 | |
| 15 | template <typename ctype> |
| 16 | void do_exec(ctype* dst, const ctype* src, size_t batch, size_t n, void* workspace) { |
| 17 | auto row_ptr = static_cast<ctype**>(workspace); |
| 18 | auto exmat = reinterpret_cast<ctype*>(row_ptr + n); |
| 19 | for (size_t b = 0; b < batch; ++b, src += n * n, dst += n * n) { |
| 20 | // exmat is [A | I] and row_ptr points to its rows |
| 21 | for (size_t i = 0; i < n; ++i) { |
| 22 | row_ptr[i] = exmat + i * n * 2; |
| 23 | memcpy(row_ptr[i], src + i * n, sizeof(ctype) * n); |
| 24 | memset(row_ptr[i] + n, 0, sizeof(ctype) * n); |
| 25 | row_ptr[i][n + i] = 1; |
| 26 | } |
| 27 | for (size_t i = 0; i < n; ++i) { |
| 28 | size_t pivot_row = 0; |
| 29 | // select pivot row that has max abs value |
| 30 | ctype pivot_row_val = static_cast<ctype>(0); |
| 31 | for (size_t j = i; j < n; ++j) { |
| 32 | ctype val = static_cast<ctype>(std::abs(row_ptr[j][i])); |
| 33 | if (val > pivot_row_val) { |
| 34 | pivot_row_val = val; |
| 35 | pivot_row = j; |
| 36 | } |
| 37 | } |
| 38 | megdnn_throw_if( |
| 39 | pivot_row_val < ctype(1e-7), megdnn_error, "pivot value too small"); |
| 40 | std::swap(row_ptr[i], row_ptr[pivot_row]); |
| 41 | |
| 42 | // substract pivot row from other rows |
| 43 | auto pivot_row_ptr = row_ptr[i]; |
| 44 | for (size_t j = 0; j < n; ++j) { |
| 45 | if (j == i) { |
| 46 | continue; |
| 47 | } |
| 48 | ctype inv_pivot = -row_ptr[j][i] / pivot_row_ptr[i]; |
| 49 | for (size_t k = i; k < n * 2; ++k) { |
| 50 | row_ptr[j][k] += pivot_row_ptr[k] * inv_pivot; |
| 51 | } |
| 52 | } |
| 53 | |
| 54 | // scale pivot row after subtracting it from other rows |
| 55 | { |
| 56 | ctype scale = (static_cast<ctype>(1)) / pivot_row_ptr[i]; |
| 57 | for (size_t j = i; j < n * 2; ++j) { |
| 58 | pivot_row_ptr[j] *= scale; |
| 59 | } |
| 60 | } |
| 61 | } |
| 62 | |
| 63 | for (size_t i = 0; i < n; ++i) { |
| 64 | memcpy(dst + i * n, row_ptr[i] + n, sizeof(ctype) * n); |
| 65 | } |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | void MatrixInverseImpl::exec( |
| 70 | _megdnn_tensor_in src, _megdnn_tensor_out dst, _megdnn_workspace workspace) { |