| 28 | |
| 29 | template<typename T> |
| 30 | void qr(Array<T> &q, Array<T> &r, Array<T> &t, const Array<T> &in) { |
| 31 | dim4 iDims = in.dims(); |
| 32 | int M = iDims[0]; |
| 33 | int N = iDims[1]; |
| 34 | |
| 35 | Array<T> in_copy = copyArray<T>(in); |
| 36 | |
| 37 | // Get workspace needed for QR |
| 38 | std::int64_t scratchpad_size = |
| 39 | ::oneapi::mkl::lapack::geqrf_scratchpad_size<compute_t<T>>( |
| 40 | getQueue(), iDims[0], iDims[1], in_copy.strides()[1]); |
| 41 | |
| 42 | auto scratchpad = memAlloc<compute_t<T>>(scratchpad_size); |
| 43 | |
| 44 | t = createEmptyArray<T>(af::dim4(min(M, N), 1, 1, 1)); |
| 45 | |
| 46 | buffer<compute_t<T>> iBuf = |
| 47 | in_copy.template getBufferWithOffset<compute_t<T>>(); |
| 48 | buffer<compute_t<T>> tBuf = t.template getBufferWithOffset<compute_t<T>>(); |
| 49 | ::oneapi::mkl::lapack::geqrf(getQueue(), M, N, iBuf, in_copy.strides()[1], |
| 50 | tBuf, *scratchpad, scratchpad->size()); |
| 51 | // SPLIT into q and r |
| 52 | dim4 rdims(M, N); |
| 53 | r = createEmptyArray<T>(rdims); |
| 54 | |
| 55 | constexpr bool is_upper = true; |
| 56 | constexpr bool is_unit_diag = false; |
| 57 | kernel::triangle<T>(r, in_copy, is_upper, is_unit_diag); |
| 58 | |
| 59 | int mn = max(M, N); |
| 60 | dim4 qdims(M, mn); |
| 61 | q = identity<T>(qdims); |
| 62 | |
| 63 | buffer<compute_t<T>> qBuf = q.template getBufferWithOffset<compute_t<T>>(); |
| 64 | if constexpr (std::is_floating_point<compute_t<T>>()) { |
| 65 | std::int64_t scratchpad_size = |
| 66 | ::oneapi::mkl::lapack::ormqr_scratchpad_size<compute_t<T>>( |
| 67 | getQueue(), ::oneapi::mkl::side::left, |
| 68 | ::oneapi::mkl::transpose::nontrans, q.dims()[0], q.dims()[1], |
| 69 | min(M, N), in_copy.strides()[1], q.strides()[1]); |
| 70 | |
| 71 | auto scratchpad_ormqr = memAlloc<compute_t<T>>(scratchpad_size); |
| 72 | ::oneapi::mkl::lapack::ormqr( |
| 73 | getQueue(), ::oneapi::mkl::side::left, |
| 74 | ::oneapi::mkl::transpose::nontrans, q.dims()[0], q.dims()[1], |
| 75 | min(M, N), iBuf, in_copy.strides()[1], tBuf, qBuf, q.strides()[1], |
| 76 | *scratchpad_ormqr, scratchpad_ormqr->size()); |
| 77 | |
| 78 | } else if constexpr (common::isComplex(static_cast<af::dtype>( |
| 79 | dtype_traits<compute_t<T>>::af_type))) { |
| 80 | std::int64_t scratchpad_size = |
| 81 | ::oneapi::mkl::lapack::unmqr_scratchpad_size<compute_t<T>>( |
| 82 | getQueue(), ::oneapi::mkl::side::left, |
| 83 | ::oneapi::mkl::transpose::nontrans, q.dims()[0], q.dims()[1], |
| 84 | min(M, N), in_copy.strides()[1], q.strides()[1]); |
| 85 | |
| 86 | auto scratchpad_ormqr = memAlloc<compute_t<T>>(scratchpad_size); |
| 87 | ::oneapi::mkl::lapack::unmqr( |