| 74 | |
| 75 | template<typename T> |
| 76 | Array<T> generalSolve(const Array<T> &a, const Array<T> &b) { |
| 77 | int batches = a.dims()[2] * a.dims()[3]; |
| 78 | |
| 79 | dim4 aDims = a.dims(); |
| 80 | dim4 bDims = b.dims(); |
| 81 | int M = aDims[0]; |
| 82 | int N = aDims[1]; |
| 83 | int K = bDims[1]; |
| 84 | int MN = std::min(M, N); |
| 85 | |
| 86 | int lda = a.strides()[1]; |
| 87 | int astride = a.strides()[2]; |
| 88 | auto ipiv = memAlloc<int64_t>(MN * batches); |
| 89 | int ipivstride = MN; |
| 90 | |
| 91 | int ldb = b.strides()[1]; |
| 92 | int bstride = b.strides()[2]; |
| 93 | |
| 94 | vector<int> info(batches, 0); |
| 95 | |
| 96 | Array<T> A = copyArray<T>(a); // A will be overwritten by L,U |
| 97 | Array<T> B = copyArray<T>(b); // will be overwritten with solution |
| 98 | |
| 99 | std::int64_t scratchpad_size = |
| 100 | ::oneapi::mkl::lapack::getrf_batch_scratchpad_size<compute_t<T>>( |
| 101 | getQueue(), M, N, lda, astride, ipivstride, batches); |
| 102 | |
| 103 | auto scratchpad = memAlloc<compute_t<T>>(scratchpad_size); |
| 104 | |
| 105 | buffer<compute_t<T>> aBuf = A.template getBufferWithOffset<compute_t<T>>(); |
| 106 | buffer<compute_t<T>> bBuf = B.template getBufferWithOffset<compute_t<T>>(); |
| 107 | ::oneapi::mkl::lapack::getrf_batch(getQueue(), M, N, aBuf, lda, astride, |
| 108 | *ipiv, ipivstride, batches, *scratchpad, |
| 109 | scratchpad->size()); |
| 110 | |
| 111 | scratchpad_size = |
| 112 | ::oneapi::mkl::lapack::getrs_batch_scratchpad_size<compute_t<T>>( |
| 113 | getQueue(), ::oneapi::mkl::transpose::nontrans, N, K, lda, astride, |
| 114 | ipivstride, ldb, bstride, batches); |
| 115 | |
| 116 | auto scratchpad_rs = memAlloc<compute_t<T>>(scratchpad_size); |
| 117 | |
| 118 | ::oneapi::mkl::lapack::getrs_batch( |
| 119 | getQueue(), ::oneapi::mkl::transpose::nontrans, N, K, aBuf, lda, |
| 120 | astride, *ipiv, ipivstride, bBuf, ldb, bstride, batches, *scratchpad_rs, |
| 121 | scratchpad_rs->size()); |
| 122 | |
| 123 | return B; |
| 124 | } |
| 125 | |
| 126 | template<typename T> |
| 127 | Array<T> leastSquares(const Array<T> &a, const Array<T> &b) { |