diff --git a/src/lapack/backends/rocsolver/rocsolver_batch.cpp b/src/lapack/backends/rocsolver/rocsolver_batch.cpp index a84be1fe6..85579b0cd 100644 --- a/src/lapack/backends/rocsolver/rocsolver_batch.cpp +++ b/src/lapack/backends/rocsolver/rocsolver_batch.cpp @@ -29,473 +29,1151 @@ namespace math { namespace lapack { namespace rocsolver { -// BATCH BUFFER API +// rocsolver does not use scratchpad memory for any of the batched routines +// below: workspace is managed internally by the rocblas handle. -void geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, - std::int64_t lda, std::int64_t stride_a, sycl::buffer& tau, - std::int64_t stride_tau, std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "geqrf_batch"); -} -void geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, - std::int64_t lda, std::int64_t stride_a, sycl::buffer& tau, - std::int64_t stride_tau, std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "geqrf_batch"); -} -void geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer>& tau, std::int64_t stride_tau, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "geqrf_batch"); -} -void geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer>& tau, std::int64_t stride_tau, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "geqrf_batch"); -} -void getri_batch(sycl::queue& queue, std::int64_t n, sycl::buffer& a, std::int64_t lda, - std::int64_t stride_a, sycl::buffer& ipiv, std::int64_t stride_ipiv, - std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getri_batch"); -} -void getri_batch(sycl::queue& queue, std::int64_t n, sycl::buffer& a, std::int64_t lda, - std::int64_t stride_a, sycl::buffer& ipiv, std::int64_t stride_ipiv, - std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getri_batch"); -} -void getri_batch(sycl::queue& queue, std::int64_t n, sycl::buffer>& a, - std::int64_t lda, std::int64_t stride_a, sycl::buffer& ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, - sycl::buffer>& scratchpad, std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getri_batch"); -} -void getri_batch(sycl::queue& queue, std::int64_t n, sycl::buffer>& a, - std::int64_t lda, std::int64_t stride_a, sycl::buffer& ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, - sycl::buffer>& scratchpad, std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getri_batch"); -} -void getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, - std::int64_t nrhs, sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer& ipiv, std::int64_t stride_ipiv, sycl::buffer& b, - std::int64_t ldb, std::int64_t stride_b, std::int64_t batch_size, - sycl::buffer& scratchpad, std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getrs_batch"); -} -void getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, - std::int64_t nrhs, sycl::buffer& a, std::int64_t lda, - std::int64_t stride_a, sycl::buffer& ipiv, std::int64_t stride_ipiv, - sycl::buffer& b, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getrs_batch"); -} -void getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, - std::int64_t nrhs, sycl::buffer>& a, std::int64_t lda, - std::int64_t stride_a, sycl::buffer& ipiv, std::int64_t stride_ipiv, - sycl::buffer>& b, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getrs_batch"); -} -void getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, - std::int64_t nrhs, sycl::buffer>& a, std::int64_t lda, - std::int64_t stride_a, sycl::buffer& ipiv, std::int64_t stride_ipiv, - sycl::buffer>& b, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getrs_batch"); -} -void getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, - std::int64_t lda, std::int64_t stride_a, sycl::buffer& ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getrf_batch"); +namespace { + +// The rocsolver legacy api indexes pivots with 32-bit ints while oneMath uses +// 64-bit ones, so a strided batch of pivots has to be converted in either +// direction around the native call. Only the leading min(m, n) entries of each +// matrix in the batch are meaningful, so the padding between two consecutive +// pivot arrays is left untouched. + +template +inline sycl::event copy_ipiv_to_32(sycl::queue& queue, IpivAcc ipiv, Ipiv32Acc ipiv32, + std::int64_t ipiv_len, std::int64_t stride_ipiv, + std::int64_t batch_size, + const std::vector& dependencies) { + return queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.parallel_for(sycl::range<2>{ static_cast(batch_size), + static_cast(ipiv_len) }, + [=](sycl::id<2> index) { + const std::int64_t offset = index[0] * stride_ipiv + index[1]; + ipiv32[offset] = static_cast(ipiv[offset]); + }); + }); } -void getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, - std::int64_t lda, std::int64_t stride_a, sycl::buffer& ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, - sycl::buffer& scratchpad, std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getrf_batch"); + +template +inline sycl::event copy_ipiv_to_64(sycl::queue& queue, Ipiv32Acc ipiv32, IpivAcc ipiv, + std::int64_t ipiv_len, std::int64_t stride_ipiv, + std::int64_t batch_size, + const std::vector& dependencies) { + return queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.parallel_for(sycl::range<2>{ static_cast(batch_size), + static_cast(ipiv_len) }, + [=](sycl::id<2> index) { + const std::int64_t offset = index[0] * stride_ipiv + index[1]; + ipiv[offset] = static_cast(ipiv32[offset]); + }); + }); } -void getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer& ipiv, std::int64_t stride_ipiv, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getrf_batch"); + +} // namespace + +// BATCH BUFFER API + +template +inline void geqrf_batch(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, + std::int64_t n, sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, + sycl::buffer& tau, std::int64_t stride_tau, std::int64_t batch_size, + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(m, n, lda, batch_size, scratchpad_size); + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + auto tau_acc = tau.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 tau_ = sc.get_mem(tau_acc); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, m, n, a_, lda, stride_a, tau_, + stride_tau, batch_size); + }); + }); } -void getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer& ipiv, std::int64_t stride_ipiv, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "getrf_batch"); + +#define GEQRF_STRIDED_BATCH_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, \ + std::int64_t lda, std::int64_t stride_a, sycl::buffer& tau, \ + std::int64_t stride_tau, std::int64_t batch_size, \ + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { \ + geqrf_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, stride_a, tau, \ + stride_tau, batch_size, scratchpad, scratchpad_size); \ + } + +GEQRF_STRIDED_BATCH_LAUNCHER(float, rocsolver_sgeqrf_strided_batched) +GEQRF_STRIDED_BATCH_LAUNCHER(double, rocsolver_dgeqrf_strided_batched) +GEQRF_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_cgeqrf_strided_batched) +GEQRF_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_zgeqrf_strided_batched) + +#undef GEQRF_STRIDED_BATCH_LAUNCHER + +template +inline void getrf_batch(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, + std::int64_t n, sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, + sycl::buffer& ipiv, std::int64_t stride_ipiv, + std::int64_t batch_size, sycl::buffer& scratchpad, + std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(m, n, lda, batch_size, scratchpad_size); + + const std::int64_t ipiv_len = std::min(m, n); + sycl::buffer ipiv32( + sycl::range<1>{ static_cast(stride_ipiv * batch_size) }); + sycl::buffer devInfo{ sycl::range<1>{ static_cast(batch_size) } }; + + auto done = queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + auto ipiv32_acc = ipiv32.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 ipiv32_ = sc.get_mem(ipiv32_acc); + auto devInfo_ = sc.get_mem(devInfo_acc); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, m, n, a_, lda, stride_a, + ipiv32_, stride_ipiv, devInfo_, batch_size); + }); + }); + + queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(done); + auto ipiv32_acc = ipiv32.template get_access(cgh); + auto ipiv_acc = ipiv.template get_access(cgh); + cgh.parallel_for(sycl::range<2>{ static_cast(batch_size), + static_cast(ipiv_len) }, + [=](sycl::id<2> index) { + const std::int64_t offset = index[0] * stride_ipiv + index[1]; + ipiv_acc[offset] = static_cast(ipiv32_acc[offset]); + }); + }); + lapack_info_check_batch(queue, devInfo, __func__, func_name, batch_size); } -void orgqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, - sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer& tau, std::int64_t stride_tau, std::int64_t batch_size, - sycl::buffer& scratchpad, std::int64_t scratchpad_size) { - throw unimplemented("lapack", "orgqr_batch"); + +#define GETRF_STRIDED_BATCH_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, \ + std::int64_t lda, std::int64_t stride_a, sycl::buffer& ipiv, \ + std::int64_t stride_ipiv, std::int64_t batch_size, \ + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { \ + getrf_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, stride_a, ipiv, \ + stride_ipiv, batch_size, scratchpad, scratchpad_size); \ + } + +GETRF_STRIDED_BATCH_LAUNCHER(float, rocsolver_sgetrf_strided_batched) +GETRF_STRIDED_BATCH_LAUNCHER(double, rocsolver_dgetrf_strided_batched) +GETRF_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_cgetrf_strided_batched) +GETRF_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_zgetrf_strided_batched) + +#undef GETRF_STRIDED_BATCH_LAUNCHER + +template +inline void getri_batch(const char* func_name, Func func, sycl::queue& queue, std::int64_t n, + sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, + sycl::buffer& ipiv, std::int64_t stride_ipiv, + std::int64_t batch_size, sycl::buffer& scratchpad, + std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(n, lda, batch_size, scratchpad_size); + + sycl::buffer ipiv32( + sycl::range<1>{ static_cast(stride_ipiv * batch_size) }); + sycl::buffer devInfo{ sycl::range<1>{ static_cast(batch_size) } }; + + queue.submit([&](sycl::handler& cgh) { + auto ipiv32_acc = ipiv32.template get_access(cgh); + auto ipiv_acc = ipiv.template get_access(cgh); + cgh.parallel_for( + sycl::range<2>{ static_cast(batch_size), static_cast(n) }, + [=](sycl::id<2> index) { + const std::int64_t offset = index[0] * stride_ipiv + index[1]; + ipiv32_acc[offset] = static_cast(ipiv_acc[offset]); + }); + }); + + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + auto ipiv32_acc = ipiv32.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 ipiv32_ = sc.get_mem(ipiv32_acc); + auto devInfo_ = sc.get_mem(devInfo_acc); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, n, a_, lda, stride_a, ipiv32_, + stride_ipiv, devInfo_, batch_size); + }); + }); + lapack_info_check_batch(queue, devInfo, __func__, func_name, batch_size); } -void orgqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, - sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer& tau, std::int64_t stride_tau, std::int64_t batch_size, - sycl::buffer& scratchpad, std::int64_t scratchpad_size) { - throw unimplemented("lapack", "orgqr_batch"); + +#define GETRI_STRIDED_BATCH_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void getri_batch(sycl::queue& queue, std::int64_t n, sycl::buffer& a, std::int64_t lda, \ + std::int64_t stride_a, sycl::buffer& ipiv, \ + std::int64_t stride_ipiv, std::int64_t batch_size, \ + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { \ + getri_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, n, a, lda, stride_a, ipiv, \ + stride_ipiv, batch_size, scratchpad, scratchpad_size); \ + } + +GETRI_STRIDED_BATCH_LAUNCHER(float, rocsolver_sgetri_strided_batched) +GETRI_STRIDED_BATCH_LAUNCHER(double, rocsolver_dgetri_strided_batched) +GETRI_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_cgetri_strided_batched) +GETRI_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_zgetri_strided_batched) + +#undef GETRI_STRIDED_BATCH_LAUNCHER + +template +inline void getrs_batch(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, std::int64_t stride_a, + sycl::buffer& ipiv, std::int64_t stride_ipiv, + sycl::buffer& b, std::int64_t ldb, std::int64_t stride_b, + std::int64_t batch_size, sycl::buffer& scratchpad, + std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(n, nrhs, lda, ldb, batch_size, scratchpad_size); + + sycl::buffer ipiv32( + sycl::range<1>{ static_cast(stride_ipiv * batch_size) }); + + queue.submit([&](sycl::handler& cgh) { + auto ipiv32_acc = ipiv32.template get_access(cgh); + auto ipiv_acc = ipiv.template get_access(cgh); + cgh.parallel_for( + sycl::range<2>{ static_cast(batch_size), static_cast(n) }, + [=](sycl::id<2> index) { + const std::int64_t offset = index[0] * stride_ipiv + index[1]; + ipiv32_acc[offset] = static_cast(ipiv_acc[offset]); + }); + }); + + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + auto ipiv32_acc = ipiv32.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 ipiv32_ = sc.get_mem(ipiv32_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, stride_a, ipiv32_, stride_ipiv, b_, ldb, + stride_b, batch_size); + }); + }); } -void potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, - std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "potrf_batch"); + +#define GETRS_STRIDED_BATCH_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, \ + std::int64_t nrhs, sycl::buffer& a, std::int64_t lda, \ + std::int64_t stride_a, sycl::buffer& ipiv, \ + std::int64_t stride_ipiv, sycl::buffer& b, std::int64_t ldb, \ + std::int64_t stride_b, std::int64_t batch_size, \ + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { \ + getrs_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, trans, n, nrhs, a, lda, \ + stride_a, ipiv, stride_ipiv, b, ldb, stride_b, batch_size, scratchpad, \ + scratchpad_size); \ + } + +GETRS_STRIDED_BATCH_LAUNCHER(float, rocsolver_sgetrs_strided_batched) +GETRS_STRIDED_BATCH_LAUNCHER(double, rocsolver_dgetrs_strided_batched) +GETRS_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_cgetrs_strided_batched) +GETRS_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_zgetrs_strided_batched) + +#undef GETRS_STRIDED_BATCH_LAUNCHER + +// rocsolver has no batched orgqr/ungqr, so the batch is walked one matrix at a +// time inside a single host task. +template +inline void orgqr_batch(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, + std::int64_t n, std::int64_t k, sycl::buffer& a, std::int64_t lda, + std::int64_t stride_a, sycl::buffer& tau, std::int64_t stride_tau, + std::int64_t batch_size, sycl::buffer& scratchpad, + std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(m, n, k, lda, batch_size, scratchpad_size); + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + auto tau_acc = tau.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 tau_ = sc.get_mem(tau_acc); + rocblas_status err; + for (std::int64_t i = 0; i < batch_size; ++i) { + rocsolver_native_named_func(func_name, func, err, handle, m, n, k, + a_ + i * stride_a, lda, tau_ + i * stride_tau); + } + }); + }); } -void potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, - std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "potrf_batch"); + +#define ORGQR_STRIDED_BATCH_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void orgqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, \ + sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, \ + sycl::buffer& tau, std::int64_t stride_tau, std::int64_t batch_size, \ + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { \ + orgqr_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, k, a, lda, stride_a, tau, \ + stride_tau, batch_size, scratchpad, scratchpad_size); \ + } + +ORGQR_STRIDED_BATCH_LAUNCHER(float, rocsolver_sorgqr) +ORGQR_STRIDED_BATCH_LAUNCHER(double, rocsolver_dorgqr) + +#undef ORGQR_STRIDED_BATCH_LAUNCHER + +#define UNGQR_STRIDED_BATCH_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void ungqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, \ + sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, \ + sycl::buffer& tau, std::int64_t stride_tau, std::int64_t batch_size, \ + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { \ + orgqr_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, k, a, lda, stride_a, tau, \ + stride_tau, batch_size, scratchpad, scratchpad_size); \ + } + +UNGQR_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_cungqr) +UNGQR_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_zungqr) + +#undef UNGQR_STRIDED_BATCH_LAUNCHER + +template +inline void potrf_batch(const char* func_name, Func func, sycl::queue& queue, + oneapi::math::uplo uplo, std::int64_t n, sycl::buffer& a, + std::int64_t lda, std::int64_t stride_a, std::int64_t batch_size, + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(n, lda, batch_size, scratchpad_size); + + sycl::buffer devInfo{ sycl::range<1>{ static_cast(batch_size) } }; + + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.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 devInfo_ = sc.get_mem(devInfo_acc); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, get_rocblas_fill_mode(uplo), + n, a_, lda, stride_a, devInfo_, batch_size); + }); + }); + lapack_info_check_batch(queue, devInfo, __func__, func_name, batch_size); } -void potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "potrf_batch"); + +#define POTRF_STRIDED_BATCH_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, \ + sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, \ + std::int64_t batch_size, sycl::buffer& scratchpad, \ + std::int64_t scratchpad_size) { \ + potrf_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, uplo, n, a, lda, stride_a, \ + batch_size, scratchpad, scratchpad_size); \ + } + +POTRF_STRIDED_BATCH_LAUNCHER(float, rocsolver_spotrf_strided_batched) +POTRF_STRIDED_BATCH_LAUNCHER(double, rocsolver_dpotrf_strided_batched) +POTRF_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_cpotrf_strided_batched) +POTRF_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_zpotrf_strided_batched) + +#undef POTRF_STRIDED_BATCH_LAUNCHER + +template +inline void potrs_batch(const char* func_name, Func func, sycl::queue& queue, + oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, + sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, + sycl::buffer& b, std::int64_t ldb, std::int64_t stride_b, + std::int64_t batch_size, sycl::buffer& scratchpad, + std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(n, nrhs, lda, ldb, batch_size, scratchpad_size); + + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.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 b_ = sc.get_mem(b_acc); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, get_rocblas_fill_mode(uplo), + n, nrhs, a_, lda, stride_a, b_, ldb, stride_b, batch_size); + }); + }); } -void potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "potrf_batch"); + +#define POTRS_STRIDED_BATCH_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void potrs_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, \ + std::int64_t nrhs, sycl::buffer& a, std::int64_t lda, \ + std::int64_t stride_a, sycl::buffer& b, std::int64_t ldb, \ + std::int64_t stride_b, std::int64_t batch_size, \ + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { \ + potrs_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, uplo, n, nrhs, a, lda, stride_a, \ + b, ldb, stride_b, batch_size, scratchpad, scratchpad_size); \ + } + +POTRS_STRIDED_BATCH_LAUNCHER(float, rocsolver_spotrs_strided_batched) +POTRS_STRIDED_BATCH_LAUNCHER(double, rocsolver_dpotrs_strided_batched) +POTRS_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_cpotrs_strided_batched) +POTRS_STRIDED_BATCH_LAUNCHER(std::complex, rocsolver_zpotrs_strided_batched) + +#undef POTRS_STRIDED_BATCH_LAUNCHER + +// BATCH USM API + +template +inline sycl::event geqrf_batch(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 stride_a, + T* tau, std::int64_t stride_tau, std::int64_t batch_size, + T* scratchpad, std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(m, n, lda, batch_size, scratchpad_size); + 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); + auto tau_ = reinterpret_cast(tau); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, m, n, a_, lda, stride_a, tau_, + stride_tau, batch_size); + }); + }); + return done; } -void potrs_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, - sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer& b, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "potrs_batch"); + +#define GEQRF_STRIDED_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, TYPE* a, \ + std::int64_t lda, std::int64_t stride_a, TYPE* tau, \ + std::int64_t stride_tau, std::int64_t batch_size, TYPE* scratchpad, \ + std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return geqrf_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, stride_a, \ + tau, stride_tau, batch_size, scratchpad, scratchpad_size, \ + dependencies); \ + } + +GEQRF_STRIDED_BATCH_LAUNCHER_USM(float, rocsolver_sgeqrf_strided_batched) +GEQRF_STRIDED_BATCH_LAUNCHER_USM(double, rocsolver_dgeqrf_strided_batched) +GEQRF_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_cgeqrf_strided_batched) +GEQRF_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_zgeqrf_strided_batched) + +#undef GEQRF_STRIDED_BATCH_LAUNCHER_USM + +// The group api is served by the rocsolver _batched entry points, which take a +// device resident array of matrix pointers. +template +inline sycl::event geqrf_batch(const char* func_name, Func func, sycl::queue& queue, + std::int64_t* m, std::int64_t* n, T** a, std::int64_t* lda, T** tau, + std::int64_t group_count, std::int64_t* group_sizes, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + + std::int64_t batch_size = 0; + overflow_check(group_count, scratchpad_size); + for (std::int64_t i = 0; i < group_count; ++i) { + overflow_check(m[i], n[i], lda[i], group_sizes[i]); + batch_size += group_sizes[i]; + } + + // The batched entry point writes the Householder scalars of a whole group + // into one contiguous array, so they are staged here and scattered back to + // the caller supplied per-matrix arrays afterwards. + std::vector tau_len(group_count); + std::int64_t tau_stage_size = 0; + for (std::int64_t i = 0; i < group_count; ++i) { + tau_len[i] = std::min(m[i], n[i]); + tau_stage_size += tau_len[i] * group_sizes[i]; + } + + T** a_dev = (T**)malloc_device(sizeof(T*) * batch_size, queue); + T* tau_stage = (T*)malloc_device(sizeof(T) * tau_stage_size, queue); + auto done_cpy = + queue.submit([&](sycl::handler& h) { h.memcpy(a_dev, a, batch_size * sizeof(T*)); }); + + auto done = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.depends_on(done_cpy); + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto** a_ = reinterpret_cast(a_dev); + auto* tau_ = reinterpret_cast(tau_stage); + std::int64_t offset = 0; + std::int64_t tau_offset = 0; + rocblas_status err; + for (std::int64_t i = 0; i < group_count; ++i) { + const std::int64_t len = std::min(m[i], n[i]); + rocsolver_native_named_func(func_name, func, err, handle, (int)m[i], (int)n[i], + a_ + offset, (int)lda[i], tau_ + tau_offset, len, + (int)group_sizes[i]); + offset += group_sizes[i]; + tau_offset += len * group_sizes[i]; + } + }); + }); + + std::vector scatter_dependencies; + scatter_dependencies.reserve(batch_size); + for (std::int64_t i = 0, global_id = 0, tau_offset = 0; i < group_count; ++i) { + const std::int64_t len = tau_len[i]; + for (std::int64_t j = 0; j < group_sizes[i]; ++j, ++global_id, tau_offset += len) { + scatter_dependencies.push_back(queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(done); + cgh.memcpy(tau[global_id], tau_stage + tau_offset, len * sizeof(T)); + })); + } + } + + auto done_scatter = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(scatter_dependencies); + cgh.host_task([]() {}); + }); + + queue.wait(); + sycl::free(a_dev, queue); + sycl::free(tau_stage, queue); + return done_scatter; } -void potrs_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, - sycl::buffer& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer& b, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "potrs_batch"); + +#define GEQRF_GROUP_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event geqrf_batch( \ + sycl::queue& queue, std::int64_t* m, std::int64_t* n, TYPE** a, std::int64_t* lda, \ + TYPE** tau, std::int64_t group_count, std::int64_t* group_sizes, TYPE* scratchpad, \ + std::int64_t scratchpad_size, const std::vector& dependencies) { \ + return geqrf_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, tau, \ + group_count, group_sizes, scratchpad, scratchpad_size, dependencies); \ + } + +GEQRF_GROUP_BATCH_LAUNCHER_USM(float, rocsolver_sgeqrf_batched) +GEQRF_GROUP_BATCH_LAUNCHER_USM(double, rocsolver_dgeqrf_batched) +GEQRF_GROUP_BATCH_LAUNCHER_USM(std::complex, rocsolver_cgeqrf_batched) +GEQRF_GROUP_BATCH_LAUNCHER_USM(std::complex, rocsolver_zgeqrf_batched) + +#undef GEQRF_GROUP_BATCH_LAUNCHER_USM + +template +inline sycl::event getrf_batch(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 stride_a, + std::int64_t* ipiv, std::int64_t stride_ipiv, + std::int64_t batch_size, T* scratchpad, std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(m, n, lda, batch_size, scratchpad_size); + + const std::int64_t ipiv_len = std::min(m, n); + int* ipiv32 = (int*)malloc_device(sizeof(int) * stride_ipiv * batch_size, queue); + int* devInfo = (int*)malloc_device(sizeof(int) * batch_size, 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, stride_a, + ipiv32, stride_ipiv, devInfo, batch_size); + }); + }); + + auto done_casting = + copy_ipiv_to_64(queue, ipiv32, ipiv, ipiv_len, stride_ipiv, batch_size, { done }); + + try { + lapack_info_check_batch(queue, devInfo, __func__, func_name, batch_size); + } + catch (...) { + sycl::free(ipiv32, queue); + sycl::free(devInfo, queue); + throw; + } + sycl::free(ipiv32, queue); + sycl::free(devInfo, queue); + return done_casting; } -void potrs_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer>& b, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "potrs_batch"); + +#define GETRF_STRIDED_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, TYPE* a, \ + std::int64_t lda, std::int64_t stride_a, std::int64_t* ipiv, \ + std::int64_t stride_ipiv, std::int64_t batch_size, TYPE* scratchpad, \ + std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return getrf_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, stride_a, \ + ipiv, stride_ipiv, batch_size, scratchpad, scratchpad_size, \ + dependencies); \ + } + +GETRF_STRIDED_BATCH_LAUNCHER_USM(float, rocsolver_sgetrf_strided_batched) +GETRF_STRIDED_BATCH_LAUNCHER_USM(double, rocsolver_dgetrf_strided_batched) +GETRF_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_cgetrf_strided_batched) +GETRF_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_zgetrf_strided_batched) + +#undef GETRF_STRIDED_BATCH_LAUNCHER_USM + +template +inline sycl::event getrf_batch(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, std::int64_t group_count, + std::int64_t* group_sizes, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + + std::int64_t batch_size = 0; + overflow_check(group_count, scratchpad_size); + for (std::int64_t i = 0; i < group_count; ++i) { + overflow_check(m[i], n[i], lda[i], group_sizes[i]); + batch_size += group_sizes[i]; + } + + // Pivots of a group are contiguous in the 32-bit staging buffer so that the + // native call can address them with a fixed stride. + std::vector ipiv_len(group_count); + std::int64_t ipiv32_size = 0; + for (std::int64_t i = 0; i < group_count; ++i) { + ipiv_len[i] = std::min(m[i], n[i]); + ipiv32_size += ipiv_len[i] * group_sizes[i]; + } + + T** a_dev = (T**)malloc_device(sizeof(T*) * batch_size, queue); + int* ipiv32 = (int*)malloc_device(sizeof(int) * ipiv32_size, queue); + int* devInfo = (int*)malloc_device(sizeof(int) * batch_size, queue); + auto done_cpy = + queue.submit([&](sycl::handler& h) { h.memcpy(a_dev, a, batch_size * sizeof(T*)); }); + + auto done = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.depends_on(done_cpy); + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto** a_ = reinterpret_cast(a_dev); + std::int64_t offset = 0; + std::int64_t ipiv_offset = 0; + rocblas_status err; + for (std::int64_t i = 0; i < group_count; ++i) { + const std::int64_t len = std::min(m[i], n[i]); + rocsolver_native_named_func(func_name, func, err, handle, (int)m[i], (int)n[i], + a_ + offset, (int)lda[i], ipiv32 + ipiv_offset, len, + devInfo + offset, (int)group_sizes[i]); + offset += group_sizes[i]; + ipiv_offset += len * group_sizes[i]; + } + }); + }); + + std::vector casting_dependencies; + casting_dependencies.reserve(batch_size); + for (std::int64_t i = 0, global_id = 0, ipiv_offset = 0; i < group_count; ++i) { + const std::int64_t len = ipiv_len[i]; + for (std::int64_t j = 0; j < group_sizes[i]; ++j, ++global_id, ipiv_offset += len) { + std::int64_t* d_ipiv = ipiv[global_id]; + const int* d_ipiv32 = ipiv32 + ipiv_offset; + casting_dependencies.push_back(queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(done); + cgh.parallel_for(sycl::range<1>{ static_cast(len) }, + [=](sycl::id<1> index) { + d_ipiv[index] = static_cast(d_ipiv32[index]); + }); + })); + } + } + + auto done_casting = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(casting_dependencies); + cgh.host_task([]() {}); + }); + + try { + lapack_info_check_batch(queue, devInfo, __func__, func_name, batch_size); + } + catch (...) { + sycl::free(a_dev, queue); + sycl::free(ipiv32, queue); + sycl::free(devInfo, queue); + throw; + } + sycl::free(a_dev, queue); + sycl::free(ipiv32, queue); + sycl::free(devInfo, queue); + return done_casting; } -void potrs_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer>& b, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "potrs_batch"); + +#define GETRF_GROUP_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event getrf_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, TYPE** a, \ + std::int64_t* lda, std::int64_t** ipiv, std::int64_t group_count, \ + std::int64_t* group_sizes, TYPE* scratchpad, \ + std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return getrf_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, ipiv, \ + group_count, group_sizes, scratchpad, scratchpad_size, dependencies); \ + } + +GETRF_GROUP_BATCH_LAUNCHER_USM(float, rocsolver_sgetrf_batched) +GETRF_GROUP_BATCH_LAUNCHER_USM(double, rocsolver_dgetrf_batched) +GETRF_GROUP_BATCH_LAUNCHER_USM(std::complex, rocsolver_cgetrf_batched) +GETRF_GROUP_BATCH_LAUNCHER_USM(std::complex, rocsolver_zgetrf_batched) + +#undef GETRF_GROUP_BATCH_LAUNCHER_USM + +template +inline sycl::event getri_batch(const char* func_name, Func func, sycl::queue& queue, std::int64_t n, + T* a, std::int64_t lda, std::int64_t stride_a, std::int64_t* ipiv, + std::int64_t stride_ipiv, std::int64_t batch_size, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(n, lda, batch_size, scratchpad_size); + + int* ipiv32 = (int*)malloc_device(sizeof(int) * stride_ipiv * batch_size, queue); + int* devInfo = (int*)malloc_device(sizeof(int) * batch_size, queue); + + auto done_casting = + copy_ipiv_to_32(queue, ipiv, ipiv32, n, stride_ipiv, batch_size, dependencies); + + auto done = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.depends_on(done_casting); + 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, n, a_, lda, stride_a, ipiv32, + stride_ipiv, devInfo, batch_size); + }); + }); + + try { + lapack_info_check_batch(queue, devInfo, __func__, func_name, batch_size); + } + catch (...) { + sycl::free(ipiv32, queue); + sycl::free(devInfo, queue); + throw; + } + sycl::free(ipiv32, queue); + sycl::free(devInfo, queue); + return done; } -void ungqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer>& tau, std::int64_t stride_tau, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "ungqr_batch"); + +#define GETRI_STRIDED_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event getri_batch( \ + sycl::queue& queue, std::int64_t n, TYPE* a, std::int64_t lda, std::int64_t stride_a, \ + std::int64_t* ipiv, std::int64_t stride_ipiv, std::int64_t batch_size, TYPE* scratchpad, \ + std::int64_t scratchpad_size, const std::vector& dependencies) { \ + return getri_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, n, a, lda, stride_a, \ + ipiv, stride_ipiv, batch_size, scratchpad, scratchpad_size, \ + dependencies); \ + } + +GETRI_STRIDED_BATCH_LAUNCHER_USM(float, rocsolver_sgetri_strided_batched) +GETRI_STRIDED_BATCH_LAUNCHER_USM(double, rocsolver_dgetri_strided_batched) +GETRI_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_cgetri_strided_batched) +GETRI_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_zgetri_strided_batched) + +#undef GETRI_STRIDED_BATCH_LAUNCHER_USM + +template +inline sycl::event getri_batch(const char* func_name, Func func, sycl::queue& queue, + std::int64_t* n, T** a, std::int64_t* lda, std::int64_t** ipiv, + std::int64_t group_count, std::int64_t* group_sizes, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + + std::int64_t batch_size = 0; + overflow_check(group_count, scratchpad_size); + for (std::int64_t i = 0; i < group_count; ++i) { + overflow_check(n[i], lda[i], group_sizes[i]); + batch_size += group_sizes[i]; + } + + std::int64_t ipiv32_size = 0; + for (std::int64_t i = 0; i < group_count; ++i) + ipiv32_size += n[i] * group_sizes[i]; + + T** a_dev = (T**)malloc_device(sizeof(T*) * batch_size, queue); + int* ipiv32 = (int*)malloc_device(sizeof(int) * ipiv32_size, queue); + int* devInfo = (int*)malloc_device(sizeof(int) * batch_size, queue); + auto done_cpy = + queue.submit([&](sycl::handler& h) { h.memcpy(a_dev, a, batch_size * sizeof(T*)); }); + + std::vector casting_dependencies; + casting_dependencies.reserve(batch_size); + for (std::int64_t i = 0, global_id = 0, ipiv_offset = 0; i < group_count; ++i) { + const std::int64_t len = n[i]; + for (std::int64_t j = 0; j < group_sizes[i]; ++j, ++global_id, ipiv_offset += len) { + const std::int64_t* d_ipiv = ipiv[global_id]; + int* d_ipiv32 = ipiv32 + ipiv_offset; + casting_dependencies.push_back(queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.parallel_for(sycl::range<1>{ static_cast(len) }, + [=](sycl::id<1> index) { + d_ipiv32[index] = static_cast(d_ipiv[index]); + }); + })); + } + } + + auto done = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.depends_on(done_cpy); + cgh.depends_on(casting_dependencies); + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto** a_ = reinterpret_cast(a_dev); + std::int64_t offset = 0; + std::int64_t ipiv_offset = 0; + rocblas_status err; + for (std::int64_t i = 0; i < group_count; ++i) { + rocsolver_native_named_func(func_name, func, err, handle, (int)n[i], a_ + offset, + (int)lda[i], ipiv32 + ipiv_offset, n[i], + devInfo + offset, (int)group_sizes[i]); + offset += group_sizes[i]; + ipiv_offset += n[i] * group_sizes[i]; + } + }); + }); + + try { + lapack_info_check_batch(queue, devInfo, __func__, func_name, batch_size); + } + catch (...) { + sycl::free(a_dev, queue); + sycl::free(ipiv32, queue); + sycl::free(devInfo, queue); + throw; + } + sycl::free(a_dev, queue); + sycl::free(ipiv32, queue); + sycl::free(devInfo, queue); + return done; } -void ungqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, - sycl::buffer>& a, std::int64_t lda, std::int64_t stride_a, - sycl::buffer>& tau, std::int64_t stride_tau, - std::int64_t batch_size, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "ungqr_batch"); + +#define GETRI_GROUP_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event getri_batch( \ + sycl::queue& queue, std::int64_t* n, TYPE** a, std::int64_t* lda, std::int64_t** ipiv, \ + std::int64_t group_count, std::int64_t* group_sizes, TYPE* scratchpad, \ + std::int64_t scratchpad_size, const std::vector& dependencies) { \ + return getri_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, n, a, lda, ipiv, \ + group_count, group_sizes, scratchpad, scratchpad_size, dependencies); \ + } + +GETRI_GROUP_BATCH_LAUNCHER_USM(float, rocsolver_sgetri_batched) +GETRI_GROUP_BATCH_LAUNCHER_USM(double, rocsolver_dgetri_batched) +GETRI_GROUP_BATCH_LAUNCHER_USM(std::complex, rocsolver_cgetri_batched) +GETRI_GROUP_BATCH_LAUNCHER_USM(std::complex, rocsolver_zgetri_batched) + +#undef GETRI_GROUP_BATCH_LAUNCHER_USM + +template +inline sycl::event getrs_batch(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 stride_a, std::int64_t* ipiv, + std::int64_t stride_ipiv, T* b, std::int64_t ldb, + std::int64_t stride_b, std::int64_t batch_size, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(n, nrhs, lda, ldb, batch_size, scratchpad_size); + + int* ipiv32 = (int*)malloc_device(sizeof(int) * stride_ipiv * batch_size, queue); + + auto done_casting = + copy_ipiv_to_32(queue, ipiv, ipiv32, n, stride_ipiv, batch_size, dependencies); + + auto done = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.depends_on(done_casting); + 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, stride_a, ipiv32, stride_ipiv, b_, ldb, + stride_b, batch_size); + }); + }); + + queue.wait(); + sycl::free(ipiv32, queue); + return done; } -// BATCH USM API +#define GETRS_STRIDED_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, \ + std::int64_t nrhs, TYPE* a, std::int64_t lda, std::int64_t stride_a, \ + std::int64_t* ipiv, std::int64_t stride_ipiv, TYPE* b, \ + std::int64_t ldb, std::int64_t stride_b, std::int64_t batch_size, \ + TYPE* scratchpad, std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return getrs_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, trans, n, nrhs, a, lda, \ + stride_a, ipiv, stride_ipiv, b, ldb, stride_b, batch_size, scratchpad, \ + scratchpad_size, dependencies); \ + } + +GETRS_STRIDED_BATCH_LAUNCHER_USM(float, rocsolver_sgetrs_strided_batched) +GETRS_STRIDED_BATCH_LAUNCHER_USM(double, rocsolver_dgetrs_strided_batched) +GETRS_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_cgetrs_strided_batched) +GETRS_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_zgetrs_strided_batched) + +#undef GETRS_STRIDED_BATCH_LAUNCHER_USM + +template +inline sycl::event getrs_batch(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, std::int64_t group_count, + std::int64_t* group_sizes, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + + std::int64_t batch_size = 0; + overflow_check(group_count, scratchpad_size); + for (std::int64_t i = 0; i < group_count; ++i) { + overflow_check(n[i], nrhs[i], lda[i], ldb[i], group_sizes[i]); + batch_size += group_sizes[i]; + } + + std::int64_t ipiv32_size = 0; + for (std::int64_t i = 0; i < group_count; ++i) + ipiv32_size += n[i] * group_sizes[i]; + + T** a_dev = (T**)malloc_device(sizeof(T*) * batch_size, queue); + T** b_dev = (T**)malloc_device(sizeof(T*) * batch_size, queue); + int* ipiv32 = (int*)malloc_device(sizeof(int) * ipiv32_size, queue); + auto done_cpy_a = + queue.submit([&](sycl::handler& h) { h.memcpy(a_dev, a, batch_size * sizeof(T*)); }); + auto done_cpy_b = + queue.submit([&](sycl::handler& h) { h.memcpy(b_dev, b, batch_size * sizeof(T*)); }); + + std::vector casting_dependencies; + casting_dependencies.reserve(batch_size); + for (std::int64_t i = 0, global_id = 0, ipiv_offset = 0; i < group_count; ++i) { + const std::int64_t len = n[i]; + for (std::int64_t j = 0; j < group_sizes[i]; ++j, ++global_id, ipiv_offset += len) { + const std::int64_t* d_ipiv = ipiv[global_id]; + int* d_ipiv32 = ipiv32 + ipiv_offset; + casting_dependencies.push_back(queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.parallel_for(sycl::range<1>{ static_cast(len) }, + [=](sycl::id<1> index) { + d_ipiv32[index] = static_cast(d_ipiv[index]); + }); + })); + } + } -sycl::event geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, float* a, - std::int64_t lda, std::int64_t stride_a, float* tau, - std::int64_t stride_tau, std::int64_t batch_size, float* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "geqrf_batch"); -} -sycl::event geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, double* a, - std::int64_t lda, std::int64_t stride_a, double* tau, - std::int64_t stride_tau, std::int64_t batch_size, double* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "geqrf_batch"); -} -sycl::event geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::complex* a, - std::int64_t lda, std::int64_t stride_a, std::complex* tau, - std::int64_t stride_tau, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "geqrf_batch"); -} -sycl::event geqrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::complex* a, - std::int64_t lda, std::int64_t stride_a, std::complex* tau, - std::int64_t stride_tau, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "geqrf_batch"); -} -sycl::event geqrf_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, float** a, - std::int64_t* lda, float** tau, std::int64_t group_count, - std::int64_t* group_sizes, float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "geqrf_batch"); -} -sycl::event geqrf_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, double** a, - std::int64_t* lda, double** tau, std::int64_t group_count, - std::int64_t* group_sizes, double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "geqrf_batch"); -} -sycl::event geqrf_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, - std::complex** a, std::int64_t* lda, std::complex** tau, - std::int64_t group_count, std::int64_t* group_sizes, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "geqrf_batch"); -} -sycl::event geqrf_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, - std::complex** a, std::int64_t* lda, std::complex** tau, - std::int64_t group_count, std::int64_t* group_sizes, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "geqrf_batch"); -} -sycl::event getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, float* a, - std::int64_t lda, std::int64_t stride_a, std::int64_t* ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, float* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrf_batch"); -} -sycl::event getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, double* a, - std::int64_t lda, std::int64_t stride_a, std::int64_t* ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, double* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrf_batch"); -} -sycl::event getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::complex* a, - std::int64_t lda, std::int64_t stride_a, std::int64_t* ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrf_batch"); -} -sycl::event getrf_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::complex* a, - std::int64_t lda, std::int64_t stride_a, std::int64_t* ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrf_batch"); -} -sycl::event getrf_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, float** a, - std::int64_t* lda, std::int64_t** ipiv, std::int64_t group_count, - std::int64_t* group_sizes, float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrf_batch"); -} -sycl::event getrf_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, double** a, - std::int64_t* lda, std::int64_t** ipiv, std::int64_t group_count, - std::int64_t* group_sizes, double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrf_batch"); -} -sycl::event getrf_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, - std::complex** a, std::int64_t* lda, std::int64_t** ipiv, - std::int64_t group_count, std::int64_t* group_sizes, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrf_batch"); -} -sycl::event getrf_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, - std::complex** a, std::int64_t* lda, std::int64_t** ipiv, - std::int64_t group_count, std::int64_t* group_sizes, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrf_batch"); -} -sycl::event getri_batch(sycl::queue& queue, std::int64_t n, float* a, std::int64_t lda, - std::int64_t stride_a, std::int64_t* ipiv, std::int64_t stride_ipiv, - std::int64_t batch_size, float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri_batch"); -} -sycl::event getri_batch(sycl::queue& queue, std::int64_t n, double* a, std::int64_t lda, - std::int64_t stride_a, std::int64_t* ipiv, std::int64_t stride_ipiv, - std::int64_t batch_size, double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri_batch"); -} -sycl::event getri_batch(sycl::queue& queue, std::int64_t n, std::complex* a, - std::int64_t lda, std::int64_t stride_a, std::int64_t* ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri_batch"); -} -sycl::event getri_batch(sycl::queue& queue, std::int64_t n, std::complex* a, - std::int64_t lda, std::int64_t stride_a, std::int64_t* ipiv, - std::int64_t stride_ipiv, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri_batch"); -} -sycl::event getri_batch(sycl::queue& queue, std::int64_t* n, float** a, std::int64_t* lda, - std::int64_t** ipiv, std::int64_t group_count, std::int64_t* group_sizes, - float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri_batch"); -} -sycl::event getri_batch(sycl::queue& queue, std::int64_t* n, double** a, std::int64_t* lda, - std::int64_t** ipiv, std::int64_t group_count, std::int64_t* group_sizes, - double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri_batch"); -} -sycl::event getri_batch(sycl::queue& queue, std::int64_t* n, std::complex** a, - std::int64_t* lda, std::int64_t** ipiv, std::int64_t group_count, - std::int64_t* group_sizes, std::complex* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri_batch"); -} -sycl::event getri_batch(sycl::queue& queue, std::int64_t* n, std::complex** a, - std::int64_t* lda, std::int64_t** ipiv, std::int64_t group_count, - std::int64_t* group_sizes, std::complex* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri_batch"); -} -sycl::event getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, - std::int64_t nrhs, float* a, std::int64_t lda, std::int64_t stride_a, - std::int64_t* ipiv, std::int64_t stride_ipiv, float* b, std::int64_t ldb, - std::int64_t stride_b, std::int64_t batch_size, float* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrs_batch"); -} -sycl::event getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, - std::int64_t nrhs, double* a, std::int64_t lda, std::int64_t stride_a, - std::int64_t* ipiv, std::int64_t stride_ipiv, double* b, std::int64_t ldb, - std::int64_t stride_b, std::int64_t batch_size, double* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrs_batch"); -} -sycl::event getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, - std::int64_t nrhs, std::complex* a, std::int64_t lda, - std::int64_t stride_a, std::int64_t* ipiv, std::int64_t stride_ipiv, - std::complex* b, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size, std::complex* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrs_batch"); -} -sycl::event getrs_batch(sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, - std::int64_t nrhs, std::complex* a, std::int64_t lda, - std::int64_t stride_a, std::int64_t* ipiv, std::int64_t stride_ipiv, - std::complex* b, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size, std::complex* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrs_batch"); -} -sycl::event getrs_batch(sycl::queue& queue, oneapi::math::transpose* trans, std::int64_t* n, - std::int64_t* nrhs, float** a, std::int64_t* lda, std::int64_t** ipiv, - float** b, std::int64_t* ldb, std::int64_t group_count, - std::int64_t* group_sizes, float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrs_batch"); -} -sycl::event getrs_batch(sycl::queue& queue, oneapi::math::transpose* trans, std::int64_t* n, - std::int64_t* nrhs, double** a, std::int64_t* lda, std::int64_t** ipiv, - double** b, std::int64_t* ldb, std::int64_t group_count, - std::int64_t* group_sizes, double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrs_batch"); -} -sycl::event getrs_batch(sycl::queue& queue, oneapi::math::transpose* trans, std::int64_t* n, - std::int64_t* nrhs, std::complex** a, std::int64_t* lda, - std::int64_t** ipiv, std::complex** b, std::int64_t* ldb, - std::int64_t group_count, std::int64_t* group_sizes, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrs_batch"); -} -sycl::event getrs_batch(sycl::queue& queue, oneapi::math::transpose* trans, std::int64_t* n, - std::int64_t* nrhs, std::complex** a, std::int64_t* lda, - std::int64_t** ipiv, std::complex** b, std::int64_t* ldb, - std::int64_t group_count, std::int64_t* group_sizes, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getrs_batch"); -} -sycl::event orgqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, - float* a, std::int64_t lda, std::int64_t stride_a, float* tau, - std::int64_t stride_tau, std::int64_t batch_size, float* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "orgqr_batch"); -} -sycl::event orgqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, - double* a, std::int64_t lda, std::int64_t stride_a, double* tau, - std::int64_t stride_tau, std::int64_t batch_size, double* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "orgqr_batch"); -} -sycl::event orgqr_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, std::int64_t* k, - float** a, std::int64_t* lda, float** tau, std::int64_t group_count, - std::int64_t* group_sizes, float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "orgqr_batch"); -} -sycl::event orgqr_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, std::int64_t* k, - double** a, std::int64_t* lda, double** tau, std::int64_t group_count, - std::int64_t* group_sizes, double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "orgqr_batch"); -} -sycl::event potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, float* a, - std::int64_t lda, std::int64_t stride_a, std::int64_t batch_size, - float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "potrf_batch"); + auto done = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.depends_on(done_cpy_a); + cgh.depends_on(done_cpy_b); + cgh.depends_on(casting_dependencies); + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto** a_ = reinterpret_cast(a_dev); + auto** b_ = reinterpret_cast(b_dev); + std::int64_t offset = 0; + std::int64_t ipiv_offset = 0; + rocblas_status err; + for (std::int64_t i = 0; i < group_count; ++i) { + rocsolver_native_named_func( + func_name, func, err, handle, get_rocblas_operation(trans[i]), (int)n[i], + (int)nrhs[i], a_ + offset, (int)lda[i], ipiv32 + ipiv_offset, n[i], b_ + offset, + (int)ldb[i], (int)group_sizes[i]); + offset += group_sizes[i]; + ipiv_offset += n[i] * group_sizes[i]; + } + }); + }); + + queue.wait(); + sycl::free(a_dev, queue); + sycl::free(b_dev, queue); + sycl::free(ipiv32, queue); + return done; } -sycl::event potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, double* a, - std::int64_t lda, std::int64_t stride_a, std::int64_t batch_size, - double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "potrf_batch"); + +#define GETRS_GROUP_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event getrs_batch( \ + 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, \ + std::int64_t group_count, std::int64_t* group_sizes, TYPE* scratchpad, \ + std::int64_t scratchpad_size, const std::vector& dependencies) { \ + return getrs_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, trans, n, nrhs, a, lda, \ + ipiv, b, ldb, group_count, group_sizes, scratchpad, scratchpad_size, \ + dependencies); \ + } + +GETRS_GROUP_BATCH_LAUNCHER_USM(float, rocsolver_sgetrs_batched) +GETRS_GROUP_BATCH_LAUNCHER_USM(double, rocsolver_dgetrs_batched) +GETRS_GROUP_BATCH_LAUNCHER_USM(std::complex, rocsolver_cgetrs_batched) +GETRS_GROUP_BATCH_LAUNCHER_USM(std::complex, rocsolver_zgetrs_batched) + +#undef GETRS_GROUP_BATCH_LAUNCHER_USM + +template +inline sycl::event orgqr_batch(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, + std::int64_t n, std::int64_t k, T* a, std::int64_t lda, + std::int64_t stride_a, T* tau, std::int64_t stride_tau, + std::int64_t batch_size, T* scratchpad, std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(m, n, k, lda, batch_size, scratchpad_size); + 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); + auto tau_ = reinterpret_cast(tau); + rocblas_status err; + for (std::int64_t i = 0; i < batch_size; ++i) { + rocsolver_native_named_func(func_name, func, err, handle, m, n, k, + a_ + i * stride_a, lda, tau_ + i * stride_tau); + } + }); + }); + return done; } -sycl::event potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - std::complex* a, std::int64_t lda, std::int64_t stride_a, - std::int64_t batch_size, std::complex* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "potrf_batch"); + +template +inline sycl::event orgqr_batch(const char* func_name, Func func, sycl::queue& queue, + std::int64_t* m, std::int64_t* n, std::int64_t* k, T** a, + std::int64_t* lda, T** tau, std::int64_t group_count, + std::int64_t* group_sizes, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + + overflow_check(group_count, scratchpad_size); + for (std::int64_t i = 0; i < group_count; ++i) + overflow_check(m[i], n[i], k[i], lda[i], group_sizes[i]); + + 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); + std::int64_t global_id = 0; + rocblas_status err; + for (std::int64_t i = 0; i < group_count; ++i) { + for (std::int64_t j = 0; j < group_sizes[i]; ++j, ++global_id) { + auto a_ = reinterpret_cast(a[global_id]); + auto tau_ = reinterpret_cast(tau[global_id]); + rocsolver_native_named_func(func_name, func, err, handle, (int)m[i], (int)n[i], + (int)k[i], a_, (int)lda[i], tau_); + } + } + }); + }); + return done; } -sycl::event potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - std::complex* a, std::int64_t lda, std::int64_t stride_a, - std::int64_t batch_size, std::complex* scratchpad, - std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "potrf_batch"); + +#define ORGQR_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event orgqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, \ + TYPE* a, std::int64_t lda, std::int64_t stride_a, TYPE* tau, \ + std::int64_t stride_tau, std::int64_t batch_size, TYPE* scratchpad, \ + std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return orgqr_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, k, a, lda, \ + stride_a, tau, stride_tau, batch_size, scratchpad, scratchpad_size, \ + dependencies); \ + } \ + sycl::event orgqr_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, std::int64_t* k, \ + TYPE** a, std::int64_t* lda, TYPE** tau, std::int64_t group_count, \ + std::int64_t* group_sizes, TYPE* scratchpad, \ + std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return orgqr_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, k, a, lda, tau, \ + group_count, group_sizes, scratchpad, scratchpad_size, dependencies); \ + } + +ORGQR_BATCH_LAUNCHER_USM(float, rocsolver_sorgqr) +ORGQR_BATCH_LAUNCHER_USM(double, rocsolver_dorgqr) + +#undef ORGQR_BATCH_LAUNCHER_USM + +#define UNGQR_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event ungqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, \ + TYPE* a, std::int64_t lda, std::int64_t stride_a, TYPE* tau, \ + std::int64_t stride_tau, std::int64_t batch_size, TYPE* scratchpad, \ + std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return orgqr_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, k, a, lda, \ + stride_a, tau, stride_tau, batch_size, scratchpad, scratchpad_size, \ + dependencies); \ + } \ + sycl::event ungqr_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, std::int64_t* k, \ + TYPE** a, std::int64_t* lda, TYPE** tau, std::int64_t group_count, \ + std::int64_t* group_sizes, TYPE* scratchpad, \ + std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return orgqr_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, k, a, lda, tau, \ + group_count, group_sizes, scratchpad, scratchpad_size, dependencies); \ + } + +UNGQR_BATCH_LAUNCHER_USM(std::complex, rocsolver_cungqr) +UNGQR_BATCH_LAUNCHER_USM(std::complex, rocsolver_zungqr) + +#undef UNGQR_BATCH_LAUNCHER_USM + +template +inline sycl::event potrf_batch(const char* func_name, Func func, sycl::queue& queue, + oneapi::math::uplo uplo, std::int64_t n, T* a, std::int64_t lda, + std::int64_t stride_a, std::int64_t batch_size, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(n, lda, batch_size, scratchpad_size); + + int* devInfo = (int*)malloc_device(sizeof(int) * batch_size, 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, get_rocblas_fill_mode(uplo), + n, a_, lda, stride_a, devInfo, batch_size); + }); + }); + + try { + lapack_info_check_batch(queue, devInfo, __func__, func_name, batch_size); + } + catch (...) { + sycl::free(devInfo, queue); + throw; + } + sycl::free(devInfo, queue); + return done; } +#define POTRF_STRIDED_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event potrf_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, TYPE* a, \ + std::int64_t lda, std::int64_t stride_a, std::int64_t batch_size, \ + TYPE* scratchpad, std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return potrf_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, uplo, n, a, lda, \ + stride_a, batch_size, scratchpad, scratchpad_size, dependencies); \ + } + +POTRF_STRIDED_BATCH_LAUNCHER_USM(float, rocsolver_spotrf_strided_batched) +POTRF_STRIDED_BATCH_LAUNCHER_USM(double, rocsolver_dpotrf_strided_batched) +POTRF_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_cpotrf_strided_batched) +POTRF_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_zpotrf_strided_batched) + +#undef POTRF_STRIDED_BATCH_LAUNCHER_USM + template inline sycl::event potrf_batch(const char* func_name, Func func, sycl::queue& queue, oneapi::math::uplo* uplo, std::int64_t* n, T** a, std::int64_t* lda, @@ -535,6 +1213,17 @@ inline sycl::event potrf_batch(const char* func_name, Func func, sycl::queue& qu } }); }); + + try { + lapack_info_check_batch(queue, info, __func__, func_name, batch_size); + } + catch (...) { + sycl::free(a_dev, queue); + sycl::free(info, queue); + throw; + } + sycl::free(a_dev, queue); + sycl::free(info, queue); return done; } @@ -555,37 +1244,48 @@ POTRF_BATCH_LAUNCHER_USM(std::complex, rocsolver_zpotrf_batched) #undef POTRF_BATCH_LAUNCHER_USM -sycl::event potrs_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - std::int64_t nrhs, float* a, std::int64_t lda, std::int64_t stride_a, - float* b, std::int64_t ldb, std::int64_t stride_b, std::int64_t batch_size, - float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "potrs_batch"); -} -sycl::event potrs_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - std::int64_t nrhs, double* a, std::int64_t lda, std::int64_t stride_a, - double* b, std::int64_t ldb, std::int64_t stride_b, std::int64_t batch_size, - double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "potrs_batch"); -} -sycl::event potrs_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - std::int64_t nrhs, std::complex* a, std::int64_t lda, - std::int64_t stride_a, std::complex* b, std::int64_t ldb, - std::int64_t stride_b, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "potrs_batch"); -} -sycl::event potrs_batch(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, - std::int64_t nrhs, std::complex* a, std::int64_t lda, - std::int64_t stride_a, std::complex* b, std::int64_t ldb, - std::int64_t stride_b, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "potrs_batch"); +template +inline sycl::event potrs_batch(const char* func_name, Func func, sycl::queue& queue, + oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, T* a, + std::int64_t lda, std::int64_t stride_a, T* b, std::int64_t ldb, + std::int64_t stride_b, std::int64_t batch_size, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(n, nrhs, lda, ldb, batch_size, scratchpad_size); + + 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); + auto b_ = reinterpret_cast(b); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, get_rocblas_fill_mode(uplo), + n, nrhs, a_, lda, stride_a, b_, ldb, stride_b, batch_size); + }); + }); + return done; } +#define POTRS_STRIDED_BATCH_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event potrs_batch( \ + sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, TYPE* a, \ + std::int64_t lda, std::int64_t stride_a, TYPE* b, std::int64_t ldb, std::int64_t stride_b, \ + std::int64_t batch_size, TYPE* scratchpad, std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return potrs_batch(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, uplo, n, nrhs, a, lda, \ + stride_a, b, ldb, stride_b, batch_size, scratchpad, scratchpad_size, \ + dependencies); \ + } + +POTRS_STRIDED_BATCH_LAUNCHER_USM(float, rocsolver_spotrs_strided_batched) +POTRS_STRIDED_BATCH_LAUNCHER_USM(double, rocsolver_dpotrs_strided_batched) +POTRS_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_cpotrs_strided_batched) +POTRS_STRIDED_BATCH_LAUNCHER_USM(std::complex, rocsolver_zpotrs_strided_batched) + +#undef POTRS_STRIDED_BATCH_LAUNCHER_USM + template inline sycl::event potrs_batch(const char* func_name, Func func, sycl::queue& queue, oneapi::math::uplo* uplo, std::int64_t* n, std::int64_t* nrhs, T** a, @@ -597,13 +1297,8 @@ inline sycl::event potrs_batch(const char* func_name, Func func, sycl::queue& qu int64_t batch_size = 0; for (int64_t i = 0; i < group_count; i++) { - overflow_check(n[i], lda[i], group_sizes[i]); + overflow_check(n[i], nrhs[i], lda[i], ldb[i], group_sizes[i]); batch_size += group_sizes[i]; - - // rocsolver function only supports nrhs = 1 - if (nrhs[i] != 1) - throw unimplemented("lapack", "potrs_batch", - "rocsolver potrs_batch only supports nrhs = 1"); } T** a_dev = (T**)malloc_device(sizeof(T*) * batch_size, queue); @@ -636,6 +1331,10 @@ inline sycl::event potrs_batch(const char* func_name, Func func, sycl::queue& qu } }); }); + + queue.wait(); + sycl::free(a_dev, queue); + sycl::free(b_dev, queue); return done; } @@ -658,361 +1357,215 @@ POTRS_BATCH_LAUNCHER_USM(std::complex, rocsolver_zpotrs_batched) #undef POTRS_BATCH_LAUNCHER_USM -sycl::event ungqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, - std::complex* a, std::int64_t lda, std::int64_t stride_a, - std::complex* tau, std::int64_t stride_tau, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "ungqr_batch"); -} -sycl::event ungqr_batch(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, - std::complex* a, std::int64_t lda, std::int64_t stride_a, - std::complex* tau, std::int64_t stride_tau, std::int64_t batch_size, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "ungqr_batch"); -} -sycl::event ungqr_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, std::int64_t* k, - std::complex** a, std::int64_t* lda, std::complex** tau, - std::int64_t group_count, std::int64_t* group_sizes, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "ungqr_batch"); -} -sycl::event ungqr_batch(sycl::queue& queue, std::int64_t* m, std::int64_t* n, std::int64_t* k, - std::complex** a, std::int64_t* lda, std::complex** tau, - std::int64_t group_count, std::int64_t* group_sizes, - std::complex* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "ungqr_batch"); -} - // BATCH SCRATCHPAD API +// +// rocsolver allocates any workspace it needs from the rocblas handle, so every +// scratchpad query below reports that no user provided memory is required. -template <> -std::int64_t getrf_batch_scratchpad_size(sycl::queue& queue, std::int64_t m, std::int64_t n, - std::int64_t lda, std::int64_t stride_a, - std::int64_t stride_ipiv, std::int64_t batch_size) { - throw unimplemented("lapack", "getrf_batch_scratchpad_size"); -} -template <> -std::int64_t getrf_batch_scratchpad_size(sycl::queue& queue, std::int64_t m, std::int64_t n, - std::int64_t lda, std::int64_t stride_a, - std::int64_t stride_ipiv, - std::int64_t batch_size) { - throw unimplemented("lapack", "getrf_batch_scratchpad_size"); -} -template <> -std::int64_t getrf_batch_scratchpad_size>(sycl::queue& queue, std::int64_t m, - std::int64_t n, std::int64_t lda, - std::int64_t stride_a, - std::int64_t stride_ipiv, - std::int64_t batch_size) { - throw unimplemented("lapack", "getrf_batch_scratchpad_size"); -} -template <> -std::int64_t getrf_batch_scratchpad_size>(sycl::queue& queue, std::int64_t m, - std::int64_t n, std::int64_t lda, - std::int64_t stride_a, - std::int64_t stride_ipiv, - std::int64_t batch_size) { - throw unimplemented("lapack", "getrf_batch_scratchpad_size"); -} -template <> -std::int64_t getri_batch_scratchpad_size(sycl::queue& queue, std::int64_t n, - std::int64_t lda, std::int64_t stride_a, - std::int64_t stride_ipiv, std::int64_t batch_size) { - throw unimplemented("lapack", "getri_batch_scratchpad_size"); -} -template <> -std::int64_t getri_batch_scratchpad_size(sycl::queue& queue, std::int64_t n, - std::int64_t lda, std::int64_t stride_a, - std::int64_t stride_ipiv, - std::int64_t batch_size) { - throw unimplemented("lapack", "getri_batch_scratchpad_size"); -} -template <> -std::int64_t getri_batch_scratchpad_size>(sycl::queue& queue, std::int64_t n, - std::int64_t lda, - std::int64_t stride_a, - std::int64_t stride_ipiv, - std::int64_t batch_size) { - throw unimplemented("lapack", "getri_batch_scratchpad_size"); -} -template <> -std::int64_t getri_batch_scratchpad_size>(sycl::queue& queue, std::int64_t n, - std::int64_t lda, - std::int64_t stride_a, - std::int64_t stride_ipiv, - std::int64_t batch_size) { - throw unimplemented("lapack", "getri_batch_scratchpad_size"); -} -template <> -std::int64_t getrs_batch_scratchpad_size(sycl::queue& queue, oneapi::math::transpose trans, - std::int64_t n, std::int64_t nrhs, std::int64_t lda, - std::int64_t stride_a, std::int64_t stride_ipiv, - std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size) { - throw unimplemented("lapack", "getrs_batch_scratchpad_size"); -} -template <> -std::int64_t getrs_batch_scratchpad_size(sycl::queue& queue, oneapi::math::transpose trans, - std::int64_t n, std::int64_t nrhs, - std::int64_t lda, std::int64_t stride_a, - std::int64_t stride_ipiv, std::int64_t ldb, - std::int64_t stride_b, std::int64_t batch_size) { - throw unimplemented("lapack", "getrs_batch_scratchpad_size"); -} -template <> -std::int64_t getrs_batch_scratchpad_size>( - sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, std::int64_t nrhs, - std::int64_t lda, std::int64_t stride_a, std::int64_t stride_ipiv, std::int64_t ldb, - std::int64_t stride_b, std::int64_t batch_size) { - throw unimplemented("lapack", "getrs_batch_scratchpad_size"); -} -template <> -std::int64_t getrs_batch_scratchpad_size>( - sycl::queue& queue, oneapi::math::transpose trans, std::int64_t n, std::int64_t nrhs, - std::int64_t lda, std::int64_t stride_a, std::int64_t stride_ipiv, std::int64_t ldb, - std::int64_t stride_b, std::int64_t batch_size) { - throw unimplemented("lapack", "getrs_batch_scratchpad_size"); -} -template <> -std::int64_t geqrf_batch_scratchpad_size(sycl::queue& queue, std::int64_t m, std::int64_t n, - std::int64_t lda, std::int64_t stride_a, - std::int64_t stride_tau, std::int64_t batch_size) { - throw unimplemented("lapack", "geqrf_batch_scratchpad_size"); -} -template <> -std::int64_t geqrf_batch_scratchpad_size(sycl::queue& queue, std::int64_t m, std::int64_t n, - std::int64_t lda, std::int64_t stride_a, - std::int64_t stride_tau, std::int64_t batch_size) { - throw unimplemented("lapack", "geqrf_batch_scratchpad_size"); -} -template <> -std::int64_t geqrf_batch_scratchpad_size>(sycl::queue& queue, std::int64_t m, - std::int64_t n, std::int64_t lda, - std::int64_t stride_a, - std::int64_t stride_tau, - std::int64_t batch_size) { - throw unimplemented("lapack", "geqrf_batch_scratchpad_size"); -} -template <> -std::int64_t geqrf_batch_scratchpad_size>(sycl::queue& queue, std::int64_t m, - std::int64_t n, std::int64_t lda, - std::int64_t stride_a, - std::int64_t stride_tau, - std::int64_t batch_size) { - throw unimplemented("lapack", "geqrf_batch_scratchpad_size"); -} +#define GEQRF_STRIDED_BATCH_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t geqrf_batch_scratchpad_size( \ + sycl::queue & queue, std::int64_t m, std::int64_t n, std::int64_t lda, \ + std::int64_t stride_a, std::int64_t stride_tau, std::int64_t batch_size) { \ + return 0; \ + } -template <> -std::int64_t potrf_batch_scratchpad_size(sycl::queue& queue, oneapi::math::uplo uplo, - std::int64_t n, std::int64_t lda, - std::int64_t stride_a, std::int64_t batch_size) { - throw unimplemented("lapack", "potrf_batch_scratchpad_size"); -} -template <> -std::int64_t potrf_batch_scratchpad_size(sycl::queue& queue, oneapi::math::uplo uplo, - std::int64_t n, std::int64_t lda, - std::int64_t stride_a, std::int64_t batch_size) { - throw unimplemented("lapack", "potrf_batch_scratchpad_size"); -} -template <> -std::int64_t potrf_batch_scratchpad_size>(sycl::queue& queue, - oneapi::math::uplo uplo, - std::int64_t n, std::int64_t lda, - std::int64_t stride_a, - std::int64_t batch_size) { - throw unimplemented("lapack", "potrf_batch_scratchpad_size"); -} -template <> -std::int64_t potrf_batch_scratchpad_size>(sycl::queue& queue, - oneapi::math::uplo uplo, - std::int64_t n, std::int64_t lda, - std::int64_t stride_a, - std::int64_t batch_size) { - throw unimplemented("lapack", "potrf_batch_scratchpad_size"); -} -template <> -std::int64_t potrs_batch_scratchpad_size(sycl::queue& queue, oneapi::math::uplo uplo, - std::int64_t n, std::int64_t nrhs, std::int64_t lda, - std::int64_t stride_a, std::int64_t ldb, - std::int64_t stride_b, std::int64_t batch_size) { - throw unimplemented("lapack", "potrs_batch_scratchpad_size"); -} -template <> -std::int64_t potrs_batch_scratchpad_size(sycl::queue& queue, oneapi::math::uplo uplo, - std::int64_t n, std::int64_t nrhs, - std::int64_t lda, std::int64_t stride_a, - std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size) { - throw unimplemented("lapack", "potrs_batch_scratchpad_size"); -} -template <> -std::int64_t potrs_batch_scratchpad_size>( - sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, - std::int64_t lda, std::int64_t stride_a, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size) { - throw unimplemented("lapack", "potrs_batch_scratchpad_size"); -} -template <> -std::int64_t potrs_batch_scratchpad_size>( - sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, - std::int64_t lda, std::int64_t stride_a, std::int64_t ldb, std::int64_t stride_b, - std::int64_t batch_size) { - throw unimplemented("lapack", "potrs_batch_scratchpad_size"); -} -template <> -std::int64_t orgqr_batch_scratchpad_size(sycl::queue& queue, std::int64_t m, std::int64_t n, - std::int64_t k, std::int64_t lda, - std::int64_t stride_a, std::int64_t stride_tau, - std::int64_t batch_size) { - throw unimplemented("lapack", "orgqr_batch_scratchpad_size"); -} -template <> -std::int64_t orgqr_batch_scratchpad_size(sycl::queue& queue, std::int64_t m, std::int64_t n, - std::int64_t k, std::int64_t lda, - std::int64_t stride_a, std::int64_t stride_tau, - std::int64_t batch_size) { - throw unimplemented("lapack", "orgqr_batch_scratchpad_size"); -} -template <> -std::int64_t ungqr_batch_scratchpad_size>( - sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, std::int64_t lda, - std::int64_t stride_a, std::int64_t stride_tau, std::int64_t batch_size) { - throw unimplemented("lapack", "ungqr_batch_scratchpad_size"); -} -template <> -std::int64_t ungqr_batch_scratchpad_size>( - sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t k, std::int64_t lda, - std::int64_t stride_a, std::int64_t stride_tau, std::int64_t batch_size) { - throw unimplemented("lapack", "ungqr_batch_scratchpad_size"); -} -template <> -std::int64_t getrf_batch_scratchpad_size(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getrf_batch_scratchpad_size"); -} -template <> -std::int64_t getrf_batch_scratchpad_size(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getrf_batch_scratchpad_size"); -} -template <> -std::int64_t getrf_batch_scratchpad_size>(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getrf_batch_scratchpad_size"); -} -template <> -std::int64_t getrf_batch_scratchpad_size>(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getrf_batch_scratchpad_size"); -} -template <> -std::int64_t getri_batch_scratchpad_size(sycl::queue& queue, std::int64_t* n, - std::int64_t* lda, std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getri_batch_scratchpad_size"); -} -template <> -std::int64_t getri_batch_scratchpad_size(sycl::queue& queue, std::int64_t* n, - std::int64_t* lda, std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getri_batch_scratchpad_size"); -} -template <> -std::int64_t getri_batch_scratchpad_size>(sycl::queue& queue, std::int64_t* n, - std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getri_batch_scratchpad_size"); -} -template <> -std::int64_t getri_batch_scratchpad_size>(sycl::queue& queue, std::int64_t* n, - std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getri_batch_scratchpad_size"); -} -template <> -std::int64_t getrs_batch_scratchpad_size(sycl::queue& queue, oneapi::math::transpose* trans, - std::int64_t* n, std::int64_t* nrhs, - std::int64_t* lda, std::int64_t* ldb, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getrs_batch_scratchpad_size"); -} -template <> -std::int64_t getrs_batch_scratchpad_size(sycl::queue& queue, oneapi::math::transpose* trans, - std::int64_t* n, std::int64_t* nrhs, - std::int64_t* lda, std::int64_t* ldb, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "getrs_batch_scratchpad_size"); -} -template <> -std::int64_t getrs_batch_scratchpad_size>( - sycl::queue& queue, oneapi::math::transpose* trans, std::int64_t* n, std::int64_t* nrhs, - std::int64_t* lda, std::int64_t* ldb, std::int64_t group_count, std::int64_t* group_sizes) { - throw unimplemented("lapack", "getrs_batch_scratchpad_size"); -} -template <> -std::int64_t getrs_batch_scratchpad_size>( - sycl::queue& queue, oneapi::math::transpose* trans, std::int64_t* n, std::int64_t* nrhs, - std::int64_t* lda, std::int64_t* ldb, std::int64_t group_count, std::int64_t* group_sizes) { - throw unimplemented("lapack", "getrs_batch_scratchpad_size"); -} -template <> -std::int64_t geqrf_batch_scratchpad_size(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "geqrf_batch_scratchpad_size"); -} -template <> -std::int64_t geqrf_batch_scratchpad_size(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "geqrf_batch_scratchpad_size"); -} -template <> -std::int64_t geqrf_batch_scratchpad_size>(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "geqrf_batch_scratchpad_size"); -} -template <> -std::int64_t geqrf_batch_scratchpad_size>(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "geqrf_batch_scratchpad_size"); -} -template <> -std::int64_t orgqr_batch_scratchpad_size(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* k, std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "orgqr_batch_scratchpad_size"); -} -template <> -std::int64_t orgqr_batch_scratchpad_size(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* k, - std::int64_t* lda, std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "orgqr_batch_scratchpad_size"); -} +GEQRF_STRIDED_BATCH_LAUNCHER_SCRATCH(float) +GEQRF_STRIDED_BATCH_LAUNCHER_SCRATCH(double) +GEQRF_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) +GEQRF_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) + +#undef GEQRF_STRIDED_BATCH_LAUNCHER_SCRATCH + +#define GETRF_STRIDED_BATCH_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t getrf_batch_scratchpad_size( \ + sycl::queue & queue, std::int64_t m, std::int64_t n, std::int64_t lda, \ + std::int64_t stride_a, std::int64_t stride_ipiv, std::int64_t batch_size) { \ + return 0; \ + } + +GETRF_STRIDED_BATCH_LAUNCHER_SCRATCH(float) +GETRF_STRIDED_BATCH_LAUNCHER_SCRATCH(double) +GETRF_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) +GETRF_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) + +#undef GETRF_STRIDED_BATCH_LAUNCHER_SCRATCH + +#define GETRI_STRIDED_BATCH_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t getri_batch_scratchpad_size( \ + sycl::queue & queue, std::int64_t n, std::int64_t lda, std::int64_t stride_a, \ + std::int64_t stride_ipiv, std::int64_t batch_size) { \ + return 0; \ + } + +GETRI_STRIDED_BATCH_LAUNCHER_SCRATCH(float) +GETRI_STRIDED_BATCH_LAUNCHER_SCRATCH(double) +GETRI_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) +GETRI_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) + +#undef GETRI_STRIDED_BATCH_LAUNCHER_SCRATCH + +#define GETRS_STRIDED_BATCH_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t getrs_batch_scratchpad_size( \ + sycl::queue & queue, oneapi::math::transpose trans, std::int64_t n, std::int64_t nrhs, \ + std::int64_t lda, std::int64_t stride_a, std::int64_t stride_ipiv, std::int64_t ldb, \ + std::int64_t stride_b, std::int64_t batch_size) { \ + return 0; \ + } + +GETRS_STRIDED_BATCH_LAUNCHER_SCRATCH(float) +GETRS_STRIDED_BATCH_LAUNCHER_SCRATCH(double) +GETRS_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) +GETRS_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) + +#undef GETRS_STRIDED_BATCH_LAUNCHER_SCRATCH + +#define ORGQR_STRIDED_BATCH_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t orgqr_batch_scratchpad_size( \ + sycl::queue & queue, std::int64_t m, std::int64_t n, std::int64_t k, std::int64_t lda, \ + std::int64_t stride_a, std::int64_t stride_tau, std::int64_t batch_size) { \ + return 0; \ + } + +ORGQR_STRIDED_BATCH_LAUNCHER_SCRATCH(float) +ORGQR_STRIDED_BATCH_LAUNCHER_SCRATCH(double) + +#undef ORGQR_STRIDED_BATCH_LAUNCHER_SCRATCH + +#define UNGQR_STRIDED_BATCH_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t ungqr_batch_scratchpad_size( \ + sycl::queue & queue, std::int64_t m, std::int64_t n, std::int64_t k, std::int64_t lda, \ + std::int64_t stride_a, std::int64_t stride_tau, std::int64_t batch_size) { \ + return 0; \ + } + +UNGQR_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) +UNGQR_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) + +#undef UNGQR_STRIDED_BATCH_LAUNCHER_SCRATCH + +#define POTRF_STRIDED_BATCH_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t potrf_batch_scratchpad_size( \ + sycl::queue & queue, oneapi::math::uplo uplo, std::int64_t n, std::int64_t lda, \ + std::int64_t stride_a, std::int64_t batch_size) { \ + return 0; \ + } + +POTRF_STRIDED_BATCH_LAUNCHER_SCRATCH(float) +POTRF_STRIDED_BATCH_LAUNCHER_SCRATCH(double) +POTRF_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) +POTRF_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) + +#undef POTRF_STRIDED_BATCH_LAUNCHER_SCRATCH + +#define POTRS_STRIDED_BATCH_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t potrs_batch_scratchpad_size( \ + sycl::queue & queue, oneapi::math::uplo uplo, std::int64_t n, std::int64_t nrhs, \ + std::int64_t lda, std::int64_t stride_a, std::int64_t ldb, std::int64_t stride_b, \ + std::int64_t batch_size) { \ + return 0; \ + } + +POTRS_STRIDED_BATCH_LAUNCHER_SCRATCH(float) +POTRS_STRIDED_BATCH_LAUNCHER_SCRATCH(double) +POTRS_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) +POTRS_STRIDED_BATCH_LAUNCHER_SCRATCH(std::complex) + +#undef POTRS_STRIDED_BATCH_LAUNCHER_SCRATCH + +#define GEQRF_GROUP_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t geqrf_batch_scratchpad_size( \ + sycl::queue & queue, std::int64_t* m, std::int64_t* n, std::int64_t* lda, \ + std::int64_t group_count, std::int64_t* group_sizes) { \ + return 0; \ + } + +GEQRF_GROUP_LAUNCHER_SCRATCH(float) +GEQRF_GROUP_LAUNCHER_SCRATCH(double) +GEQRF_GROUP_LAUNCHER_SCRATCH(std::complex) +GEQRF_GROUP_LAUNCHER_SCRATCH(std::complex) + +#undef GEQRF_GROUP_LAUNCHER_SCRATCH + +#define GETRF_GROUP_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t getrf_batch_scratchpad_size( \ + sycl::queue & queue, std::int64_t* m, std::int64_t* n, std::int64_t* lda, \ + std::int64_t group_count, std::int64_t* group_sizes) { \ + return 0; \ + } + +GETRF_GROUP_LAUNCHER_SCRATCH(float) +GETRF_GROUP_LAUNCHER_SCRATCH(double) +GETRF_GROUP_LAUNCHER_SCRATCH(std::complex) +GETRF_GROUP_LAUNCHER_SCRATCH(std::complex) + +#undef GETRF_GROUP_LAUNCHER_SCRATCH + +#define GETRI_GROUP_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t getri_batch_scratchpad_size(sycl::queue & queue, std::int64_t* n, \ + std::int64_t* lda, std::int64_t group_count, \ + std::int64_t* group_sizes) { \ + return 0; \ + } + +GETRI_GROUP_LAUNCHER_SCRATCH(float) +GETRI_GROUP_LAUNCHER_SCRATCH(double) +GETRI_GROUP_LAUNCHER_SCRATCH(std::complex) +GETRI_GROUP_LAUNCHER_SCRATCH(std::complex) + +#undef GETRI_GROUP_LAUNCHER_SCRATCH + +#define GETRS_GROUP_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t getrs_batch_scratchpad_size( \ + sycl::queue & queue, oneapi::math::transpose * trans, std::int64_t* n, std::int64_t* nrhs, \ + std::int64_t* lda, std::int64_t* ldb, std::int64_t group_count, \ + std::int64_t* group_sizes) { \ + return 0; \ + } + +GETRS_GROUP_LAUNCHER_SCRATCH(float) +GETRS_GROUP_LAUNCHER_SCRATCH(double) +GETRS_GROUP_LAUNCHER_SCRATCH(std::complex) +GETRS_GROUP_LAUNCHER_SCRATCH(std::complex) + +#undef GETRS_GROUP_LAUNCHER_SCRATCH + +#define ORGQR_GROUP_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t orgqr_batch_scratchpad_size( \ + sycl::queue & queue, std::int64_t* m, std::int64_t* n, std::int64_t* k, std::int64_t* lda, \ + std::int64_t group_count, std::int64_t* group_sizes) { \ + return 0; \ + } + +ORGQR_GROUP_LAUNCHER_SCRATCH(float) +ORGQR_GROUP_LAUNCHER_SCRATCH(double) + +#undef ORGQR_GROUP_LAUNCHER_SCRATCH + +#define UNGQR_GROUP_LAUNCHER_SCRATCH(TYPE) \ + template <> \ + std::int64_t ungqr_batch_scratchpad_size( \ + sycl::queue & queue, std::int64_t* m, std::int64_t* n, std::int64_t* k, std::int64_t* lda, \ + std::int64_t group_count, std::int64_t* group_sizes) { \ + return 0; \ + } + +UNGQR_GROUP_LAUNCHER_SCRATCH(std::complex) +UNGQR_GROUP_LAUNCHER_SCRATCH(std::complex) + +#undef UNGQR_GROUP_LAUNCHER_SCRATCH // rocsolverDnXpotrfBatched does not use scratchpad memory #define POTRF_GROUP_LAUNCHER_SCRATCH(TYPE) \ @@ -1047,23 +1600,6 @@ POTRS_GROUP_LAUNCHER_SCRATCH(std::complex) #undef POTRS_GROUP_LAUNCHER_SCRATCH -template <> -std::int64_t ungqr_batch_scratchpad_size>(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* k, - std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "ungqr_batch_scratchpad_size"); -} -template <> -std::int64_t ungqr_batch_scratchpad_size>(sycl::queue& queue, std::int64_t* m, - std::int64_t* n, std::int64_t* k, - std::int64_t* lda, - std::int64_t group_count, - std::int64_t* group_sizes) { - throw unimplemented("lapack", "ungqr_batch_scratchpad_size"); -} - } // namespace rocsolver } // namespace lapack } // namespace math diff --git a/src/lapack/backends/rocsolver/rocsolver_helper.hpp b/src/lapack/backends/rocsolver/rocsolver_helper.hpp index 5d4e6e821..a8d8c94b4 100644 --- a/src/lapack/backends/rocsolver/rocsolver_helper.hpp +++ b/src/lapack/backends/rocsolver/rocsolver_helper.hpp @@ -31,6 +31,10 @@ #include #include #include +#include +#include +#include +#include #include "oneapi/math/types.hpp" #include "runtime_support_helper.hpp" @@ -280,6 +284,44 @@ inline void lapack_info_check(sycl::queue& queue, DEVINFO_T devinfo, const char* devinfo_); } +inline void get_rocsolver_devinfo(sycl::queue& queue, sycl::buffer& devInfo, + std::vector& dev_info_) { + sycl::host_accessor dev_info_acc{ devInfo }; + for (std::size_t i = 0; i < dev_info_.size(); ++i) + dev_info_[i] = dev_info_acc[i]; +} + +inline void get_rocsolver_devinfo(sycl::queue& queue, const int* devInfo, + std::vector& dev_info_) { + queue.wait(); + queue.memcpy(dev_info_.data(), devInfo, sizeof(int) * dev_info_.size()).wait(); +} + +/* Reports every failing matrix of a batch. */ +template +inline void lapack_info_check_batch(sycl::queue& queue, DEVINFO_T devinfo, const char* func_name, + const char* rocfunc_name, std::int64_t batch_size) { + queue.wait(); + std::vector dev_info_(static_cast(batch_size)); + get_rocsolver_devinfo(queue, devinfo, dev_info_); + std::vector ids; + std::vector exceptions; + for (std::size_t i = 0; i < dev_info_.size(); ++i) { + const auto val = dev_info_[i]; + if (val > 0) { + ids.push_back(static_cast(i)); + exceptions.push_back(std::make_exception_ptr(oneapi::math::lapack::computation_error( + func_name, std::string(rocfunc_name) + " failed with info = " + std::to_string(val), + val))); + } + } + if (!ids.empty()) { + throw oneapi::math::lapack::batch_error( + func_name, std::string(rocfunc_name) + " failed for one or more matrices", + static_cast(ids.size()), std::move(ids), std::move(exceptions)); + } +} + } // namespace rocsolver } // namespace lapack } // namespace math diff --git a/src/lapack/backends/rocsolver/rocsolver_lapack.cpp b/src/lapack/backends/rocsolver/rocsolver_lapack.cpp index 5b0c265b2..140b7c481 100644 --- a/src/lapack/backends/rocsolver/rocsolver_lapack.cpp +++ b/src/lapack/backends/rocsolver/rocsolver_lapack.cpp @@ -76,28 +76,40 @@ GEBRD_LAUNCHER(std::complex, double, rocsolver_zgebrd) #undef GEBRD_LAUNCHER -void gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, - std::int64_t lda, sycl::buffer& tau, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "gerqf"); -} -void gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, - std::int64_t lda, sycl::buffer& tau, sycl::buffer& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "gerqf"); -} -void gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer>& a, - std::int64_t lda, sycl::buffer>& tau, - sycl::buffer>& scratchpad, std::int64_t scratchpad_size) { - throw unimplemented("lapack", "gerqf"); -} -void gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, - sycl::buffer>& a, std::int64_t lda, - sycl::buffer>& tau, sycl::buffer>& scratchpad, - std::int64_t scratchpad_size) { - throw unimplemented("lapack", "gerqf"); +template +inline void gerqf(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& tau, + sycl::buffer& scratchpad, std::int64_t scratchpad_size) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(m, n, lda, scratchpad_size); + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + auto tau_acc = tau.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 tau_ = sc.get_mem(tau_acc); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, m, n, a_, lda, tau_); + }); + }); } +#define GERQF_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + void gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, sycl::buffer& a, \ + std::int64_t lda, sycl::buffer& tau, sycl::buffer& scratchpad, \ + std::int64_t scratchpad_size) { \ + gerqf(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, tau, scratchpad, \ + scratchpad_size); \ + } + +GERQF_LAUNCHER(float, rocsolver_sgerqf) +GERQF_LAUNCHER(double, rocsolver_dgerqf) +GERQF_LAUNCHER(std::complex, rocsolver_cgerqf) +GERQF_LAUNCHER(std::complex, rocsolver_zgerqf) + +#undef GERQF_LAUNCHER + template inline void geqrf(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& tau, @@ -188,27 +200,59 @@ GETRF_LAUNCHER(std::complex, rocsolver_zgetrf) #undef GETRF_LAUNCHER -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) { - throw unimplemented("lapack", "getri"); -} -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) { - throw unimplemented("lapack", "getri"); -} -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) { - throw unimplemented("lapack", "getri"); -} -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) { - throw unimplemented("lapack", "getri"); +template +inline void getri(const char* func_name, Func func, 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) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(n, lda, scratchpad_size); + + // rocsolver legacy api does not accept 64-bit ints. + // To get around the limitation. + // Create new buffer and convert 64-bit values. + std::uint64_t ipiv_size = ipiv.size(); + sycl::buffer ipiv32(sycl::range<1>{ ipiv_size }); + sycl::buffer devInfo{ 1 }; + + queue.submit([&](sycl::handler& cgh) { + auto ipiv32_acc = ipiv32.template get_access(cgh); + auto ipiv_acc = ipiv.template get_access(cgh); + cgh.parallel_for(sycl::range<1>{ ipiv_size }, [=](sycl::id<1> index) { + ipiv32_acc[index] = static_cast(ipiv_acc[index]); + }); + }); + + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + auto ipiv_acc = ipiv32.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, n, a_, lda, ipiv_, devInfo_); + }); + }); + lapack_info_check(queue, devInfo, __func__, func_name); } +#define GETRI_LAUNCHER(TYPE, ROCSOLVER_ROUTINE) \ + 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) { \ + getri(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, n, a, lda, ipiv, scratchpad, \ + scratchpad_size); \ + } + +GETRI_LAUNCHER(float, rocsolver_sgetri) +GETRI_LAUNCHER(double, rocsolver_dgetri) +GETRI_LAUNCHER(std::complex, rocsolver_cgetri) +GETRI_LAUNCHER(std::complex, rocsolver_zgetri) + +#undef GETRI_LAUNCHER + 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, @@ -444,6 +488,8 @@ HETRD_LAUNCHER(std::complex, double, rocsolver_zhetrd) #undef HETRD_LAUNCHER +// rocsolver exposes sytrf/sytf2 for the symmetric factorization but has no +// Hermitian counterpart, so hetrf would have to be written from scratch. void hetrf(sycl::queue& queue, oneapi::math::uplo uplo, std::int64_t n, sycl::buffer>& a, std::int64_t lda, sycl::buffer& ipiv, sycl::buffer>& scratchpad, std::int64_t scratchpad_size) { @@ -595,6 +641,8 @@ ORMTR_LAUNCHER(double, rocsolver_dormtr) #undef ORMTR_LAUNCHER +// rocsolver ships orm2l/orm2r/ormbr/orml2/ormlq/ormql/ormqr/ormtr but no ormrq, +// so applying Q from an RQ factorization has no native entry point. void ormrq(sycl::queue& queue, oneapi::math::side side, oneapi::math::transpose trans, std::int64_t m, std::int64_t n, std::int64_t k, sycl::buffer& a, std::int64_t lda, sycl::buffer& tau, sycl::buffer& c, std::int64_t ldc, @@ -934,6 +982,9 @@ SYTRF_LAUNCHER(std::complex, rocsolver_zsytrf) #undef SYTRF_LAUNCHER +// rocsolver has no trtrs. The solve itself maps onto rocblas trsm, but the +// singularity scan that trtrs must report through info has no native +// equivalent and would need a dedicated kernel. void trtrs(sycl::queue& queue, oneapi::math::uplo uplo, oneapi::math::transpose trans, oneapi::math::diag diag, std::int64_t n, std::int64_t nrhs, sycl::buffer>& a, std::int64_t lda, @@ -1060,6 +1111,7 @@ UNGTR_LAUNCHER(std::complex, rocsolver_zungtr) #undef UNGTR_LAUNCHER +// The complex counterpart of ormrq, likewise absent from rocsolver. void unmrq(sycl::queue& queue, oneapi::math::side side, oneapi::math::transpose trans, std::int64_t m, std::int64_t n, std::int64_t k, sycl::buffer>& a, std::int64_t lda, sycl::buffer>& tau, @@ -1200,27 +1252,44 @@ GEBRD_LAUNCHER_USM(std::complex, double, rocsolver_zgebrd) #undef GEBRD_LAUNCHER_USM -sycl::event gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, float* a, std::int64_t lda, - float* tau, float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "gerqf"); -} -sycl::event gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, double* a, std::int64_t lda, - double* tau, double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "gerqf"); -} -sycl::event gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, std::complex* a, - std::int64_t lda, std::complex* tau, std::complex* scratchpad, - std::int64_t scratchpad_size, const std::vector& dependencies) { - throw unimplemented("lapack", "gerqf"); -} -sycl::event gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, std::complex* a, - std::int64_t lda, std::complex* tau, std::complex* scratchpad, - std::int64_t scratchpad_size, const std::vector& dependencies) { - throw unimplemented("lapack", "gerqf"); +template +inline sycl::event gerqf(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, + std::int64_t n, T* a, std::int64_t lda, T* tau, T* scratchpad, + std::int64_t scratchpad_size, + const std::vector& dependencies) { + using rocmDataType = typename RocmEquivalentType::Type; + overflow_check(m, n, lda, scratchpad_size); + auto done = queue.submit([&](sycl::handler& cgh) { + int64_t num_events = dependencies.size(); + for (int64_t i = 0; i < num_events; i++) { + cgh.depends_on(dependencies[i]); + } + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto a_ = reinterpret_cast(a); + auto tau_ = reinterpret_cast(tau); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, m, n, a_, lda, tau_); + }); + }); + return done; } +#define GERQF_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event gerqf(sycl::queue& queue, std::int64_t m, std::int64_t n, TYPE* a, \ + std::int64_t lda, TYPE* tau, TYPE* scratchpad, std::int64_t scratchpad_size, \ + const std::vector& dependencies) { \ + return gerqf(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, m, n, a, lda, tau, scratchpad, \ + scratchpad_size, dependencies); \ + } + +GERQF_LAUNCHER_USM(float, rocsolver_sgerqf) +GERQF_LAUNCHER_USM(double, rocsolver_dgerqf) +GERQF_LAUNCHER_USM(std::complex, rocsolver_cgerqf) +GERQF_LAUNCHER_USM(std::complex, rocsolver_zgerqf) + +#undef GERQF_LAUNCHER_USM + template inline sycl::event geqrf(const char* func_name, Func func, sycl::queue& queue, std::int64_t m, std::int64_t n, T* a, std::int64_t lda, T* tau, T* scratchpad, @@ -1298,7 +1367,14 @@ inline sycl::event getrf(const char* func_name, Func func, sycl::queue& queue, s }); }); - lapack_info_check(queue, devInfo, __func__, func_name); + try { + lapack_info_check(queue, devInfo, __func__, func_name); + } + catch (...) { + free(ipiv32, queue); + free(devInfo, queue); + throw; + } free(ipiv32, queue); free(devInfo, queue); return done_casting; @@ -1320,27 +1396,71 @@ GETRF_LAUNCHER_USM(std::complex, rocsolver_zgetrf) #undef GETRF_LAUNCHER_USM -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) { - throw unimplemented("lapack", "getri"); -} -sycl::event getri(sycl::queue& queue, std::int64_t n, double* a, std::int64_t lda, - std::int64_t* ipiv, double* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri"); -} -sycl::event getri(sycl::queue& queue, std::int64_t n, float* a, std::int64_t lda, - std::int64_t* ipiv, float* scratchpad, std::int64_t scratchpad_size, - const std::vector& dependencies) { - throw unimplemented("lapack", "getri"); -} -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) { - throw unimplemented("lapack", "getri"); +template +inline sycl::event getri(const char* func_name, Func func, sycl::queue& queue, 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; + overflow_check(n, lda, scratchpad_size); + + // rocsolver legacy api does not accept 64-bit ints. + // To get around the limitation. + // Create new buffer and convert 64-bit values. + std::uint64_t ipiv_size = n; + int* ipiv32 = (int*)malloc_device(sizeof(int) * ipiv_size, queue); + int* devInfo = (int*)malloc_device(sizeof(int), queue); + + auto done_casting = queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.parallel_for(sycl::range<1>{ ipiv_size }, [=](sycl::id<1> index) { + ipiv32[index] = static_cast(ipiv[index]); + }); + }); + + auto done = queue.submit([&](sycl::handler& cgh) { + int64_t num_events = dependencies.size(); + for (int64_t i = 0; i < num_events; i++) { + cgh.depends_on(dependencies[i]); + } + cgh.depends_on(done_casting); + onemath_rocsolver_host_task(cgh, queue, [=](RocsolverScopedContextHandler& sc) { + auto handle = sc.get_handle(queue); + auto a_ = reinterpret_cast(a); + auto ipiv_ = reinterpret_cast(ipiv32); + rocblas_status err; + rocsolver_native_named_func(func_name, func, err, handle, n, a_, lda, ipiv_, devInfo); + }); + }); + + try { + lapack_info_check(queue, devInfo, __func__, func_name); + } + catch (...) { + free(ipiv32, queue); + free(devInfo, queue); + throw; + } + free(ipiv32, queue); + free(devInfo, queue); + return done; } +#define GETRI_LAUNCHER_USM(TYPE, ROCSOLVER_ROUTINE) \ + sycl::event getri(sycl::queue& queue, 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 getri(#ROCSOLVER_ROUTINE, ROCSOLVER_ROUTINE, queue, n, a, lda, ipiv, scratchpad, \ + scratchpad_size, dependencies); \ + } + +GETRI_LAUNCHER_USM(float, rocsolver_sgetri) +GETRI_LAUNCHER_USM(double, rocsolver_dgetri) +GETRI_LAUNCHER_USM(std::complex, rocsolver_cgetri) +GETRI_LAUNCHER_USM(std::complex, rocsolver_zgetri) + +#undef GETRI_LAUNCHER_USM + 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, @@ -2386,22 +2506,22 @@ GEBRD_LAUNCHER_SCRATCH(std::complex) template <> std::int64_t gerqf_scratchpad_size(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t lda) { - throw unimplemented("lapack", "gerqf_scratchpad_size"); + return 0; } template <> std::int64_t gerqf_scratchpad_size(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t lda) { - throw unimplemented("lapack", "gerqf_scratchpad_size"); + return 0; } template <> std::int64_t gerqf_scratchpad_size>(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t lda) { - throw unimplemented("lapack", "gerqf_scratchpad_size"); + return 0; } template <> std::int64_t gerqf_scratchpad_size>(sycl::queue& queue, std::int64_t m, std::int64_t n, std::int64_t lda) { - throw unimplemented("lapack", "gerqf_scratchpad_size"); + return 0; } #define GEQRF_LAUNCHER_SCRATCH(TYPE) \ @@ -2449,21 +2569,21 @@ GETRF_LAUNCHER_SCRATCH(std::complex) template <> std::int64_t getri_scratchpad_size(sycl::queue& queue, std::int64_t n, std::int64_t lda) { - throw unimplemented("lapack", "getri_scratchpad_size"); + return 0; } template <> std::int64_t getri_scratchpad_size(sycl::queue& queue, std::int64_t n, std::int64_t lda) { - throw unimplemented("lapack", "getri_scratchpad_size"); + return 0; } template <> std::int64_t getri_scratchpad_size>(sycl::queue& queue, std::int64_t n, std::int64_t lda) { - throw unimplemented("lapack", "getri_scratchpad_size"); + return 0; } template <> std::int64_t getri_scratchpad_size>(sycl::queue& queue, std::int64_t n, std::int64_t lda) { - throw unimplemented("lapack", "getri_scratchpad_size"); + return 0; } #define GETRS_LAUNCHER_SCRATCH(TYPE) \