| 208 | |
| 209 | template<typename T> |
| 210 | Array<T> generalSolveBatched(const Array<T> &a, const Array<T> &b, |
| 211 | const af_mat_prop options) { |
| 212 | using std::vector; |
| 213 | int batches = a.dims()[2] * a.dims()[3]; |
| 214 | |
| 215 | dim4 aDims = a.dims(); |
| 216 | dim4 bDims = b.dims(); |
| 217 | int M = aDims[0]; |
| 218 | int N = aDims[1]; |
| 219 | int K = bDims[1]; |
| 220 | int MN = std::min(M, N); |
| 221 | |
| 222 | int lda = a.strides()[1]; |
| 223 | int astride = a.strides()[2]; |
| 224 | |
| 225 | vector<int> ipiv(MN * batches); |
| 226 | int ipivstride = MN; |
| 227 | |
| 228 | int ldb = b.strides()[1]; |
| 229 | int bstride = b.strides()[2]; |
| 230 | |
| 231 | vector<int> info(batches, 0); |
| 232 | |
| 233 | char trans = 'N'; |
| 234 | |
| 235 | Array<T> A = copyArray<T>(a); |
| 236 | Array<T> B = copyArray<T>(b); |
| 237 | |
| 238 | auto getrf_rs = [](char TRANS, int M, int N, int K, Param<T> a, int LDA, |
| 239 | int ASTRIDE, vector<int> IPIV, int IPIVSTRIDE, |
| 240 | Param<T> b, int LDB, int BSTRIDE, int BATCH_SIZE, |
| 241 | vector<int> INFO) { |
| 242 | getrf_batch_strided_func<typename mkl_type<T>::type>()( |
| 243 | &M, &N, reinterpret_cast<typename mkl_type<T>::type *>(a.get()), |
| 244 | &LDA, &ASTRIDE, IPIV.data(), &IPIVSTRIDE, &BATCH_SIZE, INFO.data()); |
| 245 | |
| 246 | getrs_batch_strided_func<typename mkl_type<T>::type>()( |
| 247 | &TRANS, &M, &K, |
| 248 | reinterpret_cast<typename mkl_type<T>::type *>(a.get()), &LDA, |
| 249 | &ASTRIDE, IPIV.data(), &IPIVSTRIDE, |
| 250 | reinterpret_cast<typename mkl_type<T>::type *>(b.get()), &LDB, |
| 251 | &BSTRIDE, &BATCH_SIZE, INFO.data()); |
| 252 | }; |
| 253 | |
| 254 | getQueue().enqueue(getrf_rs, trans, M, N, K, A, lda, astride, ipiv, |
| 255 | ipivstride, B, ldb, bstride, batches, info); |
| 256 | |
| 257 | return B; |
| 258 | } |
| 259 | #endif |
| 260 | |
| 261 | template<typename T> |