| 200 | |
| 201 | template<typename T> |
| 202 | Array<T> generalSolveBatched(const Array<T> &a, const Array<T> &b, |
| 203 | const af_mat_prop options) { |
| 204 | using std::vector; |
| 205 | int batches = a.dims()[2] * a.dims()[3]; |
| 206 | |
| 207 | dim4 aDims = a.dims(); |
| 208 | dim4 bDims = b.dims(); |
| 209 | int M = aDims[0]; |
| 210 | int N = aDims[1]; |
| 211 | int K = bDims[1]; |
| 212 | int MN = std::min(M, N); |
| 213 | |
| 214 | int lda = a.strides()[1]; |
| 215 | int astride = a.strides()[2]; |
| 216 | |
| 217 | vector<int> ipiv(MN * batches); |
| 218 | int ipivstride = MN; |
| 219 | |
| 220 | int ldb = b.strides()[1]; |
| 221 | int bstride = b.strides()[2]; |
| 222 | |
| 223 | vector<int> info(batches, 0); |
| 224 | |
| 225 | char trans = 'N'; |
| 226 | |
| 227 | Array<T> A = copyArray<T>(a); |
| 228 | Array<T> B = copyArray<T>(b); |
| 229 | |
| 230 | mapped_ptr<T> aPtr = A.getMappedPtr(); |
| 231 | mapped_ptr<T> bPtr = B.getMappedPtr(); |
| 232 | |
| 233 | getrf_batch_strided_func<typename mkl_type<T>::type>()( |
| 234 | &M, &N, reinterpret_cast<typename mkl_type<T>::type *>(aPtr.get()), |
| 235 | &lda, &astride, ipiv.data(), &ipivstride, &batches, info.data()); |
| 236 | |
| 237 | getrs_batch_strided_func<typename mkl_type<T>::type>()( |
| 238 | &trans, &M, &K, |
| 239 | reinterpret_cast<typename mkl_type<T>::type *>(aPtr.get()), &lda, |
| 240 | &astride, ipiv.data(), &ipivstride, |
| 241 | reinterpret_cast<typename mkl_type<T>::type *>(bPtr.get()), &ldb, |
| 242 | &bstride, &batches, info.data()); |
| 243 | |
| 244 | return B; |
| 245 | } |
| 246 | #endif |
| 247 | |
| 248 | template<typename T> |