Reduce implementation code

Author: vlad-perevezentsev

dpnp/backend/extensions/lapack/CMakeLists.txt

@@ -41,7 +41,6 @@ set(_module_src
     ${CMAKE_CURRENT_SOURCE_DIR}/potrf.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/potrf_batch.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/syevd.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/syevd_batch.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/ungqr.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/ungqr_batch.cpp
 )

dpnp/backend/extensions/lapack/lapack_py.cpp

@@ -60,7 +60,6 @@ void init_dispatch_vectors(void)
     lapack_ext::init_orgqr_dispatch_vector();
     lapack_ext::init_potrf_batch_dispatch_vector();
     lapack_ext::init_potrf_dispatch_vector();
-    lapack_ext::init_syevd_batch_dispatch_vector();
     lapack_ext::init_syevd_dispatch_vector();
     lapack_ext::init_ungqr_batch_dispatch_vector();
     lapack_ext::init_ungqr_dispatch_vector();

dpnp/backend/extensions/lapack/syevd.cpp

@@ -230,6 +230,123 @@ std::pair<sycl::event, sycl::event>
     return std::make_pair(args_ev, syevd_ev);
 }
 
+std::pair<sycl::event, sycl::event>
+    syevd_batch(sycl::queue exec_q,
+                const std::int8_t jobz,
+                const std::int8_t upper_lower,
+                dpctl::tensor::usm_ndarray eig_vecs,
+                dpctl::tensor::usm_ndarray eig_vals,
+                const std::vector<sycl::event> &depends)
+{
+    const int eig_vecs_nd = eig_vecs.get_ndim();
+    const int eig_vals_nd = eig_vals.get_ndim();
+
+    if (eig_vecs_nd != 3) {
+        throw py::value_error("Unexpected ndim=" + std::to_string(eig_vecs_nd) +
+                              " of an output array with eigenvectors");
+    }
+    else if (eig_vals_nd != 2) {
+        throw py::value_error("Unexpected ndim=" + std::to_string(eig_vals_nd) +
+                              " of an output array with eigenvalues");
+    }
+
+    const py::ssize_t *eig_vecs_shape = eig_vecs.get_shape_raw();
+    const py::ssize_t *eig_vals_shape = eig_vals.get_shape_raw();
+
+    if (eig_vecs_shape[1] != eig_vecs_shape[2]) {
+        throw py::value_error(
+            "The last two dimensions of 'eig_vecs' must be the same.");
+    }
+    else if (eig_vecs_shape[0] != eig_vals_shape[0] ||
+             eig_vecs_shape[1] != eig_vals_shape[1])
+    {
+        throw py::value_error(
+            "The shape of 'eig_vals' must be (batch_size, n), "
+            "where batch_size = " +
+            std::to_string(eig_vecs_shape[0]) +
+            " and n = " + std::to_string(eig_vecs_shape[1]));
+    }
+
+    size_t src_nelems(1);
+
+    for (int i = 0; i < eig_vecs_nd; ++i) {
+        src_nelems *= static_cast<size_t>(eig_vecs_shape[i]);
+    }
+
+    if (src_nelems == 0) {
+        // nothing to do
+        return std::make_pair(sycl::event(), sycl::event());
+    }
+
+    // check compatibility of execution queue and allocation queue
+    if (!dpctl::utils::queues_are_compatible(exec_q, {eig_vecs, eig_vals})) {
+        throw py::value_error(
+            "Execution queue is not compatible with allocation queues");
+    }
+
+    auto const &overlap = dpctl::tensor::overlap::MemoryOverlap();
+    if (overlap(eig_vecs, eig_vals)) {
+        throw py::value_error("Arrays with eigenvectors and eigenvalues are "
+                              "overlapping segments of memory");
+    }
+
+    bool is_eig_vecs_c_contig = eig_vecs.is_c_contiguous();
+    bool is_eig_vals_c_contig = eig_vals.is_c_contiguous();
+    if (!is_eig_vecs_c_contig) {
+        throw py::value_error(
+            "An array with input matrix / output eigenvectors "
+            "must be C-contiguous");
+    }
+    else if (!is_eig_vals_c_contig) {
+        throw py::value_error(
+            "An array with output eigenvalues must be C-contiguous");
+    }
+
+    auto array_types = dpctl_td_ns::usm_ndarray_types();
+    int eig_vecs_type_id =
+        array_types.typenum_to_lookup_id(eig_vecs.get_typenum());
+    int eig_vals_type_id =
+        array_types.typenum_to_lookup_id(eig_vals.get_typenum());
+
+    if (eig_vecs_type_id != eig_vals_type_id) {
+        throw py::value_error(
+            "Types of eigenvectors and eigenvalues are mismatched");
+    }
+
+    syevd_impl_fn_ptr_t syevd_fn = syevd_dispatch_vector[eig_vecs_type_id];
+    if (syevd_fn == nullptr) {
+        throw py::value_error("No syevd implementation defined for a type of "
+                              "eigenvectors and eigenvalues");
+    }
+
+    char *eig_vecs_data = eig_vecs.get_data();
+    char *eig_vals_data = eig_vals.get_data();
+
+    const std::int64_t batch_size = eig_vecs_shape[0];
+    const std::int64_t n = eig_vecs_shape[1];
+    int elemsize = eig_vecs.get_elemsize();
+
+    const oneapi::mkl::job jobz_val = static_cast<oneapi::mkl::job>(jobz);
+    const oneapi::mkl::uplo uplo_val =
+        static_cast<oneapi::mkl::uplo>(upper_lower);
+
+    std::vector<sycl::event> host_task_events;
+
+    for (std::int64_t i = 0; i < batch_size; ++i) {
+        char *eig_vecs_batch = eig_vecs_data + i * n * n * elemsize;
+        char *eig_vals_batch = eig_vals_data + i * n * elemsize;
+
+        sycl::event syevd_ev =
+            syevd_fn(exec_q, jobz_val, uplo_val, n, eig_vecs_batch,
+                     eig_vals_batch, host_task_events, depends);
+    }
+
+    sycl::event args_ev = dpctl::utils::keep_args_alive(
+        exec_q, {eig_vecs, eig_vals}, host_task_events);
+
+    return std::make_pair(args_ev, args_ev);
+}
+
 template <typename fnT, typename T>
 struct SyevdContigFactory
 {

dpnp/backend/extensions/lapack/syevd.hpp

@@ -55,7 +55,6 @@ extern std::pair<sycl::event, sycl::event>
                 const std::vector<sycl::event> &depends = {});
 
 extern void init_syevd_dispatch_vector(void);
-extern void init_syevd_batch_dispatch_vector(void);
 } // namespace lapack
 } // namespace ext
 } // namespace backend