From f779c2d7f45979e9ac9643a04b6291cb1c67b225 Mon Sep 17 00:00:00 2001 From: Luc Berger-Vergiat Date: Fri, 16 Jun 2023 00:04:04 +0000 Subject: [PATCH] ONEMKL: working version of spmv mkl backend This should be good for now but eventually more work is needed for the multivector case which will use the gemm function instead of gemv from onemkl Signed-off-by: Berger-Vergiat --- sparse/src/KokkosSparse_Utils_mkl.hpp | 9 ++ sparse/src/KokkosSparse_spmv.hpp | 10 +- ...sSparse_spgemm_symbolic_tpl_spec_avail.hpp | 2 +- .../tpls/KokkosSparse_spmv_tpl_spec_avail.hpp | 39 +++++ .../tpls/KokkosSparse_spmv_tpl_spec_decl.hpp | 150 ++++++++++++++++++ sparse/unit_test/Test_Sparse.hpp | 1 - 6 files changed, 207 insertions(+), 4 deletions(-) diff --git a/sparse/src/KokkosSparse_Utils_mkl.hpp b/sparse/src/KokkosSparse_Utils_mkl.hpp index 0afa75de0a..7a8dd0cb22 100644 --- a/sparse/src/KokkosSparse_Utils_mkl.hpp +++ b/sparse/src/KokkosSparse_Utils_mkl.hpp @@ -230,6 +230,15 @@ inline void MKLSparseMatrix>::export_data( } // namespace Impl } // namespace KokkosSparse +// Utilities for oneMKL SYCL code +#ifdef KOKKOS_ENABLE_SYCL +#include "oneapi/mkl/spblas.hpp" + +namespace KokkosSparse { +namespace Impl {} +} // namespace KokkosSparse +#endif // KOKKOS_ENABLE_SYCL + #endif // KOKKOSKERNELS_ENABLE_TPL_MKL #endif // _KOKKOSKERNELS_SPARSEUTILS_MKL_HPP diff --git a/sparse/src/KokkosSparse_spmv.hpp b/sparse/src/KokkosSparse_spmv.hpp index 60fb5331cf..f43ec0bd54 100644 --- a/sparse/src/KokkosSparse_spmv.hpp +++ b/sparse/src/KokkosSparse_spmv.hpp @@ -171,10 +171,16 @@ void spmv(KokkosKernels::Experimental::Controls controls, const char mode[], #endif #ifdef KOKKOSKERNELS_ENABLE_TPL_MKL - if (std::is_same::value) { + if (std::is_same_v) { useFallback = useFallback || (mode[0] == Conjugate[0]); } +#ifdef KOKKOS_ENABLE_SYCL + if (std::is_same_v) { + useFallback = useFallback || (mode[0] == Conjugate[0]); + } +#endif #endif if (useFallback) { diff --git a/sparse/tpls/KokkosSparse_spgemm_symbolic_tpl_spec_avail.hpp b/sparse/tpls/KokkosSparse_spgemm_symbolic_tpl_spec_avail.hpp index 1fcfa7132a..b8c545ffe2 100644 --- a/sparse/tpls/KokkosSparse_spgemm_symbolic_tpl_spec_avail.hpp +++ b/sparse/tpls/KokkosSparse_spgemm_symbolic_tpl_spec_avail.hpp @@ -141,7 +141,7 @@ SPGEMM_SYMBOLIC_AVAIL_MKL_E(Kokkos::Serial) #ifdef KOKKOS_ENABLE_OPENMP SPGEMM_SYMBOLIC_AVAIL_MKL_E(Kokkos::OpenMP) #endif -#endif +#endif // KOKKOSKERNELS_ENABLE_TPL_MKL } // namespace Impl } // namespace KokkosSparse diff --git a/sparse/tpls/KokkosSparse_spmv_tpl_spec_avail.hpp b/sparse/tpls/KokkosSparse_spmv_tpl_spec_avail.hpp index 060fef45bb..a8632263f9 100644 --- a/sparse/tpls/KokkosSparse_spmv_tpl_spec_avail.hpp +++ b/sparse/tpls/KokkosSparse_spmv_tpl_spec_avail.hpp @@ -242,6 +242,45 @@ KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_MKL(Kokkos::complex, Kokkos::OpenMP) KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_MKL(Kokkos::complex, Kokkos::OpenMP) #endif +#ifdef KOKKOS_ENABLE_SYCL +#define KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL(SCALAR, ORDINAL, MEMSPACE) \ + template <> \ + struct spmv_tpl_spec_avail< \ + const SCALAR, const ORDINAL, \ + Kokkos::Device, \ + Kokkos::MemoryTraits, const ORDINAL, const SCALAR*, \ + Kokkos::LayoutLeft, \ + Kokkos::Device, \ + Kokkos::MemoryTraits, SCALAR*, \ + Kokkos::LayoutLeft, \ + Kokkos::Device, \ + Kokkos::MemoryTraits > { \ + enum : bool { value = true }; \ + }; + +KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL( + float, std::int32_t, Kokkos::Experimental::SYCLDeviceUSMSpace) +KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL( + double, std::int32_t, Kokkos::Experimental::SYCLDeviceUSMSpace) +KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL( + Kokkos::complex, std::int32_t, + Kokkos::Experimental::SYCLDeviceUSMSpace) +KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL( + Kokkos::complex, std::int32_t, + Kokkos::Experimental::SYCLDeviceUSMSpace) + +KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL( + float, std::int64_t, Kokkos::Experimental::SYCLDeviceUSMSpace) +KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL( + double, std::int64_t, Kokkos::Experimental::SYCLDeviceUSMSpace) +KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL( + Kokkos::complex, std::int64_t, + Kokkos::Experimental::SYCLDeviceUSMSpace) +KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL( + Kokkos::complex, std::int64_t, + Kokkos::Experimental::SYCLDeviceUSMSpace) +#endif + #endif // KOKKOSKERNELS_ENABLE_TPL_MKL } // namespace Impl diff --git a/sparse/tpls/KokkosSparse_spmv_tpl_spec_decl.hpp b/sparse/tpls/KokkosSparse_spmv_tpl_spec_decl.hpp index ecbe45c7fd..9b37361e65 100644 --- a/sparse/tpls/KokkosSparse_spmv_tpl_spec_decl.hpp +++ b/sparse/tpls/KokkosSparse_spmv_tpl_spec_decl.hpp @@ -755,6 +755,156 @@ KOKKOSSPARSE_SPMV_MKL(Kokkos::complex, Kokkos::OpenMP, #endif #undef KOKKOSSPARSE_SPMV_MKL + +#ifdef KOKKOS_ENABLE_SYCL +inline oneapi::mkl::transpose mode_kk_to_onemkl(char mode_kk) { + switch (toupper(mode_kk)) { + case 'N': return oneapi::mkl::transpose::nontrans; + case 'T': return oneapi::mkl::transpose::trans; + case 'H': return oneapi::mkl::transpose::conjtrans; + default:; + } + throw std::invalid_argument( + "Invalid mode for oneMKL (should be one of N, T, H)"); +} + +template +struct spmv_onemkl_wrapper {}; + +template <> +struct spmv_onemkl_wrapper { + template + static void spmv(const execution_space& exec, + oneapi::mkl::transpose const mkl_mode, + typename matrix_type::non_const_value_type const alpha, + const matrix_type& A, const xview_type& x, + typename matrix_type::non_const_value_type const beta, + const yview_type& y) { + using scalar_type = typename matrix_type::non_const_value_type; + using ordinal_type = typename matrix_type::non_const_ordinal_type; + + oneapi::mkl::sparse::matrix_handle_t handle = nullptr; + oneapi::mkl::sparse::init_matrix_handle(&handle); + auto ev_set = oneapi::mkl::sparse::set_csr_data( + exec.sycl_queue(), handle, A.numRows(), A.numCols(), + oneapi::mkl::index_base::zero, + const_cast(A.graph.row_map.data()), + const_cast(A.graph.entries.data()), + const_cast(A.values.data())); + auto ev_opt = oneapi::mkl::sparse::optimize_gemv( + exec.sycl_queue(), mkl_mode, handle, {ev_set}); + auto ev_gemv = + oneapi::mkl::sparse::gemv(exec.sycl_queue(), mkl_mode, alpha, handle, + x.data(), beta, y.data(), {ev_opt}); + auto ev_release = oneapi::mkl::sparse::release_matrix_handle( + exec.sycl_queue(), &handle, {ev_gemv}); + ev_release.wait(); + } +}; + +template <> +struct spmv_onemkl_wrapper { + template + static void spmv(const execution_space& exec, + oneapi::mkl::transpose const mkl_mode, + typename matrix_type::non_const_value_type const alpha, + const matrix_type& A, const xview_type& x, + typename matrix_type::non_const_value_type const beta, + const yview_type& y) { + using scalar_type = typename matrix_type::non_const_value_type; + using ordinal_type = typename matrix_type::non_const_ordinal_type; + using mag_type = typename Kokkos::ArithTraits::mag_type; + + oneapi::mkl::sparse::matrix_handle_t handle = nullptr; + oneapi::mkl::sparse::init_matrix_handle(&handle); + auto ev_set = oneapi::mkl::sparse::set_csr_data( + exec.sycl_queue(), handle, static_cast(A.numRows()), + static_cast(A.numCols()), oneapi::mkl::index_base::zero, + const_cast(A.graph.row_map.data()), + const_cast(A.graph.entries.data()), + reinterpret_cast*>( + const_cast(A.values.data()))); + auto ev_opt = oneapi::mkl::sparse::optimize_gemv( + exec.sycl_queue(), mkl_mode, handle, {ev_set}); + auto ev_gemv = oneapi::mkl::sparse::gemv( + exec.sycl_queue(), mkl_mode, alpha, handle, + reinterpret_cast*>( + const_cast(x.data())), + beta, reinterpret_cast*>(y.data()), {ev_opt}); + auto ev_release = oneapi::mkl::sparse::release_matrix_handle( + exec.sycl_queue(), &handle, {ev_gemv}); + ev_release.wait(); + } +}; + +#define KOKKOSSPARSE_SPMV_ONEMKL(SCALAR, ORDINAL, MEMSPACE, COMPILE_LIBRARY) \ + template <> \ + struct SPMV, \ + Kokkos::MemoryTraits, ORDINAL const, \ + SCALAR const*, Kokkos::LayoutLeft, \ + Kokkos::Device, \ + Kokkos::MemoryTraits, \ + SCALAR*, Kokkos::LayoutLeft, \ + Kokkos::Device, \ + Kokkos::MemoryTraits, true, \ + COMPILE_LIBRARY> { \ + using execution_space = Kokkos::Experimental::SYCL; \ + using device_type = Kokkos::Device; \ + using AMatrix = \ + CrsMatrix, ORDINAL const>; \ + using XVector = Kokkos::View< \ + SCALAR const*, Kokkos::LayoutLeft, device_type, \ + Kokkos::MemoryTraits>; \ + using YVector = Kokkos::View>; \ + using coefficient_type = typename YVector::non_const_value_type; \ + using Controls = KokkosKernels::Experimental::Controls; \ + \ + static void spmv(const Controls&, const char mode[], \ + const coefficient_type& alpha, const AMatrix& A, \ + const XVector& x, const coefficient_type& beta, \ + const YVector& y) { \ + std::string label = "KokkosSparse::spmv[TPL_ONEMKL," + \ + Kokkos::ArithTraits::name() + "]"; \ + Kokkos::Profiling::pushRegion(label); \ + oneapi::mkl::transpose mkl_mode = mode_kk_to_onemkl(mode[0]); \ + execution_space exec{}; \ + spmv_onemkl_wrapper::is_complex>::spmv( \ + exec, mkl_mode, alpha, A, x, beta, y); \ + Kokkos::Profiling::popRegion(); \ + } \ + }; + +KOKKOSSPARSE_SPMV_ONEMKL(float, std::int32_t, + Kokkos::Experimental::SYCLDeviceUSMSpace, + KOKKOSKERNELS_IMPL_COMPILE_LIBRARY) +KOKKOSSPARSE_SPMV_ONEMKL(double, std::int32_t, + Kokkos::Experimental::SYCLDeviceUSMSpace, + KOKKOSKERNELS_IMPL_COMPILE_LIBRARY) +KOKKOSSPARSE_SPMV_ONEMKL(Kokkos::complex, std::int32_t, + Kokkos::Experimental::SYCLDeviceUSMSpace, + KOKKOSKERNELS_IMPL_COMPILE_LIBRARY) +KOKKOSSPARSE_SPMV_ONEMKL(Kokkos::complex, std::int32_t, + Kokkos::Experimental::SYCLDeviceUSMSpace, + KOKKOSKERNELS_IMPL_COMPILE_LIBRARY) + +KOKKOSSPARSE_SPMV_ONEMKL(float, std::int64_t, + Kokkos::Experimental::SYCLDeviceUSMSpace, + KOKKOSKERNELS_IMPL_COMPILE_LIBRARY) +KOKKOSSPARSE_SPMV_ONEMKL(double, std::int64_t, + Kokkos::Experimental::SYCLDeviceUSMSpace, + KOKKOSKERNELS_IMPL_COMPILE_LIBRARY) +KOKKOSSPARSE_SPMV_ONEMKL(Kokkos::complex, std::int64_t, + Kokkos::Experimental::SYCLDeviceUSMSpace, + KOKKOSKERNELS_IMPL_COMPILE_LIBRARY) +KOKKOSSPARSE_SPMV_ONEMKL(Kokkos::complex, std::int64_t, + Kokkos::Experimental::SYCLDeviceUSMSpace, + KOKKOSKERNELS_IMPL_COMPILE_LIBRARY) +#endif } // namespace Impl } // namespace KokkosSparse #endif diff --git a/sparse/unit_test/Test_Sparse.hpp b/sparse/unit_test/Test_Sparse.hpp index e0d0085be1..0dbf7bc759 100644 --- a/sparse/unit_test/Test_Sparse.hpp +++ b/sparse/unit_test/Test_Sparse.hpp @@ -52,7 +52,6 @@ // to actually define tests. #include "Test_Sparse_Utils_cusparse.hpp" - #include "Test_Sparse_rocsparse.hpp" #endif // TEST_SPARSE_HPP