| 60 | |
| 61 | template<typename T> |
| 62 | Array<T> generalSolve(const Array<T> &a, const Array<T> &b) { |
| 63 | dim4 aDims = a.dims(); |
| 64 | int batchz = aDims[2]; |
| 65 | int batchw = aDims[3]; |
| 66 | |
| 67 | Array<T> A = copyArray<T>(a); |
| 68 | Array<T> B = copyArray<T>(b); |
| 69 | |
| 70 | for (int i = 0; i < batchw; i++) { |
| 71 | for (int j = 0; j < batchz; j++) { |
| 72 | int M = aDims[0]; |
| 73 | int N = aDims[1]; |
| 74 | int MN = min(M, N); |
| 75 | vector<int> ipiv(MN); |
| 76 | |
| 77 | Buffer *A_buf = A.get(); |
| 78 | int info = 0; |
| 79 | cl_command_queue q = getQueue()(); |
| 80 | auto aoffset = |
| 81 | A.getOffset() + j * A.strides()[2] + i * A.strides()[3]; |
| 82 | magma_getrf_gpu<T>(M, N, (*A_buf)(), aoffset, A.strides()[1], |
| 83 | &ipiv[0], q, &info); |
| 84 | |
| 85 | Buffer *B_buf = B.get(); |
| 86 | int K = B.dims()[1]; |
| 87 | |
| 88 | auto boffset = |
| 89 | B.getOffset() + j * B.strides()[2] + i * B.strides()[3]; |
| 90 | magma_getrs_gpu<T>(MagmaNoTrans, M, K, (*A_buf)(), aoffset, |
| 91 | A.strides()[1], &ipiv[0], (*B_buf)(), boffset, |
| 92 | B.strides()[1], q, &info); |
| 93 | } |
| 94 | } |
| 95 | return B; |
| 96 | } |
| 97 | |
| 98 | template<typename T> |
| 99 | Array<T> leastSquares(const Array<T> &a, const Array<T> &b) { |