| 95 | |
| 96 | template<typename T> |
| 97 | Array<T> qr_inplace(Array<T> &in) { |
| 98 | dim4 iDims = in.dims(); |
| 99 | dim4 iStrides = in.strides(); |
| 100 | int M = iDims[0]; |
| 101 | int N = iDims[1]; |
| 102 | |
| 103 | Array<T> t = createEmptyArray<T>(af::dim4(min(M, N), 1, 1, 1)); |
| 104 | |
| 105 | // Get workspace needed for QR |
| 106 | std::int64_t scratchpad_size = |
| 107 | ::oneapi::mkl::lapack::geqrf_scratchpad_size<compute_t<T>>( |
| 108 | getQueue(), iDims[0], iDims[1], iStrides[1]); |
| 109 | |
| 110 | auto scratchpad = memAlloc<compute_t<T>>(scratchpad_size); |
| 111 | |
| 112 | buffer<compute_t<T>> iBuf = in.template getBufferWithOffset<compute_t<T>>(); |
| 113 | buffer<compute_t<T>> tBuf = t.template getBufferWithOffset<compute_t<T>>(); |
| 114 | // In place Perform in place QR |
| 115 | ::oneapi::mkl::lapack::geqrf(getQueue(), iDims[0], iDims[1], iBuf, |
| 116 | iStrides[1], tBuf, *scratchpad, |
| 117 | scratchpad->size()); |
| 118 | return t; |
| 119 | } |
| 120 | |
| 121 | #define INSTANTIATE_QR(T) \ |
| 122 | template Array<T> qr_inplace<T>(Array<T> & in); \ |