arrayfire · 9prady9 · Jul 31, 2021 · Jan 9, 2017 · Jul 14, 2020 · Jul 30, 2021
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -23,6 +23,7 @@ include(Version)
 include(platform)
 include(GetPrerequisites)
 include(CheckCXXCompilerFlag)
+include(CheckSymbolExists)
 include(SplitDebugInfo)
 
 # Use the function generate_product_version on Windows
@@ -168,7 +169,9 @@ mark_as_advanced(CLEAR CUDA_VERSION)
 # Note that the default value of AF_COMPUTE_LIBRARY is Intel-MKL.
 # Also, cmake doesn't have short-circuit of OR/AND conditions in if
 if(${AF_BUILD_CPU} OR ${AF_BUILD_OPENCL})
-  if("${AF_COMPUTE_LIBRARY}" STREQUAL "Intel-MKL")
+  if("${AF_COMPUTE_LIBRARY}" STREQUAL "Intel-MKL"
+      OR "${AF_COMPUTE_LIBRARY}" STREQUAL "MKL")
+    af_mkl_batch_check()
     dependency_check(MKL_FOUND "Please ensure Intel-MKL / oneAPI-oneMKL is installed")
     set(BUILD_WITH_MKL ON)
   elseif("${AF_COMPUTE_LIBRARY}" STREQUAL "FFTW/LAPACK/BLAS")

diff --git a/CMakeModules/InternalUtils.cmake b/CMakeModules/InternalUtils.cmake
@@ -218,6 +218,11 @@ macro(set_policies)
   endforeach()
 endmacro()
 
+macro(af_mkl_batch_check)
+  set(CMAKE_REQUIRED_LIBRARIES "MKL::RT")
+  check_symbol_exists(sgetrf_batch_strided "mkl_lapack.h" MKL_BATCH)
+endmacro()
+
 mark_as_advanced(
     pkgcfg_lib_PC_CBLAS_cblas
     pkgcfg_lib_PC_LAPACKE_lapacke

diff --git a/src/api/c/solve.cpp b/src/api/c/solve.cpp
@@ -34,10 +34,6 @@ af_err af_solve(af_array* out, const af_array a, const af_array b,
         const ArrayInfo& a_info = getInfo(a);
         const ArrayInfo& b_info = getInfo(b);
 
-        if (a_info.ndims() > 2 || b_info.ndims() > 2) {
-            AF_ERROR("solve can not be used in batch mode", AF_ERR_BATCH);
-        }
-
         af_dtype a_type = a_info.getType();
         af_dtype b_type = b_info.getType();
 

diff --git a/src/backend/cpu/Array.hpp b/src/backend/cpu/Array.hpp
@@ -22,6 +22,7 @@
 #include <af/dim4.hpp>
 #include <af/seq.h>
 
+#include <algorithm>
 #include <cstddef>
 #include <memory>
 #include <vector>
@@ -153,6 +154,9 @@ class Array {
     }
 
     void resetInfo(const af::dim4 &dims) { info.resetInfo(dims); }
+
+    // Modifies the dimensions of the array without modifing the underlying
+    // data
     void resetDims(const af::dim4 &dims) { info.resetDims(dims); }
     void modDims(const af::dim4 &newDims) { info.modDims(newDims); }
     void modStrides(const af::dim4 &newStrides) { info.modStrides(newStrides); }

diff --git a/src/backend/cpu/CMakeLists.txt b/src/backend/cpu/CMakeLists.txt
@@ -314,6 +314,11 @@ target_link_libraries(afcpu
   )
 if(BUILD_WITH_MKL)
   target_compile_definitions(afcpu PRIVATE USE_MKL)
+
+  if(MKL_BATCH)
+    target_compile_definitions(afcpu PRIVATE AF_USE_MKL_BATCH)
+  endif()
+
   if(AF_WITH_STATIC_MKL)
       target_link_libraries(afcpu PRIVATE MKL::Static)
   else()

diff --git a/src/backend/cpu/lapack_helper.hpp b/src/backend/cpu/lapack_helper.hpp
@@ -18,6 +18,7 @@
 #define LAPACK_NAME(fn) LAPACKE_##fn
 
 #ifdef USE_MKL
+#include <mkl_lapack.h>
 #include <mkl_lapacke.h>
 #else
 #ifdef __APPLE__

diff --git a/src/backend/cpu/solve.cpp b/src/backend/cpu/solve.cpp
@@ -17,6 +17,9 @@
 #include <math.hpp>
 #include <queue.hpp>
 #include <af/dim4.hpp>
+#include <algorithm>
+#include <complex>
+#include <vector>
 
 using af::dim4;
 
@@ -29,6 +32,21 @@ template<typename T>
 using gels_func_def = int (*)(ORDER_TYPE, char, int, int, int, T *, int, T *,
                               int);
 
+#ifdef AF_USE_MKL_BATCH
+template<typename T>
+using getrf_batch_strided_func_def =
+    void (*)(const MKL_INT *m, const MKL_INT *n, T *a, const MKL_INT *lda,
+             const MKL_INT *stride_a, MKL_INT *ipiv, const MKL_INT *stride_ipiv,
+             const MKL_INT *batch_size, MKL_INT *info);
+
+template<typename T>
+using getrs_batch_strided_func_def =
+    void (*)(const char *trans, const MKL_INT *n, const MKL_INT *nrhs, T *a,
+             const MKL_INT *lda, const MKL_INT *stride_a, MKL_INT *ipiv,
+             const MKL_INT *stride_ipiv, T *b, const MKL_INT *ldb,
+             const MKL_INT *stride_b, const MKL_INT *batch_size, MKL_INT *info);
+#endif
+
 template<typename T>
 using getrs_func_def = int (*)(ORDER_TYPE, char, int, int, const T *, int,
                                const int *, T *, int);
@@ -59,6 +77,70 @@ SOLVE_FUNC(gels, double, d)
 SOLVE_FUNC(gels, cfloat, c)
 SOLVE_FUNC(gels, cdouble, z)
 
+#ifdef AF_USE_MKL_BATCH
+
+template<typename T>
+struct mkl_type {
+    using type = T;
+};
+template<>
+struct mkl_type<std::complex<float>> {
+    using type = MKL_Complex8;
+};
+template<>
+struct mkl_type<std::complex<double>> {
+    using type = MKL_Complex16;
+};
+
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wnoexcept-type"
+template<typename T>
+getrf_batch_strided_func_def<T> getrf_batch_strided_func();
+
+template<>
+getrf_batch_strided_func_def<float> getrf_batch_strided_func<float>() {
+    return &sgetrf_batch_strided;
+}
+template<>
+getrf_batch_strided_func_def<double> getrf_batch_strided_func<double>() {
+    return &dgetrf_batch_strided;
+}
+template<>
+getrf_batch_strided_func_def<MKL_Complex8>
+getrf_batch_strided_func<MKL_Complex8>() {
+    return &cgetrf_batch_strided;
+}
+template<>
+getrf_batch_strided_func_def<MKL_Complex16>
+getrf_batch_strided_func<MKL_Complex16>() {
+    return &zgetrf_batch_strided;
+}
+
+template<typename T>
+getrs_batch_strided_func_def<T> getrs_batch_strided_func();
+
+template<>
+getrs_batch_strided_func_def<float> getrs_batch_strided_func<float>() {
+    return &sgetrs_batch_strided;
+}
+template<>
+getrs_batch_strided_func_def<double> getrs_batch_strided_func<double>() {
+    return &dgetrs_batch_strided;
+}
+template<>
+getrs_batch_strided_func_def<MKL_Complex8>
+getrs_batch_strided_func<MKL_Complex8>() {
+    return &cgetrs_batch_strided;
+}
+template<>
+getrs_batch_strided_func_def<MKL_Complex16>
+getrs_batch_strided_func<MKL_Complex16>() {
+    return &zgetrs_batch_strided;
+}
+
+#pragma GCC diagnostic pop
+#endif
+
 SOLVE_FUNC_DEF(getrs)
 SOLVE_FUNC(getrs, float, s)
 SOLVE_FUNC(getrs, double, d)
@@ -109,17 +191,81 @@ Array<T> triangleSolve(const Array<T> &A, const Array<T> &b,
     return B;
 }
 
+#ifdef AF_USE_MKL_BATCH
+
+template<typename T>
+Array<T> generalSolveBatched(const Array<T> &a, const Array<T> &b,
+                             const af_mat_prop options) {
+    using std::vector;
+    int batches = a.dims()[2] * a.dims()[3];
+
+    dim4 aDims = a.dims();
+    dim4 bDims = b.dims();
+    int M      = aDims[0];
+    int N      = aDims[1];
+    int K      = bDims[1];
+    int MN     = std::min(M, N);
+
+    int lda     = a.strides()[1];
+    int astride = a.strides()[2];
+
+    vector<int> ipiv(MN * batches);
+    int ipivstride = MN;
+
+    int ldb     = b.strides()[1];
+    int bstride = b.strides()[2];
+
+    vector<int> info(batches, 0);
+
+    char trans = 'N';
+
+    Array<T> A = copyArray<T>(a);
+    Array<T> B = copyArray<T>(b);
+
+    auto getrf_rs = [](char TRANS, int M, int N, int K, Param<T> a, int LDA,
+                       int ASTRIDE, vector<int> IPIV, int IPIVSTRIDE,
+                       Param<T> b, int LDB, int BSTRIDE, int BATCH_SIZE,
+                       vector<int> INFO) {
+        getrf_batch_strided_func<typename mkl_type<T>::type>()(
+            &M, &N, reinterpret_cast<typename mkl_type<T>::type *>(a.get()),
+            &LDA, &ASTRIDE, IPIV.data(), &IPIVSTRIDE, &BATCH_SIZE, INFO.data());
+
+        getrs_batch_strided_func<typename mkl_type<T>::type>()(
+            &TRANS, &M, &K,
+            reinterpret_cast<typename mkl_type<T>::type *>(a.get()), &LDA,
+            &ASTRIDE, IPIV.data(), &IPIVSTRIDE,
+            reinterpret_cast<typename mkl_type<T>::type *>(b.get()), &LDB,
+            &BSTRIDE, &BATCH_SIZE, INFO.data());
+    };
+
+    getQueue().enqueue(getrf_rs, trans, M, N, K, A, lda, astride, ipiv,
+                       ipivstride, B, ldb, bstride, batches, info);
+
+    return B;
+}
+#endif
+
 template<typename T>
 Array<T> solve(const Array<T> &a, const Array<T> &b,
                const af_mat_prop options) {
     if (options & AF_MAT_UPPER || options & AF_MAT_LOWER) {
         return triangleSolve<T>(a, b, options);
     }
 
+#ifdef AF_USE_MKL_BATCH
+    if (a.dims()[2] > 1 || a.dims()[3] > 1) {
+        return generalSolveBatched(a, b, options);
+    }
+#endif
+
     const dim4 NullShape(0, 0, 0, 0);
 
-    int M = a.dims()[0];
-    int N = a.dims()[1];
+    dim4 aDims = a.dims();
+    int batchz = aDims[2];
+    int batchw = aDims[3];
+
+    int M = aDims[0];
+    int N = aDims[1];
     int K = b.dims()[1];
 
     Array<T> A = copyArray<T>(a);
@@ -129,27 +275,37 @@ Array<T> solve(const Array<T> &a, const Array<T> &b,
                       ? copyArray(b)
                       : padArrayBorders(b, NullShape, endPadding, AF_PAD_ZERO));
 
-    if (M == N) {
-        Array<int> pivot = createEmptyArray<int>(dim4(N, 1, 1));
-
-        auto func = [=](Param<T> A, Param<T> B, Param<int> pivot, int N,
-                        int K) {
-            gesv_func<T>()(AF_LAPACK_COL_MAJOR, N, K, A.get(), A.strides(1),
-                           pivot.get(), B.get(), B.strides(1));
-        };
-        getQueue().enqueue(func, A, B, pivot, N, K);
-    } else {
-        auto func = [=](Param<T> A, Param<T> B, int M, int N, int K) {
-            int sM = A.strides(1);
-            int sN = A.strides(2) / sM;
-
-            gels_func<T>()(AF_LAPACK_COL_MAJOR, 'N', M, N, K, A.get(),
-                           A.strides(1), B.get(), max(sM, sN));
-        };
-        B.resetDims(dim4(N, K));
-        getQueue().enqueue(func, A, B, M, N, K);
+    for (int i = 0; i < batchw; i++) {
+        for (int j = 0; j < batchz; j++) {
+            Param<T> pA(A.get() + A.strides()[2] * j + A.strides()[3] * i,
+                        A.dims(), A.strides());
+            Param<T> pB(B.get() + B.strides()[2] * j + B.strides()[3] * i,
+                        B.dims(), B.strides());
+            if (M == N) {
+                Array<int> pivot = createEmptyArray<int>(dim4(N, 1, 1));
+
+                auto func = [](Param<T> A, Param<T> B, Param<int> pivot, int N,
+                               int K) {
+                    gesv_func<T>()(AF_LAPACK_COL_MAJOR, N, K, A.get(),
+                                   A.strides(1), pivot.get(), B.get(),
+                                   B.strides(1));
+                };
+                getQueue().enqueue(func, pA, pB, pivot, N, K);
+            } else {
+                auto func = [=](Param<T> A, Param<T> B, int M, int N, int K) {
+                    int sM = A.dims(0);
+                    int sN = A.dims(1);
+
+                    gels_func<T>()(AF_LAPACK_COL_MAJOR, 'N', M, N, K, A.get(),
+                                   A.strides(1), B.get(), max(sM, sN));
+                };
+                getQueue().enqueue(func, pA, pB, M, N, K);
+            }
+        }
     }
 
+    if (M != N) { B.resetDims(dim4(N, K, B.dims()[2], B.dims()[3])); }
+
     return B;
 }
 
@@ -162,13 +318,17 @@ namespace cpu {
 template<typename T>
 Array<T> solveLU(const Array<T> &A, const Array<int> &pivot, const Array<T> &b,
                  const af_mat_prop options) {
-    AF_ERROR("Linear Algebra is disabled on CPU", AF_ERR_NOT_CONFIGURED);
+    AF_ERROR(
+        "This version of ArrayFire was built without linear algebra routines",
+        AF_ERR_NOT_CONFIGURED);
 }
 
 template<typename T>
 Array<T> solve(const Array<T> &a, const Array<T> &b,
                const af_mat_prop options) {
-    AF_ERROR("Linear Algebra is disabled on CPU", AF_ERR_NOT_CONFIGURED);
+    AF_ERROR(
+        "This version of ArrayFire was built without linear algebra routines",
+        AF_ERR_NOT_CONFIGURED);
 }
 
 }  // namespace cpu

diff --git a/src/backend/cuda/memory.cpp b/src/backend/cuda/memory.cpp
@@ -24,8 +24,8 @@
 #include <types.hpp>
 #include <af/dim4.hpp>
 
+#include <cstdlib>
 #include <mutex>
-#include <utility>
 
 using af::dim4;
 using common::bytesToString;