From 902438f0f12e8490fa0a84effec7fc63b476b44b Mon Sep 17 00:00:00 2001 From: Zheming Jin Date: Fri, 14 Aug 2026 08:51:53 -0700 Subject: [PATCH] [lapack][rocsolver] Use the 64-bit pivot API for getrf and getrs oneMath types ipiv as int64_t while the legacy rocSOLVER entry points take rocblas_int, so getrf and getrs each allocated a temporary 32-bit array and ran a cast kernel around the call. In the USM getrs that temporary also forced a queue.wait() before it could be released. rocSOLVER has accepted int64_t pivots since 3.26 (ROCm 6.2), so both routines now pass the user's ipiv straight through, which removes the four temporary allocations, their cast kernels and the blocking wait. The 64-bit getrf reports info as int64_t rather than int, so the devinfo readback is templated on the info type. Earlier rocSOLVER versions keep the conversion path, leaving the minimum supported ROCm unchanged. Because the 64-bit entry points take the dimensions as int64_t, getrf and getrs no longer need the overflow_check that rejected sizes above the 32-bit limit. Note that actually factorizing a matrix with more than 2^28 rows additionally requires the rocBLAS fix from ROCm/rocm-libraries#10831: its 64-bit scal launcher advances the alpha scalar pointer once per n-chunk, so getrf_64 either faults or returns a silently wrong factorization for longer columns. With that fix applied, an m = 2^31 + 1024 factorization returns the correct pivot index 2147484161, which is the case the legacy 32-bit array could not represent. Co-authored-by: Cursor --- .../backends/rocsolver/rocsolver_helper.hpp | 29 ++- .../backends/rocsolver/rocsolver_lapack.cpp | 180 ++++++++++++++++++ 2 files changed, 203 insertions(+), 6 deletions(-) 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,