diff --git a/src/lapack/backends/rocsolver/rocsolver_helper.hpp b/src/lapack/backends/rocsolver/rocsolver_helper.hpp index 5d4e6e821..ffd3c0d1a 100644 --- a/src/lapack/backends/rocsolver/rocsolver_helper.hpp +++ b/src/lapack/backends/rocsolver/rocsolver_helper.hpp @@ -255,16 +255,33 @@ struct RocmEquivalentType> { using Type = rocblas_double_complex; }; +/* 64-bit pivot API */ + +#if !defined(ROCSOLVER_VERSION_MAJOR) +#define ONEMATH_ROCSOLVER_VERSION 0 +#else +#define ONEMATH_ROCSOLVER_VERSION (ROCSOLVER_VERSION_MAJOR * 100 + ROCSOLVER_VERSION_MINOR) +#endif + +// rocSOLVER exposes getrf and getrs entry points taking int64_t pivots, which lets oneMath hand +// the user's ipiv array straight through instead of converting it. They were added in rocSOLVER +// 3.26 (ROCm 6.2); earlier versions keep the conversion path, so the minimum supported ROCm is +// unchanged. +#define ONEMATH_ROCSOLVER_HAS_64BIT_PIVOTS (ONEMATH_ROCSOLVER_VERSION >= 326) + /* devinfo */ -inline int get_rocsolver_devinfo(sycl::queue& queue, sycl::buffer& devInfo) { - sycl::host_accessor dev_info_{ devInfo }; +// The 64-bit entry points report info as int64_t, the legacy ones as int. +template +inline INFO_T get_rocsolver_devinfo(sycl::queue& queue, sycl::buffer& devInfo) { + sycl::host_accessor dev_info_{ devInfo }; return dev_info_[0]; } -inline int get_rocsolver_devinfo(sycl::queue& queue, const int* devInfo) { - int dev_info_; - queue.memcpy(&dev_info_, devInfo, sizeof(int)); +template +inline INFO_T get_rocsolver_devinfo(sycl::queue& queue, const INFO_T* devInfo) { + INFO_T dev_info_; + queue.memcpy(&dev_info_, devInfo, sizeof(INFO_T)); queue.wait(); return dev_info_; } @@ -273,7 +290,7 @@ template inline void lapack_info_check(sycl::queue& queue, DEVINFO_T devinfo, const char* func_name, const char* cufunc_name) { queue.wait(); - const int devinfo_ = get_rocsolver_devinfo(queue, devinfo); + const auto devinfo_ = get_rocsolver_devinfo(queue, devinfo); if (devinfo_ > 0) throw oneapi::math::lapack::computation_error( func_name, std::string(cufunc_name) + " failed with info = " + std::to_string(devinfo_), diff --git a/src/lapack/backends/rocsolver/rocsolver_lapack.cpp b/src/lapack/backends/rocsolver/rocsolver_lapack.cpp index 5b0c265b2..c0bc5e723 100644 --- a/src/lapack/backends/rocsolver/rocsolver_lapack.cpp +++ b/src/lapack/backends/rocsolver/rocsolver_lapack.cpp @@ -132,6 +132,49 @@ GEQRF_LAUNCHER(std::complex, rocsolver_zgeqrf) #undef GEQRF_LAUNCHER +#if ONEMATH_ROCSOLVER_HAS_64BIT_PIVOTS + +template +void getrf(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, std::int64_t n, + sycl::buffer& a, std::int64_t lda, sycl::buffer& ipiv, + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + sycl::buffer devInfo{ 1 }; + + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + auto ipiv_acc = ipiv.template get_access(cgh); + auto devInfo_acc = devInfo.template get_access(cgh); + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto a_ = sc.get_mem(a_acc); + auto ipiv_ = sc.get_mem(ipiv_acc); + auto devInfo_ = sc.get_mem(devInfo_acc); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, m, n, a_, lda, ipiv_, + devInfo_); + }); + }); + lapack_info_check(queue, devInfo, __func__, func_name); +} + +#define GETRF_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void getrf(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, \ + std::int64_t lda, sycl::buffer& ipiv, sycl::buffer& scratchpad, \ + std::int64_t scratchpad_size) { \ + getrf(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, ipiv, scratchpad, \ + scratchpad_size); \ + } + +GETRF_LAUNCHER(float, rocsolver_sgetrf_64) +GETRF_LAUNCHER(double, rocsolver_dgetrf_64) +GETRF_LAUNCHER(std::complex, rocsolver_cgetrf_64) +GETRF_LAUNCHER(std::complex, rocsolver_zgetrf_64) + +#undef GETRF_LAUNCHER + +#else // no 64-bit pivot API: factorise into a temporary 32-bit array and widen it + template void getrf(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, std::int64_t lda, sycl::buffer& ipiv, @@ -188,6 +231,8 @@ GETRF_LAUNCHER(std::complex, rocsolver_zgetrf) #undef GETRF_LAUNCHER +#endif // ONEMATH_ROCSOLVER_HAS_64BIT_PIVOTS + void getri(sycl::queue& queue, std::int64_t n, sycl::buffer>& a, std::int64_t lda, sycl::buffer& ipiv, sycl::buffer>& scratchpad, std::int64_t scratchpad_size) { @@ -209,6 +254,50 @@ void getri(sycl::queue& queue, std::int64_t n, sycl::buffer throw unimplemented("lapack", "getri"); } +#if ONEMATH_ROCSOLVER_HAS_64BIT_PIVOTS + +template +inline void getrs(const char* func_name, Func func, sycl::queue& queue, + oneapi::math::transpose trans, std::int64_t n, std::int64_t nrhs, + sycl::buffer& a, std::int64_t lda, sycl::buffer& ipiv, + sycl::buffer& b, std::int64_t ldb, sycl::buffer& scratchpad, + std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + auto ipiv_acc = ipiv.template get_access(cgh); + auto b_acc = b.template get_access(cgh); + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto a_ = sc.get_mem(a_acc); + auto ipiv_ = sc.get_mem(ipiv_acc); + auto b_ = sc.get_mem(b_acc); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, get_rocblas_operation(trans), + n, nrhs, a_, lda, ipiv_, b_, ldb); + }); + }); +} + +#define GETRS_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void getrs(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, \ + std::int64_t nrhs, sycl::buffer& a, std::int64_t lda, \ + sycl::buffer& ipiv, sycl::buffer& b, std::int64_t ldb, \ + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { \ + getrs(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, trans, n, nrhs, a, lda, ipiv, b, ldb, \ + scratchpad, scratchpad_size); \ + } + +GETRS_LAUNCHER(float, rocsolver_sgetrs_64) +GETRS_LAUNCHER(double, rocsolver_dgetrs_64) +GETRS_LAUNCHER(std::complex, rocsolver_cgetrs_64) +GETRS_LAUNCHER(std::complex, rocsolver_zgetrs_64) + +#undef GETRS_LAUNCHER + +#else // no 64-bit pivot API: narrow the pivots into a temporary 32-bit array + template inline void getrs(const char* func_name, Func func, sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, std::int64_t nrhs, @@ -264,6 +353,8 @@ GETRS_LAUNCHER(std::complex, rocsolver_zgetrs) #undef GETRS_LAUNCHER +#endif // ONEMATH_ROCSOLVER_HAS_64BIT_PIVOTS + template inline void gesvd(const char* func_name, Func func, sycl::queue& queue, oneapi::math::jobsvd jobu, oneapi::math::jobsvd jobvt, std::int64_t m, std::int64_t n, sycl::buffer& a, @@ -1259,6 +1350,50 @@ GEQRF_LAUNCHER_USM(std::complex, rocsolver_zgeqrf) #undef GEQRF_LAUNCHER_USM +#if ONEMATH_ROCSOLVER_HAS_64BIT_PIVOTS + +template +inline sycl::event getrf(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, + std::int64_t n, T* a, std::int64_t lda, std::int64_t* ipiv, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + + std::int64_t* devInfo = (std::int64_t*)malloc_device(sizeof(std::int64_t), queue); + auto done = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto a_ = reinterpret_cast(a); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, m, n, a_, lda, ipiv, devInfo); + }); + }); + + // lapack_info_check calls queue.wait() + lapack_info_check(queue, devInfo, __func__, func_name); + free(devInfo, queue); + return done; +} + +#define GETRF_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event getrf(sycl::queue& queue, std::int64_t m, std::int64_t n, TYPE* a, \ + std::int64_t lda, std::int64_t* ipiv, TYPE* scratchpad, \ + std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return getrf(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, ipiv, scratchpad, \ + scratchpad_size, dependencies); \ + } + +GETRF_LAUNCHER_USM(float, rocsolver_sgetrf_64) +GETRF_LAUNCHER_USM(double, rocsolver_dgetrf_64) +GETRF_LAUNCHER_USM(std::complex, rocsolver_cgetrf_64) +GETRF_LAUNCHER_USM(std::complex, rocsolver_zgetrf_64) + +#undef GETRF_LAUNCHER_USM + +#else // no 64-bit pivot API: factorise into a temporary 32-bit array and widen it + template inline sycl::event getrf(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, std::int64_t n, T* a, std::int64_t lda, std::int64_t* ipiv, T* scratchpad, @@ -1320,6 +1455,8 @@ GETRF_LAUNCHER_USM(std::complex, rocsolver_zgetrf) #undef GETRF_LAUNCHER_USM +#endif // ONEMATH_ROCSOLVER_HAS_64BIT_PIVOTS + sycl::event getri(sycl::queue& queue, std::int64_t n, std::complex* a, std::int64_t lda, std::int64_t* ipiv, std::complex* scratchpad, std::int64_t scratchpad_size, const std::vector& dependencies) { @@ -1341,6 +1478,47 @@ sycl::event getri(sycl::queue& queue, std::int64_t n, std::complex* a, s throw unimplemented("lapack", "getri"); } +#if ONEMATH_ROCSOLVER_HAS_64BIT_PIVOTS + +template +inline sycl::event getrs(const char* func_name, Func func, sycl::queue& queue, + oneapi::math::transpose trans, std::int64_t n, std::int64_t nrhs, T* a, + std::int64_t lda, std::int64_t* ipiv, T* b, std::int64_t ldb, + T* scratchpad, std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + + return queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto a_ = reinterpret_cast(a); + auto b_ = reinterpret_cast(b); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, get_rocblas_operation(trans), + n, nrhs, a_, lda, ipiv, b_, ldb); + }); + }); +} + +#define GETRS_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event getrs(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, \ + std::int64_t nrhs, TYPE* a, std::int64_t lda, std::int64_t* ipiv, TYPE* b, \ + std::int64_t ldb, TYPE* scratchpad, std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return getrs(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, trans, n, nrhs, a, lda, ipiv, \ + b, ldb, scratchpad, scratchpad_size, dependencies); \ + } + +GETRS_LAUNCHER_USM(float, rocsolver_sgetrs_64) +GETRS_LAUNCHER_USM(double, rocsolver_dgetrs_64) +GETRS_LAUNCHER_USM(std::complex, rocsolver_cgetrs_64) +GETRS_LAUNCHER_USM(std::complex, rocsolver_zgetrs_64) + +#undef GETRS_LAUNCHER_USM + +#else // no 64-bit pivot API: narrow the pivots into a temporary 32-bit array + template inline sycl::event getrs(const char* func_name, Func func, sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, std::int64_t nrhs, T* a, @@ -1402,6 +1580,8 @@ GETRS_LAUNCHER_USM(std::complex, rocsolver_zgetrs) #undef GETRS_LAUNCHER_USM +#endif // ONEMATH_ROCSOLVER_HAS_64BIT_PIVOTS + template inline sycl::event gesvd(const char* func_name, Func func, sycl::queue& queue, oneapi::math::jobsvd jobu, oneapi::math::jobsvd jobvt, std::int64_t m,