Merge branch 'lanczos-modernize' into 'master'

Lanczos modernize

See merge request SLai/ChASE!31

Author: brunowu

ChASE-MPI/chase_mpi.hpp

@@ -554,74 +554,13 @@ class ChaseMpi : public chase::Chase<T>
     void Lanczos(std::size_t m, Base<T>* upperb) override
     {
         // todo
-        std::size_t n = N_;
         Base<T>* d = new Base<T>[m]();
         Base<T>* e = new Base<T>[m]();
 
         int idx_ = -1;
         Base<T> real_beta;
+        dla_->Lanczos(m, idx_, d, e, &real_beta);
 
-        T alpha = T(1.0);
-        T beta = T(0.0);
-        T One = T(1.0);
-        T Zero = T(0.0);
-
-        T* V1;
-        T* V2;
-        std::size_t ld;
-        T* v0;
-        T* v1;
-        T* w;
-#ifdef USE_NSIGHT
-        nvtxRangePushA("getLanczosBuffer2");
-#endif
-        dla_->getLanczosBuffer2(&v0, &v1, &w);
-#ifdef USE_NSIGHT
-        nvtxRangePop();
-#endif
-
-#ifdef USE_NSIGHT
-        nvtxRangePushA("Lanczos: loop");
-#endif
-        // ENSURE that v1 has one norm
-        Base<T> real_alpha = dla_->nrm2(n, v1, 1);
-        alpha = T(1 / real_alpha);
-        dla_->scal(n, &alpha, v1, 1);
-
-        for (std::size_t k = 0; k < m; k = k + 1)
-        {
-            dla_->applyVec(v1, w);
-
-            alpha = dla_->dot(n, v1, 1, w, 1);
-
-            alpha = -alpha;
-            dla_->axpy(n, &alpha, v1, 1, w, 1);
-            alpha = -alpha;
-
-            d[k] = std::real(alpha);
-
-            if (k == m - 1)
-                break;
-
-            beta = T(-real_beta);
-            dla_->axpy(n, &beta, v0, 1, w, 1);
-            beta = -beta;
-
-            real_beta = dla_->nrm2(n, w, 1);
-
-            beta = T(1.0 / real_beta);
-
-            dla_->scal(n, &beta, w, 1);
-
-            e[k] = real_beta;
-
-            std::swap(v1, v0);
-            std::swap(v1, w);
-        }
-#ifdef USE_NSIGHT
-        nvtxRangePop();
-#endif
-        dla_->preApplication(v1, 0, 1);
 #ifdef USE_NSIGHT
         nvtxRangePushA("Stemr");
 #endif
@@ -663,71 +602,8 @@ class ChaseMpi : public chase::Chase<T>
         int idx_ = static_cast<int>(idx);
         Base<T> real_beta;
 
-        std::size_t n = N_;
+        dla_->Lanczos(m, idx_, d, e, &real_beta);
 
-        T alpha = T(1.0);
-        T beta = T(0.0);
-        T One = T(1.0);
-        T Zero = T(0.0);
-
-        T* V1;
-        T* V2;
-        std::size_t ld;
-        T* v0;
-        T* v1;
-        T* w;
-#ifdef USE_NSIGHT
-        nvtxRangePushA("getLanczosBuffer");
-#endif
-        dla_->getLanczosBuffer(&V1, &V2, &ld, &v0, &v1, &w);
-#ifdef USE_NSIGHT
-        nvtxRangePop();
-        nvtxRangePushA("C2V");
-#endif
-        dla_->C2V(V2, idx, v1, 0, 1);
-#ifdef USE_NSIGHT
-        nvtxRangePop();
-#endif
-        // ENSURE that v1 has one norm
-#ifdef USE_NSIGHT
-        nvtxRangePushA("Lanczos: loop");
-#endif
-        Base<T> real_alpha = dla_->nrm2(n, v1, 1);
-        alpha = T(1 / real_alpha);
-        dla_->scal(n, &alpha, v1, 1);
-        for (std::size_t k = 0; k < m; k = k + 1)
-        {
-            dla_->V2C(v1, 0, V1, k, 1);
-            dla_->applyVec(v1, w);
-            alpha = dla_->dot(n, v1, 1, w, 1);
-            alpha = -alpha;
-            dla_->axpy(n, &alpha, v1, 1, w, 1);
-            alpha = -alpha;
-
-            d[k] = std::real(alpha);
-
-            if (k == m - 1)
-                break;
-
-            beta = T(-real_beta);
-            dla_->axpy(n, &beta, v0, 1, w, 1);
-            beta = -beta;
-
-            real_beta = dla_->nrm2(n, w, 1);
-
-            beta = T(1.0 / real_beta);
-
-            dla_->scal(n, &beta, w, 1);
-
-            e[k] = real_beta;
-
-            std::swap(v1, v0);
-            std::swap(v1, w);
-        }
-#ifdef USE_NSIGHT
-        nvtxRangePop();
-#endif
-        dla_->preApplication(v1, 0, 1);
 #ifdef USE_NSIGHT
         nvtxRangePushA("Stemr");
 #endif
@@ -738,7 +614,7 @@ class ChaseMpi : public chase::Chase<T>
         int* isuppz = new int[2 * m];
         t_stemr(LAPACK_COL_MAJOR, 'V', 'A', m, d, e, ul, ll, vl, vu,
                 &notneeded_m, ritzv, ritzV, m, m, isuppz, &tryrac);
-        *upperb = std::max(std::abs(ritzv[0]), std::abs(ritzv[m - 1])) +
+        *upperb = std::max(std::abs(ritzv[0]), std::abs(ritzv[m - 1])) + 
                   std::abs(real_beta);
 #ifdef USE_NSIGHT
         nvtxRangePop();
@@ -775,26 +651,6 @@ class ChaseMpi : public chase::Chase<T>
     }
 #endif
 
-    // if distributed ChASE is used, collecting the distributed ritz vectors
-    // into V which is redundant across all MPI ranks. if non-distributed ChASE
-    // is used, copying Ritz vectors directly to V
-    //! When distributed ChASE is used, this member function collects the
-    //! partially distributed Ritz vectors into a redundant vectors `V` on all
-    //! MPI procs.
-    //! @param V: the buffer of size `N_xnev_` which stores the collected
-    //! redundant Ritz vectors.
-    void collectRitzVecs(T* V)
-    {
-#ifdef USE_NSIGHT
-        nvtxRangePushA("collectRitzVecs");
-#endif
-        T* Vv = matrices_.get_V1();
-        dla_->C2V(Vv, 0, V, 0, nev_);
-#ifdef USE_NSIGHT
-        nvtxRangePop();
-#endif
-    }
-
     //! \return `H_`: A pointer to the memory allocated to store (local part if
     //! applicable) of matrix `A`.
     T* GetMatrixPtr() { return matrices_.get_H(); }

ChASE-MPI/chase_mpidla_interface.hpp

@@ -129,42 +129,7 @@ class ChaseMpiDLAInterface
     */
     virtual void asynCxHGatherC(std::size_t locked, std::size_t block,
                                 bool isCcopied = false) = 0;
-
-    //! Copy from buffer rectangular matrix `v1` to `v2`.
-    //! For the implementation of distributed-memory ChASE, this operation
-    //! performs a `copy` from a matrix redundantly distributed across all MPI
-    //! procs to a matrix distributed within each column communicator and
-    //! redundant among different column communicators. This operation is
-    //! reciprocal to V2C().
-    /*!
-     *  @param v1: the buffer to copy from
-     *  @param off1: the offset for the starting column index of `v1` to copy
-     * from
-     *  @param v2: the buffer to copy to
-     *  @param off2: the offset for the starting column index of `v2` to copy to
-     *  @param block: number of columns to copy from `v1` to `v2`
-     */
-    virtual void C2V(T* v1, std::size_t off1, T* v2, std::size_t off2,
-                     std::size_t block) = 0;
-    //! Copy from buffer rectangular matrix `v1` to `v2`.
-    //! For the implementation of distributed-memory ChASE, this operation
-    //! performs a `copy` from a matrix distributed within each column
-    //! communicator and redundant among different column communicators to a
-    //! matrix redundantly distributed across all MPI procs. This operation is
-    //! reciprocal to C2V(). It requires the `dim_[0]` MPI broadcasting
-    //! operations, in which `dim_[0]` is the size of each MPI column
-    //! communicatior.
-    /*!
-     *  @param v1: the buffer to copy from
-     *  @param off1: the offset for the starting column index of `v1` to copy
-     * from
-     *  @param v2: the buffer to copy to
-     *  @param off2: the offset for the starting column index of `v2` to copy to
-     *  @param block: number of columns to copy from `v1` to `v2`
-     */
-    virtual void V2C(T* v1, std::size_t off1, T* v2, std::size_t off2,
-                     std::size_t block) = 0;
-
+    
     //! Swap the columns indexing `i` and `j` in a rectangular matrix
     //! The operated matrices maybe different in different implementations
     /*!
@@ -371,25 +336,13 @@ class ChaseMpiDLAInterface
     //! Cholesky QR factorization on the rectangular matrix `V1`.
     //!  @param locked: number of converged eigenvectors.
     virtual void cholQR(std::size_t locked, Base<T> cond) = 0;
-    //! Return the required buffers of Lanczos which are allocated within each
-    //! individual implementation of DLA. This operation is required, since
-    //! Lanczos algorithm is implemented in ChaseMpi class, which has no direct
-    //! access to these buffers. Depending on the implementation and targeting
-    //! architectures, these buffers can be on CPU or GPUs.
-    //! **This function will be removed in short future in which all the buffers
-    //! will be allocated within ChaseMpiMatrices**.
-    virtual void getLanczosBuffer(T** V1, T** V2, std::size_t* ld, T** v0,
-                                  T** v1, T** w) = 0;
-    //! Return the required buffers of Lanczos which are allocated within each
-    //! individual implementation of DLA. This operation is required, since
-    //! Lanczos algorithm is implemented in ChaseMpi class, which has no direct
-    //! access to these buffers. Depending on the implementation and targeting
-    //! architectures, these buffers can be on CPU or GPUs.
-    //! **This function will be removed in short future in which all the buffers
-    //! will be allocated within ChaseMpiMatrices**.
-    virtual void getLanczosBuffer2(T** v0, T** v1, T** w) = 0;
     //! Lanczos DOS to estimate the \mu_{nev+nex} for ChASE
     virtual void LanczosDos(std::size_t idx, std::size_t m, T* ritzVc) = 0;
+
+    virtual void Lanczos(std::size_t M, int idx, Base<T>* d, Base<T>* e, Base<T> *r_beta) = 0;
+
+    virtual void B2C(T* B, std::size_t off1, T* C, std::size_t off2, std::size_t block) = 0;
+
 };
 } // namespace mpi
 } // namespace chase