restructured Lanczos for ChaseMpiMultiGPU

Author: brunowu

ChASE-MPI/impl/chase_mpidla.hpp

@@ -1573,8 +1573,10 @@ class ChaseMpiDLA : public ChaseMpiDLAInterface<T>
 #ifdef USE_NSIGHT
         nvtxRangePushA("Lanczos: loop");
 #endif
-        Base<T> real_alpha = t_norm_p2(m_, v1_);
-        MPI_Allreduce(MPI_IN_PLACE, &real_alpha, 1, getMPI_Type<Base<T>>(),
+        //Base<T> real_alpha = t_norm_p2(m_, v1_);
+        Base<T> real_alpha = t_nrm2(m_, v1_, 1);
+	real_alpha = std::pow(real_alpha,2);
+	MPI_Allreduce(MPI_IN_PLACE, &real_alpha, 1, getMPI_Type<Base<T>>(),
                     MPI_SUM, col_comm_);
         real_alpha = std::sqrt(real_alpha);
         alpha = T(1 / real_alpha);
@@ -1603,8 +1605,10 @@ class ChaseMpiDLA : public ChaseMpiDLAInterface<T>
             t_axpy(m_, &beta, v0_, 1, v2_, 1);
             beta = -beta;
 
-            real_beta = t_norm_p2(m_, v2_);
-            MPI_Allreduce(MPI_IN_PLACE, &real_beta, 1, getMPI_Type<Base<T>>(),
+            //real_beta = t_norm_p2(m_, v2_);
+            real_beta = t_nrm2(m_, v2_, 1);
+	    real_beta = std::pow(real_beta, 2);
+	    MPI_Allreduce(MPI_IN_PLACE, &real_beta, 1, getMPI_Type<Base<T>>(),
                           MPI_SUM, col_comm_);
             real_beta = std::sqrt(real_beta);
 

ChASE-MPI/impl/chase_mpidla_mgpu.hpp

@@ -259,6 +259,8 @@ class ChaseMpiDLAMultiGPU : public ChaseMpiDLAInterface<T>
             cudaMalloc((void**)&d_ritz_, sizeof(Base<T>) * (nev_ + nex_)));
         cuda_exec(
             cudaMalloc((void**)&states_, sizeof(curandStatePhilox4_32_10_t) * (256 * 32)));
+        cuda_exec(cudaMalloc((void**)&d_v_, sizeof(T) * m_));
+        cuda_exec(cudaMalloc((void**)&d_w_, sizeof(T) * n_));
 
         cublasCreate(&cublasH_);
         cusolverDnCreate(&cusolverH_);
@@ -486,7 +488,21 @@ class ChaseMpiDLAMultiGPU : public ChaseMpiDLAInterface<T>
     //! - All required operations for this function has been done in for
     //! ChaseMpiDLA::applyVec().
     //! - This function contains nothing in this class.
-    void applyVec(T* B, T* C) override {}
+    void applyVec(T* v, T* w) override 
+    {
+        T alpha = T(1.0);
+        T beta = T(0.0);
+	    
+        cuda_exec(cudaMemcpy(d_v_, v, m_ * sizeof(T), cudaMemcpyHostToDevice));
+	std::size_t k = 1;
+	cublas_status_ = cublasTgemm(
+            cublasH_, CUBLAS_OP_C, CUBLAS_OP_N, n_, k, m_, &alpha, d_H_, m_,
+            d_v_, m_, &beta, d_w_, n_);
+        assert(cublas_status_ == CUBLAS_STATUS_SUCCESS);
+
+        cuda_exec(cudaMemcpy(w, d_w_, n_ * sizeof(T), cudaMemcpyDeviceToHost));
+
+    }
     int get_nprocs() const override { return matrix_properties_->get_nprocs(); }
     void Start() override {}
     void End() override {}
@@ -786,6 +802,9 @@ class ChaseMpiDLAMultiGPU : public ChaseMpiDLAInterface<T>
     T* d_work_ =
         NULL; //!< a pointer to a local buffer on GPU, which is reserved for the
               //!< extra buffer required for any cuSOLVER routines
+    T *d_v_;
+    T *d_w_;
+
     std::size_t pitchB;  //!< pitch for `B_` and `d_B_`
     std::size_t pitchB2; //!< pitch for `B2_` and `d_B2_`
     std::size_t pitchC;  //!< pitch for `C_` and `d_C_`