[FEA] Add sparse LOBPCG solver

cpp/include/raft/linalg/detail/cusolver_wrappers.hpp

@@ -774,6 +774,66 @@ inline cusolverStatus_t cusolverDnpotrf(cusolverDnHandle_t handle,  // NOLINT
 }
 /** @} */
 
+/**
+ * @defgroup potri cusolver potri operations: inverse of a matrix A using Cholesky
+ * @{
+ */
+template <typename T>
+cusolverStatus_t cusolverDnpotri_bufferSize(
+  cusolverDnHandle_t handle, cublasFillMode_t uplo, int n, T* A, int lda, int* Lwork);
+template <>
+inline cusolverStatus_t cusolverDnpotri_bufferSize(
+  cusolverDnHandle_t handle, cublasFillMode_t uplo, int n, float* A, int lda, int* Lwork)
+{
+  return cusolverDnSpotri_bufferSize(handle, uplo, n, A, lda, Lwork);
+}
+template <>
+inline cusolverStatus_t cusolverDnpotri_bufferSize(
+  cusolverDnHandle_t handle, cublasFillMode_t uplo, int n, double* A, int lda, int* Lwork)
+{
+  return cusolverDnDpotri_bufferSize(handle, uplo, n, A, lda, Lwork);
+}
+
+template <typename T>
+cusolverStatus_t cusolverDnpotri(cusolverDnHandle_t handle,
+                                 cublasFillMode_t uplo,
+                                 int n,
+                                 T* A,
+                                 int lda,
+                                 T* Workspace,
+                                 int Lwork,
+                                 int* devInfo,
+                                 cudaStream_t stream);
+template <>
+inline cusolverStatus_t cusolverDnpotri(cusolverDnHandle_t handle,
+                                        cublasFillMode_t uplo,
+                                        int n,
+                                        float* A,
+                                        int lda,
+                                        float* Workspace,
+                                        int Lwork,
+                                        int* devInfo,
+                                        cudaStream_t stream)
+{
+  RAFT_CUSOLVER_TRY(cusolverDnSetStream(handle, stream));
+  return cusolverDnSpotri(handle, uplo, n, A, lda, Workspace, Lwork, devInfo);
+}
+template <>
+inline cusolverStatus_t cusolverDnpotri(cusolverDnHandle_t handle,
+                                        cublasFillMode_t uplo,
+                                        int n,
+                                        double* A,
+                                        int lda,
+                                        double* Workspace,
+                                        int Lwork,
+                                        int* devInfo,
+                                        cudaStream_t stream)
+{
+  RAFT_CUSOLVER_TRY(cusolverDnSetStream(handle, stream));
+  return cusolverDnDpotri(handle, uplo, n, A, lda, Workspace, Lwork, devInfo);
+}
+/** @} */
+
 /**
  * @defgroup potrs cusolver potrs operations
  * @{

cpp/include/raft/linalg/eig.cuh

@@ -133,7 +133,7 @@ void eig_dc(raft::device_resources const& handle,
             raft::device_vector_view<ValueType, IndexType> eig_vals)
 {
   RAFT_EXPECTS(in.size() == eig_vectors.size(), "Size mismatch between Input and Eigen Vectors");
-  RAFT_EXPECTS(eig_vals.size() == in.extent(1), "Size mismatch between Input and Eigen Values");
+  RAFT_EXPECTS(eig_vals.extent(0) == in.extent(1), "Size mismatch between Input and Eigen Values");
 
   eigDC(handle,
         in.data_handle(),

cpp/include/raft/linalg/gemv.cuh

@@ -233,8 +233,8 @@ void gemv(raft::device_resources const& handle,
  * @tparam LayoutPolicyZ layout of Z
  * @param[in] handle raft handle
  * @param[in] A input raft::device_matrix_view of size (M, N)
- * @param[in] x input raft::device_matrix_view of size (N, 1) if A is raft::col_major, else (M, 1)
- * @param[out] y output raft::device_matrix_view of size (M, 1) if A is raft::col_major, else (N, 1)
+ * @param[in] x input raft::device_vector_view of size (N, 1) if A is raft::col_major, else (M, 1)
+ * @param[out] y output raft::device_vector_view of size (M, 1) if A is raft::col_major, else (N, 1)
  * @param[in] alpha optional raft::host_scalar_view or raft::device_scalar_view, default 1.0
  * @param[in] beta optional raft::host_scalar_view or raft::device_scalar_view, default 0.0
  */

cpp/include/raft/matrix/detail/matrix.cuh

@@ -197,6 +197,47 @@ void sliceMatrix(const m_t* in,
   slice<<<grid, block, 0, stream>>>(in, n_rows, n_cols, out, x1, y1, x2, y2);
 }
 
+/**
+ * @brief Kernel for copying a small matrix inside of a bigger matrix with a
+ * size matches that slice
+ * @param src_d: input matrix
+ * @param m: number of rows of input matrix
+ * @param n: number of columns of input matrix
+ * @param dst_d: output matrix
+ * @param x1, y1: coordinate of the top-left point of the wanted area (0-based)
+ * @param x2, y2: coordinate of the bottom-right point of the wanted area
+ * (1-based)
+ */
+template <typename m_t, typename idx_t = int>
+__global__ void slice_insert(
+  const m_t* src_d, idx_t n_rows, idx_t n_cols, m_t* dst_d, idx_t x1, idx_t y1, idx_t x2, idx_t y2)
+{
+  idx_t idx = threadIdx.x + blockDim.x * blockIdx.x;
+  idx_t dm = x2 - x1, dn = y2 - y1;
+  if (idx < dm * dn) {
+    idx_t i = idx % dm, j = idx / dm;
+    idx_t is = i + x1, js = j + y1;
+    dst_d[is + js * n_rows] = src_d[idx];
+  }
+}
+
+template <typename m_t, typename idx_t = int>
+void sliceMatrix_insert(const m_t* in,
+                        idx_t n_rows,
+                        idx_t n_cols,
+                        m_t* out,
+                        idx_t x1,
+                        idx_t y1,
+                        idx_t x2,
+                        idx_t y2,
+                        cudaStream_t stream)
+{
+  // Slicing
+  dim3 block(64);
+  dim3 grid(((x2 - x1) * (y2 - y1) + block.x - 1) / block.x);
+  slice_insert<<<grid, block, 0, stream>>>(in, n_rows, n_cols, out, x1, y1, x2, y2);
+}
+
 /**
  * @brief Kernel for copying the upper triangular part of a matrix to another
  * @param src: input matrix with a size of mxn
@@ -226,6 +267,60 @@ void copyUpperTriangular(const m_t* src, m_t* dst, idx_t n_rows, idx_t n_cols, c
   getUpperTriangular<<<grid, block, 0, stream>>>(src, dst, m, n, k);
 }
 
+/**
+ * @brief Kernel for copying the lower triangular part of a matrix to another
+ * @param src: input matrix with a size of mxn
+ * @param dst: output matrix with a size of kxk
+ * @param n_rows: number of rows of input matrix
+ * @param n_cols: number of columns of input matrix
+ * @param k: min(n_rows, n_cols)
+ */
+template <typename m_t, typename idx_t = int>
+__global__ void getLowerTriangular(const m_t* src, m_t* dst, idx_t n_rows, idx_t n_cols, idx_t k)
+{
+  idx_t idx = threadIdx.x + blockDim.x * blockIdx.x;
+  idx_t m = n_rows, n = n_cols;
+  if (idx < m * n) {
+    idx_t i = idx % m, j = idx / m;
+    if (i < k && j < k && j <= i) { dst[i + j * k] = src[idx]; }
+  }
+}
+
+template <typename m_t, typename idx_t = int>
+void copyLowerTriangular(const m_t* src, m_t* dst, idx_t n_rows, idx_t n_cols, cudaStream_t stream)
+{
+  idx_t m = n_rows, n = n_cols;
+  idx_t k = std::min(m, n);
+  dim3 block(64);
+  dim3 grid((m * n + block.x - 1) / block.x);
+  getLowerTriangular<<<grid, block, 0, stream>>>(src, dst, m, n, k);
+}
+
+/**
+ * @brief Create a diagonal identity matrix
+ * @param matrix: matrix of size n_rows x n_cols
+ * @param n_rows: number of rows of the matrix
+ * @param n_cols: number of columns of the matrix
+ */
+template <typename m_t, typename idx_t = int>
+__global__ void createEyeKernel(m_t* matrix, idx_t n_rows, idx_t n_cols)
+{
+  idx_t idx = threadIdx.x + blockDim.x * blockIdx.x;
+  if (idx < n_rows * n_cols) {
+    idx_t i = idx % n_rows, j = idx / n_rows;
+    matrix[idx] = m_t(j == i);
+  }
+}
+
+template <typename m_t, typename idx_t = int>
+void createEye(m_t* matrix, idx_t n_rows, idx_t n_cols, cudaStream_t stream)
+{
+  idx_t m = n_rows, n = n_cols;
+  dim3 block(64);
+  dim3 grid((m * n + block.x - 1) / block.x);
+  createEyeKernel<<<grid, block, 0, stream>>>(matrix, n_rows, n_cols);
+}
+
 /**
  * @brief Copy a vector to the diagonal of a matrix
  * @param vec: vector of length k = min(n_rows, n_cols)

cpp/include/raft/matrix/init.cuh

@@ -69,6 +69,13 @@ void fill(raft::device_resources const& handle,
   linalg::map(handle, inout, raft::const_op{scalar});
 }
 
+template <typename math_t, typename idx_t>
+void eye(const raft::handle_t& handle,
+         raft::device_matrix_view<math_t, idx_t, raft::col_major> inout)
+{
+  detail::createEye(inout.data_handle(), inout.extent(0), inout.extent(1), handle.get_stream());
+}
+
 /** @} */  // end of group matrix_init
 
 }  // namespace raft::matrix

cpp/include/raft/matrix/slice.cuh

@@ -74,6 +74,40 @@ void slice(raft::device_resources const& handle,
                       handle.get_stream());
 }
 
+/**
+ * @brief Insert a small matrix into a bigger matrix using a slice (in-place)
+ * @tparam m_t type of matrix elements
+ * @tparam idx_t integer type used for indexing
+ * @param[in] handle: raft handle
+ * @param[in] in: input matrix (column-major)
+ * @param[out] out: output matrix (column-major)
+ * @param[in] coords: coordinates of the insertion slice
+ * example: Slice the 2nd and 3rd columns of a 4x3 matrix: slice(handle, in, out, {0, 1, 4, 3});
+ */
+template <typename m_t, typename idx_t>
+void slice_insert(raft::device_resources const& handle,
+                  raft::device_matrix_view<m_t, idx_t, col_major> in,
+                  raft::device_matrix_view<m_t, idx_t, col_major> out,
+                  slice_coordinates<idx_t> coords)
+{
+  RAFT_EXPECTS(coords.row2 > coords.row1, "row2 must be > row1");
+  RAFT_EXPECTS(coords.col2 > coords.col1, "col2 must be > col1");
+  RAFT_EXPECTS(coords.row1 >= 0, "row1 must be >= 0");
+  RAFT_EXPECTS(coords.row2 <= out.extent(0), "row2 must be <= number of rows in the output matrix");
+  RAFT_EXPECTS(coords.col1 >= 0, "col1 must be >= 0");
+  RAFT_EXPECTS(coords.col2 <= out.extent(1),
+               "col2 must be <= number of columns in the output matrix");
+
+  detail::sliceMatrix_insert(in.data_handle(),
+                             out.extent(0),
+                             out.extent(1),
+                             out.data_handle(),
+                             coords.row1,
+                             coords.col1,
+                             coords.row2,
+                             coords.col2,
+                             handle.get_stream());
+}
 /** @} */  // end group matrix_slice
 
 }  // namespace raft::matrix

cpp/include/raft/matrix/triangular.cuh

@@ -44,6 +44,24 @@ void upper_triangular(raft::device_resources const& handle,
     src.data_handle(), dst.data_handle(), src.extent(0), src.extent(1), handle.get_stream());
 }
 
+/**
+ * @brief Copy the lower triangular part of a matrix to another
+ * @param[in] handle: raft handle
+ * @param[in] src: input matrix with a size of n_rows x n_cols
+ * @param[out] dst: output matrix with a size of kxk, k = min(n_rows, n_cols)
+ */
+template <typename m_t, typename idx_t>
+void lower_triangular(const raft::handle_t& handle,
+                      raft::device_matrix_view<const m_t, idx_t, col_major> src,
+                      raft::device_matrix_view<m_t, idx_t, col_major> dst)
+{
+  auto k = std::min(src.extent(0), src.extent(1));
+  RAFT_EXPECTS(k == dst.extent(0) && k == dst.extent(1),
+               "dst should be of size kxk, k = min(n_rows, n_cols)");
+  detail::copyLowerTriangular(
+    src.data_handle(), dst.data_handle(), src.extent(0), src.extent(1), handle.get_stream());
+}
+
 /** @} */  // end group matrix_triangular
 
 }  // namespace raft::matrix