Clean the code and add docs

Author: Cstandardlib

source/source_base/module_container/ATen/kernels/lapack.cpp

@@ -3,7 +3,6 @@
 
 #include <base/third_party/lapack.h>
 
-// #include <cstring> // std::memcpy
 #include <algorithm> // std::copy
 #include <complex>
 #include <stdexcept>
@@ -208,10 +207,6 @@ struct lapack_hegvd<T, DEVICE_CPU> {
         // first copy Mat_A to eigen_vec
         // then pass as argument "A" in lapack hegvd
         // and this block of memory will be overwritten by eigenvectors
-        // for (int i = 0; i < dim * lda; ++i){
-        //     eigen_vec[i] = Mat_A[i];
-        // }
-        // std::memcpy(eigen_vec, Mat_A, sizeof(T) * dim * lda);
         // eigen_vec = Mat_A
         std::copy(Mat_A, Mat_A + dim*lda, eigen_vec);
 

source/source_base/module_container/ATen/kernels/lapack.h

@@ -40,7 +40,18 @@ struct lapack_potrf {
         const int& lda);
 };
 
-
+// ============================================================================
+// Standard Hermitian Eigenvalue Problem Solvers
+// ============================================================================
+// The following structures (lapack_heevd and lapack_heevx) implement solvers
+// for standard Hermitian eigenvalue problems of the form:
+//      A * x = lambda * x
+// where:
+//   - A is a Hermitian matrix
+//   - lambda are the eigenvalues to be computed
+//   - x are the corresponding eigenvectors
+//
+// ============================================================================
 template <typename T, typename Device>
 struct lapack_heevd {
     using Real = typename GetTypeReal<T>::type;
@@ -61,18 +72,15 @@ struct lapack_heevx {
      * This function solves the problem A*x = lambda*x, where A is a Hermitian matrix.
      * It computes a subset of eigenvalues and, optionally, the corresponding eigenvectors.
      *
-     * @param jobz  'N': Compute eigenvalues only; 'V': Compute eigenvalues and eigenvectors.
-     * @param range 'A': All eigenvalues; 'V': Eigenvalues in the half-open interval (vl, vu]; 'I': Eigenvalues with indices il through iu.
-     * @param uplo  'U': Upper triangle of A is stored; 'L': Lower triangle is stored.
      * @param dim   The order of the matrix A. dim >= 0.
-     * @param Mat   On entry, the Hermitian matrix A. On exit, it may be overwritten.
-     * @param vl    Lower bound of the interval to search for eigenvalues if range == 'V'.
-     * @param vu    Upper bound of the interval to search for eigenvalues if range == 'V'.
-     * @param il    Index of the smallest eigenvalue to be returned if range == 'I'.
-     * @param iu    Index of the largest eigenvalue to be returned if range == 'I'.
-     * @param m     Output: The total number of found eigenvalues.
-     * @param eigen_val Array to store the computed eigenvalues in ascending order.
-     * @param eigen_vec If not nullptr and jobz == 'V', array to store the computed eigenvectors.
+     * @param lda   The leading dimension of the array Mat. lda >= max(1, dim).
+     * @param Mat   On entry, the Hermitian matrix A. On exit, A is kept.
+     * @param neig  The number of eigenvalues to be found. 0 <= neig <= dim.
+     * @param eigen_val On normal exit, the first \p neig elements contain the selected
+     *                  eigenvalues in ascending order.
+     * @param eigen_vec If eigen_vec is not nullptr, then on exit it contains the
+     *                  orthonormal eigenvectors of the matrix A. The eigenvectors are stored in
+     *                  the columns of eigen_vec, in the same order as the eigenvalues.
      *
      * @note
      * See LAPACK ZHEEVX or CHEEVX documentation for more details.
@@ -87,6 +95,21 @@ struct lapack_heevx {
         T *eigen_vec);
 };
 
+
+// ============================================================================
+// Generalized Hermitian-definite Eigenvalue Problem Solvers
+// ============================================================================
+// The following structures (lapack_hegvd and lapack_hegvx) implement solvers
+// for generalized Hermitian-definite eigenvalue problems of the form:
+//      A * x = lambda * B * x
+// where:
+//   - A is a Hermitian matrix
+//   - B is a Hermitian positive definite matrix
+//   - lambda are the eigenvalues to be computed
+//   - x are the corresponding eigenvectors
+//
+// ============================================================================
+
 template <typename T, typename Device>
 struct lapack_hegvd {
     using Real = typename GetTypeReal<T>::type;

source/source_base/module_container/ATen/kernels/test/lapack_test.cpp

@@ -92,6 +92,9 @@ TYPED_TEST(LapackTest, Potrf) {
     EXPECT_EQ(A, C);
 }
 
+// Test for lapack_heevd and lapack_heevx:
+// Solve a standard eigenvalue problem
+// and check that A*V = V*E
 TYPED_TEST(LapackTest, heevd) {
     using Type = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using Real = typename GetTypeReal<Type>::type;
@@ -179,19 +182,19 @@ TYPED_TEST(LapackTest, heevx) {
 
     // Check the eigenvalues and eigenvectors
     // A * x = lambda * x for the first neig eigenvectors
-    // get A*V
+    // check that A * V = V * E
+    // get A * V
     gemmCalculator(trans, trans, m, n, k, &alpha, A.data<Type>(), m, V.data<Type>(), k, &beta, expected_C1.data<Type>(), m);
-    // get E*V
+    // get V * E
     for (int ii = 0; ii < neig; ii++) {
         axpyCalculator(dim, Alpha.data<Type>() + ii, V.data<Type>() + ii * dim, 1, expected_C2.data<Type>() + ii * dim, 1);
-    }
-    // check that A*V = E*V
-    E = E.to_device<DEVICE_CPU>();
-    V = V.to_device<DEVICE_CPU>();
 
     EXPECT_EQ(expected_C1, expected_C2);
 }
 
+// Test for lapack_hegvd and lapack_hegvx
+// Solve a generalized eigenvalue problem
+// and check that A * v = e * B * v
 TYPED_TEST(LapackTest, hegvd) {
     using Type = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using Real = typename GetTypeReal<Type>::type;
@@ -288,7 +291,8 @@ TYPED_TEST(LapackTest, hegvx) {
 
     // Check the eigenvalues and eigenvectors
     // A * x = lambda * B * x for the first neig eigenvectors
-    // get A*V
+    // check that A * V = E * B * V
+    // get A * V
     gemmCalculator(trans, trans, m, n, k, &alpha, A.data<Type>(), m, V.data<Type>(), k, &beta, expected_C1.data<Type>(), m);
     // get E * B * V
     // where B is 2 * eye(3,3)
@@ -298,10 +302,6 @@ TYPED_TEST(LapackTest, hegvx) {
     for (int ii = 0; ii < neig; ii++) {
         axpyCalculator(dim, Alpha.data<Type>() + ii, C_temp.data<Type>() + ii * dim, 1, expected_C2.data<Type>() + ii * dim, 1);
     }
-    // check that A*V = E*V
-    E = E.to_device<DEVICE_CPU>();
-    V = V.to_device<DEVICE_CPU>();
-
 
     EXPECT_EQ(expected_C1, expected_C2);
 }