| 125 | |
| 126 | template<typename T> |
| 127 | Array<T> leastSquares(const Array<T> &a, const Array<T> &b) { |
| 128 | int64_t M = a.dims()[0]; |
| 129 | int64_t N = a.dims()[1]; |
| 130 | int64_t K = b.dims()[1]; |
| 131 | int64_t MN = min(M, N); |
| 132 | |
| 133 | Array<T> B = createEmptyArray<T>(dim4()); |
| 134 | |
| 135 | if (M < N) { |
| 136 | const dim4 NullShape(0, 0, 0, 0); |
| 137 | |
| 138 | // Least squres for this case is solved using the following |
| 139 | // solve(A, B) == matmul(Q, Xpad); |
| 140 | // Where: |
| 141 | // Xpad == pad(Xt, N - M, 1); |
| 142 | // Xt == tri_solve(R1, B); |
| 143 | // R1 == R(seq(M), seq(M)); |
| 144 | // transpose(A) == matmul(Q, R); |
| 145 | |
| 146 | // QR is performed on the transpose of A |
| 147 | Array<T> A = transpose<T>(a, true); |
| 148 | dim4 endPadding(N - b.dims()[0], K - b.dims()[1], 0, 0); |
| 149 | B = (endPadding == NullShape |
| 150 | ? copyArray(b) |
| 151 | : padArrayBorders(b, NullShape, endPadding, AF_PAD_ZERO)); |
| 152 | |
| 153 | // Get workspace needed for QR |
| 154 | std::int64_t scratchpad_size = |
| 155 | ::oneapi::mkl::lapack::geqrf_scratchpad_size<compute_t<T>>( |
| 156 | getQueue(), A.dims()[0], A.dims()[1], A.strides()[1]); |
| 157 | |
| 158 | auto scratchpad = memAlloc<compute_t<T>>(scratchpad_size); |
| 159 | auto t = memAlloc<compute_t<T>>(MN); |
| 160 | |
| 161 | buffer<compute_t<T>> aBuf = |
| 162 | A.template getBufferWithOffset<compute_t<T>>(); |
| 163 | // In place Perform in place QR |
| 164 | ::oneapi::mkl::lapack::geqrf(getQueue(), A.dims()[0], A.dims()[1], aBuf, |
| 165 | A.strides()[1], *t, *scratchpad, |
| 166 | scratchpad->size()); |
| 167 | |
| 168 | // R1 = R(seq(M), seq(M)); |
| 169 | A.resetDims(dim4(M, M)); |
| 170 | |
| 171 | // Bt = tri_solve(R1, B); |
| 172 | B.resetDims(dim4(M, K)); |
| 173 | |
| 174 | buffer<compute_t<T>> bBuf = |
| 175 | B.template getBufferWithOffset<compute_t<T>>(); |
| 176 | // TODO: move to helper? trsm<T>(A, B, AF_MAT_CTRANS, true, true, |
| 177 | // false); |
| 178 | compute_t<T> alpha = scalar<compute_t<T>>(1); |
| 179 | ::oneapi::mkl::blas::trsm( |
| 180 | getQueue(), ::oneapi::mkl::side::left, ::oneapi::mkl::uplo::upper, |
| 181 | ::oneapi::mkl::transpose::conjtrans, ::oneapi::mkl::diag::nonunit, |
| 182 | B.dims()[0], B.dims()[1], alpha, aBuf, A.strides()[1], bBuf, |
| 183 | B.strides()[1]); |
| 184 | |