| 133 | |
| 134 | template<typename T> |
| 135 | void qr(Array<T> &q, Array<T> &r, Array<T> &t, const Array<T> &in) { |
| 136 | dim4 iDims = in.dims(); |
| 137 | int M = iDims[0]; |
| 138 | int N = iDims[1]; |
| 139 | |
| 140 | Array<T> in_copy = copyArray<T>(in); |
| 141 | |
| 142 | int lwork = 0; |
| 143 | |
| 144 | CUSOLVER_CHECK(geqrf_buf_func<T>()(solverDnHandle(), M, N, in_copy.get(), |
| 145 | in_copy.strides()[1], &lwork)); |
| 146 | |
| 147 | auto workspace = memAlloc<T>(lwork); |
| 148 | |
| 149 | t = createEmptyArray<T>(af::dim4(min(M, N), 1, 1, 1)); |
| 150 | auto info = memAlloc<int>(1); |
| 151 | |
| 152 | CUSOLVER_CHECK(geqrf_func<T>()(solverDnHandle(), M, N, in_copy.get(), |
| 153 | in_copy.strides()[1], t.get(), |
| 154 | workspace.get(), lwork, info.get())); |
| 155 | |
| 156 | // SPLIT into q and r |
| 157 | dim4 rdims(M, N); |
| 158 | r = createEmptyArray<T>(rdims); |
| 159 | |
| 160 | kernel::triangle<T>(r, in_copy, true, false); |
| 161 | |
| 162 | int mn = max(M, N); |
| 163 | dim4 qdims(M, mn); |
| 164 | q = identity<T>(qdims); |
| 165 | |
| 166 | CUSOLVER_CHECK(mqr_buf_func<T>()( |
| 167 | solverDnHandle(), CUBLAS_SIDE_LEFT, CUBLAS_OP_N, q.dims()[0], |
| 168 | q.dims()[1], min(M, N), in_copy.get(), in_copy.strides()[1], t.get(), |
| 169 | q.get(), q.strides()[1], &lwork)); |
| 170 | |
| 171 | workspace = memAlloc<T>(lwork); |
| 172 | |
| 173 | CUSOLVER_CHECK(mqr_func<T>()( |
| 174 | solverDnHandle(), CUBLAS_SIDE_LEFT, CUBLAS_OP_N, q.dims()[0], |
| 175 | q.dims()[1], min(M, N), in_copy.get(), in_copy.strides()[1], t.get(), |
| 176 | q.get(), q.strides()[1], workspace.get(), lwork, info.get())); |
| 177 | |
| 178 | q.resetDims(dim4(M, M)); |
| 179 | } |
| 180 | |
| 181 | template<typename T> |
| 182 | Array<T> qr_inplace(Array<T> &in) { |