From 0d6a3b4de432b10367bfa9bbf7641c39c1e9d30b Mon Sep 17 00:00:00 2001 From: Pavan Yalamanchili Date: Sun, 8 Jan 2017 22:18:58 -0800 Subject: [PATCH 1/7] Adding support for batched solve in CUDA backend --- src/api/c/solve.cpp | 4 -- src/backend/cuda/solve.cu | 136 +++++++++++++++++++++++++++++++++++++- 2 files changed, 135 insertions(+), 5 deletions(-) diff --git a/src/api/c/solve.cpp b/src/api/c/solve.cpp index 6328e90f01..ec17aafaba 100644 --- 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/cuda/solve.cu b/src/backend/cuda/solve.cu index 92cdb64b2e..988061ba12 100644 --- a/src/backend/cuda/solve.cu +++ b/src/backend/cuda/solve.cu @@ -12,8 +12,9 @@ #include #include #include -#include +#include #include +#include #include #include #include @@ -24,6 +25,64 @@ namespace cuda { +// cublasStatus_t cublas<>getrsBatched( cublasHandle_t handle, +// cublasOperation_t trans, +// int n, +// int nrhs, +// const <> *Aarray[], +// int lda, +// const int *devIpiv, +// <> *Barray[], +// int ldb, +// int *info, +// int batchSize); + +template +struct getrsBatched_func_def_t { + typedef cublasStatus_t (*getrsBatched_func_def)(cublasHandle_t, + cublasOperation_t, int, int, + const T **, int, + const int *, T **, int, + int *, int); +}; + +// cublasStatus_t cublas<>getrfBatched(cublasHandle_t handle, +// int n, +// float *A[], +// int lda, +// int *P, +// int *info, +// int batchSize); + +template +struct getrfBatched_func_def_t { + typedef cublasStatus_t (*getrfBatched_func_def)(cublasHandle_t, int, T **, + int, int *, int *, int); +}; + +#define SOLVE_BATCH_FUNC_DEF(FUNC) \ + template \ + typename FUNC##_func_def_t::FUNC##_func_def FUNC##_func(); + +#define SOLVE_BATCH_FUNC(FUNC, TYPE, PREFIX) \ + template<> \ + typename FUNC##_func_def_t::FUNC##_func_def FUNC##_func() { \ + return (FUNC##_func_def_t::FUNC##_func_def) & \ + cublas##PREFIX##FUNC; \ + } + +SOLVE_BATCH_FUNC_DEF(getrfBatched) +SOLVE_BATCH_FUNC(getrfBatched, float, S) +SOLVE_BATCH_FUNC(getrfBatched, double, D) +SOLVE_BATCH_FUNC(getrfBatched, cfloat, C) +SOLVE_BATCH_FUNC(getrfBatched, cdouble, Z) + +SOLVE_BATCH_FUNC_DEF(getrsBatched) +SOLVE_BATCH_FUNC(getrsBatched, float, S) +SOLVE_BATCH_FUNC(getrsBatched, double, D) +SOLVE_BATCH_FUNC(getrsBatched, cfloat, C) +SOLVE_BATCH_FUNC(getrsBatched, cdouble, Z) + // cusolverStatus_t cusolverDn<>getrs( // cusolverDnHandle_t handle, // cublasOperation_t trans, @@ -172,8 +231,83 @@ Array solveLU(const Array &A, const Array &pivot, const Array &b, return B; } +template +Array generalSolveBatched(const Array &a, const Array &b) { + Array A = copyArray(a); + Array B = copyArray(b); + + dim4 aDims = a.dims(); + int M = aDims[0]; + int N = aDims[1]; + int NRHS = b.dims()[1]; + + if (M != N) { + AF_ERROR("Batched solve requires square matrices", AF_ERR_ARG); + } + + int batchz = aDims[2]; + int batchw = aDims[3]; + int batch = batchz * batchw; + + size_t bytes = batch * sizeof(T *); + using unique_mem_ptr = std::unique_ptr; + + unique_mem_ptr aBatched_host_mem(pinnedAlloc(bytes), + pinnedFree); + unique_mem_ptr bBatched_host_mem(pinnedAlloc(bytes), + pinnedFree); + + T *a_ptr = A.get(); + T *b_ptr = B.get(); + T **aBatched_host_ptrs = (T **)aBatched_host_mem.get(); + T **bBatched_host_ptrs = (T **)bBatched_host_mem.get(); + + for (int i = 0; i < batchw; i++) { + for (int j = 0; j < batchz; j++) { + aBatched_host_ptrs[i * batchz + j] = + a_ptr + j * A.strides()[2] + i * A.strides()[3]; + bBatched_host_ptrs[i * batchz + j] = + b_ptr + j * B.strides()[2] + i * B.strides()[3]; + } + } + + auto aBatched_device_mem = memAlloc(bytes); + auto bBatched_device_mem = memAlloc(bytes); + + T **aBatched_device_ptrs = (T **)aBatched_device_mem.get(); + T **bBatched_device_ptrs = (T **)bBatched_device_mem.get(); + + CUDA_CHECK(cudaMemcpyAsync(aBatched_device_ptrs, aBatched_host_ptrs, bytes, + cudaMemcpyHostToDevice, + getStream(getActiveDeviceId()))); + + // Perform batched LU + // getrf requires pivot and info to be device pointers + Array pivots = createEmptyArray(af::dim4(N, batch, 1, 1)); + Array info = createEmptyArray(af::dim4(batch, 1, 1, 1)); + + CUBLAS_CHECK(getrfBatched_func()(blasHandle(), N, aBatched_device_ptrs, + A.strides()[1], pivots.get(), + info.get(), batch)); + + CUDA_CHECK(cudaMemcpyAsync(bBatched_device_ptrs, bBatched_host_ptrs, bytes, + cudaMemcpyHostToDevice, + getStream(getActiveDeviceId()))); + + // getrs requires info to be host pointer + unique_mem_ptr info_host_mem(pinnedAlloc(batch * sizeof(int)), + pinnedFree); + CUBLAS_CHECK(getrsBatched_func()( + blasHandle(), CUBLAS_OP_N, N, NRHS, (const T **)aBatched_device_ptrs, + A.strides()[1], pivots.get(), bBatched_device_ptrs, B.strides()[1], + (int *)info_host_mem.get(), batch)); + return B; +} + template Array generalSolve(const Array &a, const Array &b) { + if (a.dims()[2] > 1 || a.dims()[3] > 1) return generalSolveBatched(a, b); + int M = a.dims()[0]; int N = a.dims()[1]; int K = b.dims()[1]; From 1b0074fc4c756c808b0d9cb1cd9e947552c5064e Mon Sep 17 00:00:00 2001 From: Umar Arshad Date: Tue, 14 Jul 2020 15:55:53 -0400 Subject: [PATCH 2/7] Add batch support for CPU and OpenCL solve. Add solver batch tests Add support for batching to the CPU and OpenCL backends. Uses the MKL batching functions when MKL is enabled otherwise it iterates over all of the slices if using LAPACK. --- src/backend/cpu/Array.hpp | 4 + src/backend/cpu/lapack_helper.hpp | 1 + src/backend/cpu/solve.cpp | 198 ++++++++++++++++++++++--- src/backend/opencl/cpu/cpu_helper.hpp | 1 + src/backend/opencl/cpu/cpu_solve.cpp | 169 +++++++++++++++++++-- src/backend/opencl/solve.cpp | 95 +++++++----- test/solve_dense.cpp | 204 +++++++++++++++++++++++--- 7 files changed, 585 insertions(+), 87 deletions(-) diff --git a/src/backend/cpu/Array.hpp b/src/backend/cpu/Array.hpp index 8335e325c9..fd8ca3dce3 100644 --- a/src/backend/cpu/Array.hpp +++ b/src/backend/cpu/Array.hpp @@ -22,6 +22,7 @@ #include #include +#include #include #include #include @@ -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/lapack_helper.hpp b/src/backend/cpu/lapack_helper.hpp index a7bc77aaf3..e9b509f921 100644 --- 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 #include #else #ifdef __APPLE__ diff --git a/src/backend/cpu/solve.cpp b/src/backend/cpu/solve.cpp index d9fb586782..feac9737d5 100644 --- a/src/backend/cpu/solve.cpp +++ b/src/backend/cpu/solve.cpp @@ -17,6 +17,9 @@ #include #include #include +#include +#include +#include using af::dim4; @@ -29,6 +32,21 @@ template using gels_func_def = int (*)(ORDER_TYPE, char, int, int, int, T *, int, T *, int); +#ifdef USE_MKL +template +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 +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 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 USE_MKL + +template +struct mkl_type { + using type = T; +}; +template<> +struct mkl_type> { + using type = MKL_Complex8; +}; +template<> +struct mkl_type> { + using type = MKL_Complex16; +}; + +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wnoexcept-type" +template +getrf_batch_strided_func_def getrf_batch_strided_func(); + +template<> +getrf_batch_strided_func_def getrf_batch_strided_func() { + return &sgetrf_batch_strided; +} +template<> +getrf_batch_strided_func_def getrf_batch_strided_func() { + return &dgetrf_batch_strided; +} +template<> +getrf_batch_strided_func_def +getrf_batch_strided_func() { + return &cgetrf_batch_strided; +} +template<> +getrf_batch_strided_func_def +getrf_batch_strided_func() { + return &zgetrf_batch_strided; +} + +template +getrs_batch_strided_func_def getrs_batch_strided_func(); + +template<> +getrs_batch_strided_func_def getrs_batch_strided_func() { + return &sgetrs_batch_strided; +} +template<> +getrs_batch_strided_func_def getrs_batch_strided_func() { + return &dgetrs_batch_strided; +} +template<> +getrs_batch_strided_func_def +getrs_batch_strided_func() { + return &cgetrs_batch_strided; +} +template<> +getrs_batch_strided_func_def +getrs_batch_strided_func() { + return &zgetrs_batch_strided; +} + +#pragma GCC diagnostic pop +#endif + SOLVE_FUNC_DEF(getrs) SOLVE_FUNC(getrs, float, s) SOLVE_FUNC(getrs, double, d) @@ -109,6 +191,60 @@ Array triangleSolve(const Array &A, const Array &b, return B; } +#ifdef USE_MKL + +template +Array generalSolveBatched(const Array &a, const Array &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 ipiv(MN * batches); + int ipivstride = MN; + + int ldb = b.strides()[1]; + int bstride = b.strides()[2]; + + vector info(batches, 0); + + char trans = 'N'; + + Array A = copyArray(a); + Array B = copyArray(b); + + auto getrf_rs = [](char TRANS, int M, int N, int K, Param a, int LDA, + int ASTRIDE, vector IPIV, int IPIVSTRIDE, + Param b, int LDB, int BSTRIDE, int BATCH_SIZE, + vector INFO) { + getrf_batch_strided_func::type>()( + &M, &N, reinterpret_cast::type *>(a.get()), + &LDA, &ASTRIDE, IPIV.data(), &IPIVSTRIDE, &BATCH_SIZE, INFO.data()); + + getrs_batch_strided_func::type>()( + &TRANS, &M, &K, + reinterpret_cast::type *>(a.get()), &LDA, + &ASTRIDE, IPIV.data(), &IPIVSTRIDE, + reinterpret_cast::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 Array solve(const Array &a, const Array &b, const af_mat_prop options) { @@ -116,10 +252,20 @@ Array solve(const Array &a, const Array &b, return triangleSolve(a, b, options); } +#ifdef USE_MKL + 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 A = copyArray(a); @@ -129,27 +275,37 @@ Array solve(const Array &a, const Array &b, ? copyArray(b) : padArrayBorders(b, NullShape, endPadding, AF_PAD_ZERO)); - if (M == N) { - Array pivot = createEmptyArray(dim4(N, 1, 1)); - - auto func = [=](Param A, Param B, Param pivot, int N, - int K) { - gesv_func()(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 A, Param B, int M, int N, int K) { - int sM = A.strides(1); - int sN = A.strides(2) / sM; - - gels_func()(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 pA(A.get() + A.strides()[2] * j + A.strides()[3] * i, + A.dims(), A.strides()); + Param pB(B.get() + B.strides()[2] * j + B.strides()[3] * i, + B.dims(), B.strides()); + if (M == N) { + Array pivot = createEmptyArray(dim4(N, 1, 1)); + + auto func = [](Param A, Param B, Param pivot, int N, + int K) { + gesv_func()(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 A, Param B, int M, int N, int K) { + int sM = A.dims(0); + int sN = A.dims(1); + + gels_func()(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; } diff --git a/src/backend/opencl/cpu/cpu_helper.hpp b/src/backend/opencl/cpu/cpu_helper.hpp index 8ca6a4928c..b614e53be1 100644 --- a/src/backend/opencl/cpu/cpu_helper.hpp +++ b/src/backend/opencl/cpu/cpu_helper.hpp @@ -28,6 +28,7 @@ #define LAPACK_NAME(fn) LAPACKE_##fn #ifdef USE_MKL +#include #include #else #ifdef __APPLE__ diff --git a/src/backend/opencl/cpu/cpu_solve.cpp b/src/backend/opencl/cpu/cpu_solve.cpp index b9f2fc9933..31fbaddc62 100644 --- a/src/backend/opencl/cpu/cpu_solve.cpp +++ b/src/backend/opencl/cpu/cpu_solve.cpp @@ -12,6 +12,8 @@ #include #include #include +#include +#include namespace opencl { namespace cpu { @@ -23,6 +25,21 @@ template using gels_func_def = int (*)(ORDER_TYPE, char, int, int, int, T *, int, T *, int); +#ifdef USE_MKL +template +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 +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 using getrs_func_def = int (*)(ORDER_TYPE, char, int, int, const T *, int, const int *, T *, int); @@ -53,6 +70,70 @@ SOLVE_FUNC(gels, double, d) SOLVE_FUNC(gels, cfloat, c) SOLVE_FUNC(gels, cdouble, z) +#ifdef USE_MKL + +template +struct mkl_type { + using type = T; +}; +template<> +struct mkl_type { + using type = MKL_Complex8; +}; +template<> +struct mkl_type { + using type = MKL_Complex16; +}; + +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wnoexcept-type" +template +getrf_batch_strided_func_def getrf_batch_strided_func(); + +template<> +getrf_batch_strided_func_def getrf_batch_strided_func() { + return &sgetrf_batch_strided; +} +template<> +getrf_batch_strided_func_def getrf_batch_strided_func() { + return &dgetrf_batch_strided; +} +template<> +getrf_batch_strided_func_def +getrf_batch_strided_func() { + return &cgetrf_batch_strided; +} +template<> +getrf_batch_strided_func_def +getrf_batch_strided_func() { + return &zgetrf_batch_strided; +} + +template +getrs_batch_strided_func_def getrs_batch_strided_func(); + +template<> +getrs_batch_strided_func_def getrs_batch_strided_func() { + return &sgetrs_batch_strided; +} +template<> +getrs_batch_strided_func_def getrs_batch_strided_func() { + return &dgetrs_batch_strided; +} +template<> +getrs_batch_strided_func_def +getrs_batch_strided_func() { + return &cgetrs_batch_strided; +} +template<> +getrs_batch_strided_func_def +getrs_batch_strided_func() { + return &zgetrs_batch_strided; +} + +#pragma GCC diagnostic pop +#endif + SOLVE_FUNC_DEF(getrs) SOLVE_FUNC(getrs, float, s) SOLVE_FUNC(getrs, double, d) @@ -102,6 +183,55 @@ Array triangleSolve(const Array &A, const Array &b, return B; } +#ifdef USE_MKL + +template +Array generalSolveBatched(const Array &a, const Array &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 ipiv(MN * batches); + int ipivstride = MN; + + int ldb = b.strides()[1]; + int bstride = b.strides()[2]; + + vector info(batches, 0); + + char trans = 'N'; + + Array A = copyArray(a); + Array B = copyArray(b); + + mapped_ptr aPtr = A.getMappedPtr(); + mapped_ptr bPtr = B.getMappedPtr(); + + getrf_batch_strided_func::type>()( + &M, &N, reinterpret_cast::type *>(aPtr.get()), + &lda, &astride, ipiv.data(), &ipivstride, &batches, info.data()); + + getrs_batch_strided_func::type>()( + &trans, &M, &K, + reinterpret_cast::type *>(aPtr.get()), &lda, + &astride, ipiv.data(), &ipivstride, + reinterpret_cast::type *>(bPtr.get()), &ldb, + &bstride, &batches, info.data()); + + return B; +} +#endif + template Array solve(const Array &a, const Array &b, const af_mat_prop options) { @@ -109,8 +239,18 @@ Array solve(const Array &a, const Array &b, return triangleSolve(a, b, options); } +#ifdef USE_MKL + if (a.dims()[2] > 1 || a.dims()[3] > 1) { + return generalSolveBatched(a, b, options); + } +#endif + const dim4 NullShape(0, 0, 0, 0); + dim4 aDims = a.dims(); + int batchz = aDims[2]; + int batchw = aDims[3]; + int M = a.dims()[0]; int N = a.dims()[1]; int K = b.dims()[1]; @@ -124,18 +264,25 @@ Array solve(const Array &a, const Array &b, mapped_ptr aPtr = A.getMappedPtr(); mapped_ptr bPtr = B.getMappedPtr(); - if (M == N) { - std::vector pivot(N); - gesv_func()(AF_LAPACK_COL_MAJOR, N, K, aPtr.get(), A.strides()[1], - &pivot.front(), bPtr.get(), B.strides()[1]); - } else { - int sM = a.strides()[1]; - int sN = a.strides()[2] / sM; - - gels_func()(AF_LAPACK_COL_MAJOR, 'N', M, N, K, aPtr.get(), - A.strides()[1], bPtr.get(), max(sM, sN)); - B.resetDims(dim4(N, K)); + for (int i = 0; i < batchw; i++) { + for (int j = 0; j < batchz; j++) { + auto pA = aPtr.get() + A.strides()[2] * j + A.strides()[3] * i; + auto pB = bPtr.get() + B.strides()[2] * j + B.strides()[3] * i; + + if (M == N) { + std::vector pivot(N); + gesv_func()(AF_LAPACK_COL_MAJOR, N, K, pA, A.strides()[1], + &pivot.front(), pB, B.strides()[1]); + } else { + int sM = a.strides()[1]; + int sN = a.strides()[2] / sM; + + gels_func()(AF_LAPACK_COL_MAJOR, 'N', M, N, K, pA, + A.strides()[1], pB, max(sM, sN)); + } + } } + if (M != N) { B.resetDims(dim4(N, K, B.dims()[2], B.dims()[3])); } return B; } diff --git a/src/backend/opencl/solve.cpp b/src/backend/opencl/solve.cpp index bedd987287..ad73e21d27 100644 --- a/src/backend/opencl/solve.cpp +++ b/src/backend/opencl/solve.cpp @@ -25,6 +25,13 @@ #include #include +#include +#include + +using cl::Buffer; +using std::min; +using std::vector; + namespace opencl { template @@ -35,13 +42,13 @@ Array solveLU(const Array &A, const Array &pivot, const Array &b, int N = A.dims()[0]; int NRHS = b.dims()[1]; - std::vector ipiv(N); + vector ipiv(N); copyData(&ipiv[0], pivot); Array B = copyArray(b); - const cl::Buffer *A_buf = A.get(); - cl::Buffer *B_buf = B.get(); + const Buffer *A_buf = A.get(); + Buffer *B_buf = B.get(); int info = 0; magma_getrs_gpu(MagmaNoTrans, N, NRHS, (*A_buf)(), A.getOffset(), @@ -52,26 +59,38 @@ Array solveLU(const Array &A, const Array &pivot, const Array &b, template Array generalSolve(const Array &a, const Array &b) { - dim4 iDims = a.dims(); - int M = iDims[0]; - int N = iDims[1]; - int MN = std::min(M, N); - std::vector ipiv(MN); + dim4 aDims = a.dims(); + int batchz = aDims[2]; + int batchw = aDims[3]; Array A = copyArray(a); Array B = copyArray(b); - cl::Buffer *A_buf = A.get(); - int info = 0; - cl_command_queue q = getQueue()(); - magma_getrf_gpu(M, N, (*A_buf)(), A.getOffset(), A.strides()[1], - &ipiv[0], q, &info); - - cl::Buffer *B_buf = B.get(); - int K = B.dims()[1]; - magma_getrs_gpu(MagmaNoTrans, M, K, (*A_buf)(), A.getOffset(), - A.strides()[1], &ipiv[0], (*B_buf)(), B.getOffset(), - B.strides()[1], q, &info); + for (int i = 0; i < batchw; i++) { + for (int j = 0; j < batchz; j++) { + int M = aDims[0]; + int N = aDims[1]; + int MN = min(M, N); + vector ipiv(MN); + + Buffer *A_buf = A.get(); + int info = 0; + cl_command_queue q = getQueue()(); + auto aoffset = + A.getOffset() + j * A.strides()[2] + i * A.strides()[3]; + magma_getrf_gpu(M, N, (*A_buf)(), aoffset, A.strides()[1], + &ipiv[0], q, &info); + + Buffer *B_buf = B.get(); + int K = B.dims()[1]; + + auto boffset = + B.getOffset() + j * B.strides()[2] + i * B.strides()[3]; + magma_getrs_gpu(MagmaNoTrans, M, K, (*A_buf)(), aoffset, + A.strides()[1], &ipiv[0], (*B_buf)(), boffset, + B.strides()[1], q, &info); + } + } return B; } @@ -80,7 +99,7 @@ Array leastSquares(const Array &a, const Array &b) { int M = a.dims()[0]; int N = a.dims()[1]; int K = b.dims()[1]; - int MN = std::min(M, N); + int MN = min(M, N); Array B = createEmptyArray(dim4()); gpu_blas_trsm_func gpu_blas_trsm; @@ -117,12 +136,12 @@ Array leastSquares(const Array &a, const Array &b) { int NUM = (2 * MN + ((M + 31) / 32) * 32) * NB; Array tmp = createEmptyArray(dim4(NUM)); - std::vector h_tau(MN); + vector h_tau(MN); - int info = 0; - cl::Buffer *dA = A.get(); - cl::Buffer *dT = tmp.get(); - cl::Buffer *dB = B.get(); + int info = 0; + Buffer *dA = A.get(); + Buffer *dT = tmp.get(); + Buffer *dB = B.get(); magma_geqrf3_gpu(A.dims()[0], A.dims()[1], (*dA)(), A.getOffset(), A.strides()[1], &h_tau[0], (*dT)(), tmp.getOffset(), @@ -147,7 +166,7 @@ Array leastSquares(const Array &a, const Array &b) { #if UNMQR int lwork = (B.dims()[0] - A.dims()[0] + NB) * (B.dims()[1] + 2 * NB); - std::vector h_work(lwork); + vector h_work(lwork); B.resetDims(dim4(N, K)); magma_unmqr_gpu(MagmaLeft, MagmaNoTrans, B.dims()[0], B.dims()[1], A.dims()[0], (*dA)(), A.getOffset(), A.strides()[1], @@ -156,7 +175,7 @@ Array leastSquares(const Array &a, const Array &b) { queue, &info); #else A.resetDims(dim4(N, M)); - magma_ungqr_gpu(A.dims()[0], A.dims()[1], std::min(M, N), (*dA)(), + magma_ungqr_gpu(A.dims()[0], A.dims()[1], min(M, N), (*dA)(), A.getOffset(), A.strides()[1], &h_tau[0], (*dT)(), tmp.getOffset(), NB, queue, &info); @@ -178,18 +197,18 @@ Array leastSquares(const Array &a, const Array &b) { Array A = copyArray(a); B = copyArray(b); - int MN = std::min(M, N); + int MN = min(M, N); int NB = magma_get_geqrf_nb(M); int NUM = (2 * MN + ((N + 31) / 32) * 32) * NB; Array tmp = createEmptyArray(dim4(NUM)); - std::vector h_tau(NUM); + vector h_tau(NUM); - int info = 0; - cl::Buffer *A_buf = A.get(); - cl::Buffer *B_buf = B.get(); - cl::Buffer *dT = tmp.get(); + int info = 0; + Buffer *A_buf = A.get(); + Buffer *B_buf = B.get(); + Buffer *dT = tmp.get(); magma_geqrf3_gpu(M, N, (*A_buf)(), A.getOffset(), A.strides()[1], &h_tau[0], (*dT)(), tmp.getOffset(), getQueue()(), @@ -198,7 +217,7 @@ Array leastSquares(const Array &a, const Array &b) { int NRHS = B.dims()[1]; int lhwork = (M - N + NB) * (NRHS + NB) + NRHS * NB; - std::vector h_work(lhwork); + vector h_work(lhwork); h_work[0] = scalar(lhwork); magma_unmqr_gpu(MagmaLeft, MagmaConjTrans, M, NRHS, N, (*A_buf)(), @@ -211,8 +230,8 @@ Array leastSquares(const Array &a, const Array &b) { tmp.getOffset() + NB * MN, NB, 0, queue); if (getActivePlatform() == AFCL_PLATFORM_NVIDIA) { - Array AT = transpose(A, true); - cl::Buffer *AT_buf = AT.get(); + Array AT = transpose(A, true); + Buffer *AT_buf = AT.get(); OPENCL_BLAS_CHECK(gpu_blas_trsm( OPENCL_BLAS_SIDE_LEFT, OPENCL_BLAS_TRIANGLE_LOWER, OPENCL_BLAS_CONJ_TRANS, OPENCL_BLAS_NON_UNIT_DIAGONAL, N, NRHS, @@ -243,8 +262,8 @@ Array triangleSolve(const Array &A, const Array &b, int N = B.dims()[0]; int NRHS = B.dims()[1]; - const cl::Buffer *A_buf = A.get(); - cl::Buffer *B_buf = B.get(); + const Buffer *A_buf = A.get(); + Buffer *B_buf = B.get(); cl_event event = 0; cl_command_queue queue = getQueue()(); diff --git a/test/solve_dense.cpp b/test/solve_dense.cpp index 5014357566..0820ed51da 100644 --- a/test/solve_dense.cpp +++ b/test/solve_dense.cpp @@ -12,9 +12,163 @@ // issue https://github.com/arrayfire/arrayfire/issues/1617 #include + #include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include #include -#include "solve_common.hpp" +#include +#include + +using af::array; +using af::cdouble; +using af::cfloat; +using af::deviceGC; +using af::dim4; +using af::dtype_traits; +using af::setDevice; +using af::sum; +using std::abs; +using std::cout; +using std::endl; +using std::string; +using std::vector; + +template +void solveTester(const int m, const int n, const int k, const int b, double eps, + int targetDevice = -1) { + if (targetDevice >= 0) setDevice(targetDevice); + + deviceGC(); + + SUPPORTED_TYPE_CHECK(T); + if (noLAPACKTests()) return; + +#if 1 + array A = cpu_randu(dim4(m, n, b)); + array X0 = cpu_randu(dim4(n, k, b)); +#else + array A = randu(m, n, (dtype)dtype_traits::af_type); + array X0 = randu(n, k, (dtype)dtype_traits::af_type); +#endif + array B0 = matmul(A, X0); + + //! [ex_solve] + array X1 = solve(A, B0); + //! [ex_solve] + + //! [ex_solve_recon] + array B1 = matmul(A, X1); + //! [ex_solve_recon] + + ASSERT_NEAR( + 0, + sum::base_type>(abs(real(B0 - B1))) / (m * k), + eps); + ASSERT_NEAR( + 0, + sum::base_type>(abs(imag(B0 - B1))) / (m * k), + eps); +} + +template +void solveLUTester(const int n, const int k, double eps, + int targetDevice = -1) { + if (targetDevice >= 0) setDevice(targetDevice); + + deviceGC(); + + SUPPORTED_TYPE_CHECK(T); + if (noLAPACKTests()) return; + +#if 1 + array A = cpu_randu(dim4(n, n)); + array X0 = cpu_randu(dim4(n, k)); +#else + array A = randu(n, n, (dtype)dtype_traits::af_type); + array X0 = randu(n, k, (dtype)dtype_traits::af_type); +#endif + array B0 = matmul(A, X0); + + //! [ex_solve_lu] + array A_lu, pivot; + lu(A_lu, pivot, A); + array X1 = solveLU(A_lu, pivot, B0); + //! [ex_solve_lu] + + array B1 = matmul(A, X1); + + ASSERT_NEAR( + 0, + sum::base_type>(abs(real(B0 - B1))) / (n * k), + eps); + ASSERT_NEAR( + 0, + sum::base_type>(abs(imag(B0 - B1))) / (n * k), + eps); +} + +template +void solveTriangleTester(const int n, const int k, bool is_upper, double eps, + int targetDevice = -1) { + if (targetDevice >= 0) setDevice(targetDevice); + + deviceGC(); + + SUPPORTED_TYPE_CHECK(T); + if (noLAPACKTests()) return; + +#if 1 + array A = cpu_randu(dim4(n, n)); + array X0 = cpu_randu(dim4(n, k)); +#else + array A = randu(n, n, (dtype)dtype_traits::af_type); + array X0 = randu(n, k, (dtype)dtype_traits::af_type); +#endif + + array L, U, pivot; + lu(L, U, pivot, A); + + array AT = is_upper ? U : L; + array B0 = matmul(AT, X0); + array X1; + + if (is_upper) { + //! [ex_solve_upper] + array X = solve(AT, B0, AF_MAT_UPPER); + //! [ex_solve_upper] + + X1 = X; + } else { + //! [ex_solve_lower] + array X = solve(AT, B0, AF_MAT_LOWER); + //! [ex_solve_lower] + + X1 = X; + } + + array B1 = matmul(AT, X1); + + ASSERT_NEAR( + 0, + sum::base_type>(af::abs(real(B0 - B1))) / + (n * k), + eps); + ASSERT_NEAR( + 0, + sum::base_type>(af::abs(imag(B0 - B1))) / + (n * k), + eps); +} template class Solve : public ::testing::Test {}; @@ -37,7 +191,7 @@ double eps() { template<> double eps() { - return 0.01f; + return 0.015f; } template<> @@ -46,51 +200,67 @@ double eps() { } TYPED_TEST(Solve, Square) { - solveTester(100, 100, 10, eps()); + solveTester(100, 100, 10, 1, eps()); } TYPED_TEST(Solve, SquareMultipleOfTwo) { - solveTester(96, 96, 16, eps()); + solveTester(96, 96, 16, 1, eps()); } TYPED_TEST(Solve, SquareLarge) { - solveTester(1000, 1000, 10, eps()); + solveTester(1000, 1000, 10, 1, eps()); } TYPED_TEST(Solve, SquareMultipleOfTwoLarge) { - solveTester(2048, 2048, 32, eps()); + solveTester(2048, 2048, 32, 1, eps()); +} + +TYPED_TEST(Solve, SquareBatch) { + solveTester(100, 100, 10, 10, eps()); +} + +TYPED_TEST(Solve, SquareMultipleOfTwoBatch) { + solveTester(96, 96, 16, 10, eps()); +} + +TYPED_TEST(Solve, SquareLargeBatch) { + solveTester(1000, 1000, 10, 10, eps()); +} + +TYPED_TEST(Solve, SquareMultipleOfTwoLargeBatch) { + solveTester(2048, 2048, 32, 10, eps()); } TYPED_TEST(Solve, LeastSquaresUnderDetermined) { - solveTester(80, 100, 20, eps()); + solveTester(80, 100, 20, 1, eps()); } TYPED_TEST(Solve, LeastSquaresUnderDeterminedMultipleOfTwo) { - solveTester(96, 128, 40, eps()); + solveTester(96, 128, 40, 1, eps()); } TYPED_TEST(Solve, LeastSquaresUnderDeterminedLarge) { - solveTester(800, 1000, 200, eps()); + solveTester(800, 1000, 200, 1, eps()); } TYPED_TEST(Solve, LeastSquaresUnderDeterminedMultipleOfTwoLarge) { - solveTester(1536, 2048, 400, eps()); + solveTester(1536, 2048, 400, 1, eps()); } TYPED_TEST(Solve, LeastSquaresOverDetermined) { - solveTester(80, 60, 20, eps()); + solveTester(80, 60, 20, 1, eps()); } TYPED_TEST(Solve, LeastSquaresOverDeterminedMultipleOfTwo) { - solveTester(96, 64, 1, eps()); + solveTester(96, 64, 1, 1, eps()); } TYPED_TEST(Solve, LeastSquaresOverDeterminedLarge) { - solveTester(800, 600, 64, eps()); + solveTester(800, 600, 64, 1, eps()); } TYPED_TEST(Solve, LeastSquaresOverDeterminedMultipleOfTwoLarge) { - solveTester(1536, 1024, 1, eps()); + solveTester(1536, 1024, 1, 1, eps()); } TYPED_TEST(Solve, LU) { solveLUTester(100, 10, eps()); } @@ -152,11 +322,11 @@ int nextTargetDeviceId() { nextTargetDeviceId() % numDevices); \ tests.emplace_back(solveTriangleTester, 1000, 100, false, eps, \ nextTargetDeviceId() % numDevices); \ - tests.emplace_back(solveTester, 1000, 1000, 100, eps, \ + tests.emplace_back(solveTester, 1000, 1000, 100, 1, eps, \ nextTargetDeviceId() % numDevices); \ - tests.emplace_back(solveTester, 800, 1000, 200, eps, \ + tests.emplace_back(solveTester, 800, 1000, 200, 1, eps, \ nextTargetDeviceId() % numDevices); \ - tests.emplace_back(solveTester, 800, 600, 64, eps, \ + tests.emplace_back(solveTester, 800, 600, 64, 1, eps, \ nextTargetDeviceId() % numDevices); TEST(Solve, Threading) { From 6b287792977618304eb4b416c3351892bd4a133d Mon Sep 17 00:00:00 2001 From: Umar Arshad Date: Thu, 29 Jul 2021 22:47:32 -0400 Subject: [PATCH 3/7] Use pinned memory to copy device pointers in CUDA solve --- src/backend/cuda/solve.cu | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/backend/cuda/solve.cu b/src/backend/cuda/solve.cu index 988061ba12..f9e80efdf0 100644 --- a/src/backend/cuda/solve.cu +++ b/src/backend/cuda/solve.cu @@ -271,8 +271,8 @@ Array generalSolveBatched(const Array &a, const Array &b) { } } - auto aBatched_device_mem = memAlloc(bytes); - auto bBatched_device_mem = memAlloc(bytes); + unique_mem_ptr aBatched_device_mem(pinnedAlloc(bytes), pinnedFree); + unique_mem_ptr bBatched_device_mem(pinnedAlloc(bytes), pinnedFree); T **aBatched_device_ptrs = (T **)aBatched_device_mem.get(); T **bBatched_device_ptrs = (T **)bBatched_device_mem.get(); @@ -306,7 +306,9 @@ Array generalSolveBatched(const Array &a, const Array &b) { template Array generalSolve(const Array &a, const Array &b) { - if (a.dims()[2] > 1 || a.dims()[3] > 1) return generalSolveBatched(a, b); + if (a.dims()[2] > 1 || a.dims()[3] > 1) { + return generalSolveBatched(a, b); + } int M = a.dims()[0]; int N = a.dims()[1]; From 1efd5d0035065d07913f30753036812bdefbfafd Mon Sep 17 00:00:00 2001 From: Umar Arshad Date: Thu, 29 Jul 2021 23:05:21 -0400 Subject: [PATCH 4/7] Allow MKL as a valid entry for AF_COMPUTE_LIBRARY --- CMakeLists.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 0515e9f74f..5f1685d72e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -168,7 +168,8 @@ 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") 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") From 991f69d14a796e3de5183c3952cfdc4c5e0b48a6 Mon Sep 17 00:00:00 2001 From: Umar Arshad Date: Thu, 29 Jul 2021 23:06:30 -0400 Subject: [PATCH 5/7] Improve disabled linear algebra error message. minor header updates --- src/backend/cpu/solve.cpp | 8 ++++++-- src/backend/cuda/memory.cpp | 2 +- src/backend/opencl/Array.cpp | 3 ++- test/solve_dense.cpp | 2 +- 4 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/backend/cpu/solve.cpp b/src/backend/cpu/solve.cpp index feac9737d5..0113a8ec7d 100644 --- a/src/backend/cpu/solve.cpp +++ b/src/backend/cpu/solve.cpp @@ -318,13 +318,17 @@ namespace cpu { template Array solveLU(const Array &A, const Array &pivot, const Array &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 Array solve(const Array &a, const Array &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 index a914f9f151..969574a1c4 100644 --- a/src/backend/cuda/memory.cpp +++ b/src/backend/cuda/memory.cpp @@ -24,8 +24,8 @@ #include #include +#include #include -#include using af::dim4; using common::bytesToString; diff --git a/src/backend/opencl/Array.cpp b/src/backend/opencl/Array.cpp index 5935d51ec9..d47a0e7bec 100644 --- a/src/backend/opencl/Array.cpp +++ b/src/backend/opencl/Array.cpp @@ -23,8 +23,9 @@ #include #include +#include #include -#include +#include using af::dim4; using af::dtype_traits; diff --git a/test/solve_dense.cpp b/test/solve_dense.cpp index 0820ed51da..a63a8eede1 100644 --- a/test/solve_dense.cpp +++ b/test/solve_dense.cpp @@ -23,11 +23,11 @@ #include #include +#include #include #include #include #include -#include using af::array; using af::cdouble; From e5a1567f8e432b57fb3b9ab68d9d094d4f58abda Mon Sep 17 00:00:00 2001 From: pradeep Date: Wed, 7 Jul 2021 14:59:10 +0530 Subject: [PATCH 6/7] Use correct assert macros in solve_dense tests --- test/solve_common.hpp | 38 +++++--------------------------------- 1 file changed, 5 insertions(+), 33 deletions(-) diff --git a/test/solve_common.hpp b/test/solve_common.hpp index 341d0afc49..c464bfdc47 100644 --- a/test/solve_common.hpp +++ b/test/solve_common.hpp @@ -8,10 +8,12 @@ ********************************************************/ #pragma once + #include #include #include #include + #include #include #include @@ -25,9 +27,6 @@ using std::endl; using std::string; using std::vector; -///////////////////////////////// CPP //////////////////////////////////// -// - template void solveTester(const int m, const int n, const int k, double eps, int targetDevice = -1) { @@ -55,16 +54,7 @@ void solveTester(const int m, const int n, const int k, double eps, af::array B1 = af::matmul(A, X1); //! [ex_solve_recon] - ASSERT_NEAR(0, - af::sum::base_type>( - af::abs(real(B0 - B1))) / - (m * k), - eps); - ASSERT_NEAR(0, - af::sum::base_type>( - af::abs(imag(B0 - B1))) / - (m * k), - eps); + ASSERT_ARRAYS_NEAR(B0, B1, eps); } template @@ -94,16 +84,7 @@ void solveLUTester(const int n, const int k, double eps, af::array B1 = af::matmul(A, X1); - ASSERT_NEAR(0, - af::sum::base_type>( - af::abs(real(B0 - B1))) / - (n * k), - eps); - ASSERT_NEAR(0, - af::sum::base_type>( - af::abs(imag(B0 - B1))) / - (n * k), - eps); + ASSERT_ARRAYS_NEAR(B0, B1, eps); } template @@ -147,14 +128,5 @@ void solveTriangleTester(const int n, const int k, bool is_upper, double eps, af::array B1 = af::matmul(AT, X1); - ASSERT_NEAR(0, - af::sum::base_type>( - af::abs(real(B0 - B1))) / - (n * k), - eps); - ASSERT_NEAR(0, - af::sum::base_type>( - af::abs(imag(B0 - B1))) / - (n * k), - eps); + ASSERT_ARRAYS_NEAR(B0, B1, eps); } From 1ee9b962340f2555799c97cf52ce3abaab2f2e2d Mon Sep 17 00:00:00 2001 From: Umar Arshad Date: Fri, 30 Jul 2021 16:32:10 -0400 Subject: [PATCH 7/7] Check symbols in MKL to enable solve batch functionality We will first make sure that the getrf_batch_strided function is available in MKL to determine if the batch functionality can be used in ArrayFire. If it is available we will define the AF_USE_MKL_BATCH function to enable the batching functions. --- CMakeLists.txt | 2 ++ CMakeModules/InternalUtils.cmake | 5 +++++ src/backend/cpu/CMakeLists.txt | 5 +++++ src/backend/cpu/solve.cpp | 8 ++++---- src/backend/opencl/CMakeLists.txt | 3 +++ src/backend/opencl/cpu/cpu_solve.cpp | 8 ++++---- 6 files changed, 23 insertions(+), 8 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 5f1685d72e..06bfcdd995 100644 --- 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 @@ -170,6 +171,7 @@ mark_as_advanced(CLEAR CUDA_VERSION) if(${AF_BUILD_CPU} OR ${AF_BUILD_OPENCL}) 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 index fdb4a1bbe0..1c1a8e5f5f 100644 --- 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/backend/cpu/CMakeLists.txt b/src/backend/cpu/CMakeLists.txt index b899d6f887..7282d611ac 100644 --- 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/solve.cpp b/src/backend/cpu/solve.cpp index 0113a8ec7d..4d43405d55 100644 --- a/src/backend/cpu/solve.cpp +++ b/src/backend/cpu/solve.cpp @@ -32,7 +32,7 @@ template using gels_func_def = int (*)(ORDER_TYPE, char, int, int, int, T *, int, T *, int); -#ifdef USE_MKL +#ifdef AF_USE_MKL_BATCH template using getrf_batch_strided_func_def = void (*)(const MKL_INT *m, const MKL_INT *n, T *a, const MKL_INT *lda, @@ -77,7 +77,7 @@ SOLVE_FUNC(gels, double, d) SOLVE_FUNC(gels, cfloat, c) SOLVE_FUNC(gels, cdouble, z) -#ifdef USE_MKL +#ifdef AF_USE_MKL_BATCH template struct mkl_type { @@ -191,7 +191,7 @@ Array triangleSolve(const Array &A, const Array &b, return B; } -#ifdef USE_MKL +#ifdef AF_USE_MKL_BATCH template Array generalSolveBatched(const Array &a, const Array &b, @@ -252,7 +252,7 @@ Array solve(const Array &a, const Array &b, return triangleSolve(a, b, options); } -#ifdef USE_MKL +#ifdef AF_USE_MKL_BATCH if (a.dims()[2] > 1 || a.dims()[3] > 1) { return generalSolveBatched(a, b, options); } diff --git a/src/backend/opencl/CMakeLists.txt b/src/backend/opencl/CMakeLists.txt index b04572f2f3..5385f4fa1f 100644 --- a/src/backend/opencl/CMakeLists.txt +++ b/src/backend/opencl/CMakeLists.txt @@ -466,6 +466,9 @@ if(LAPACK_FOUND OR BUILD_WITH_MKL) if(BUILD_WITH_MKL) target_compile_definitions(afopencl PRIVATE USE_MKL) + if(MKL_BATCH) + target_compile_definitions(afopencl PRIVATE AF_USE_MKL_BATCH) + endif() if(AF_WITH_STATIC_MKL) target_link_libraries(afopencl PRIVATE MKL::Static) diff --git a/src/backend/opencl/cpu/cpu_solve.cpp b/src/backend/opencl/cpu/cpu_solve.cpp index 31fbaddc62..f5f2510597 100644 --- a/src/backend/opencl/cpu/cpu_solve.cpp +++ b/src/backend/opencl/cpu/cpu_solve.cpp @@ -25,7 +25,7 @@ template using gels_func_def = int (*)(ORDER_TYPE, char, int, int, int, T *, int, T *, int); -#ifdef USE_MKL +#ifdef AF_USE_MKL_BATCH template using getrf_batch_strided_func_def = void (*)(const MKL_INT *m, const MKL_INT *n, T *a, const MKL_INT *lda, @@ -70,7 +70,7 @@ SOLVE_FUNC(gels, double, d) SOLVE_FUNC(gels, cfloat, c) SOLVE_FUNC(gels, cdouble, z) -#ifdef USE_MKL +#ifdef AF_USE_MKL_BATCH template struct mkl_type { @@ -183,7 +183,7 @@ Array triangleSolve(const Array &A, const Array &b, return B; } -#ifdef USE_MKL +#ifdef AF_USE_MKL_BATCH template Array generalSolveBatched(const Array &a, const Array &b, @@ -239,7 +239,7 @@ Array solve(const Array &a, const Array &b, return triangleSolve(a, b, options); } -#ifdef USE_MKL +#ifdef AF_USE_MKL_BATCH if (a.dims()[2] > 1 || a.dims()[3] > 1) { return generalSolveBatched(a, b, options); }