replace einsum by gemm in diag_hsub

Author: Cstandardlib

source/module_hsolver/diago_bpcg.cpp

@@ -106,12 +106,12 @@ void DiagoBPCG<T, Device>::orth_cholesky(
                 this->n_band,      //m
                 this->n_band,      //n
                 this->n_dim,       //k
-                this->one,
+                this->one,         //1.0
                 psi_out.data<T>(),
                 this->n_basis,     //lda
                 psi_out.data<T>(),
                 this->n_basis,     //ldb
-                this->zero,
+                this->zero,        //0.0
                 hsub_out.data<T>(),
                 this->n_band);     //ldc
 
@@ -175,12 +175,12 @@ void DiagoBPCG<T, Device>::orth_projection(
                 this->n_band,      //m
                 this->n_band,      //n
                 this->n_dim,       //k
-                this->one,
+                this->one,         //1.0
                 psi_in.data<T>(),
                 this->n_basis,     //lda
                 grad_out.data<T>(),
                 this->n_basis,     //ldb
-                this->zero,
+                this->zero,        //0.0
                 hsub_in.data<T>(),
                 this->n_band);     //ldc
 
@@ -189,21 +189,21 @@ void DiagoBPCG<T, Device>::orth_projection(
         /*conj_x=*/false, /*conj_y=*/false, /*alpha=*/-1.0, /*beta=*/1.0, /*Tensor out=*/&grad_out);
     // grad_out = ct::op::einsum("ij,jk->ik", hsub_in, psi_in, option);
 
-    // grad_out(n_basis x n_band) = psi_in(n_basis x n_band) * hsub_in(n_band x n_band)
+    // grad_out(n_basis x n_band) = 1.0 * grad_out(n_basis x n_band) - psi_in(n_basis x n_band) * hsub_in(n_band x n_band)
     gemm_op()(this->ctx,
                 'N',
                 'N',
-                this->n_dim,    //m
+                this->n_dim,      //m
                 this->n_band,     //n
                 this->n_band,     //k
-                this->neg_one,
+                this->neg_one,    //-1.0
                 psi_in.data<T>(),
                 this->n_basis,    //lda
                 hsub_in.data<T>(),
                 this->n_band,     //ldb
-                this->one,
+                this->one,        //1.0
                 grad_out.data<T>(),
-                this->n_basis);  //ldc
+                this->n_basis);   //ldc
 
     return;
 }
@@ -224,15 +224,15 @@ void DiagoBPCG<T, Device>::rotate_wf(
     gemm_op()(this->ctx,
                 'N',
                 'N',
-                this->n_dim,      //m
-                this->n_band,      //n
+                this->n_dim,        //m
+                this->n_band,       //n
                 this->n_band,       //k
-                this->one,
+                this->one,          //1.0
                 psi_out.data<T>(),
-                this->n_basis,     //lda
+                this->n_basis,      //lda
                 hsub_in.data<T>(),
-                this->n_band,     //ldb
-                this->zero,
+                this->n_band,       //ldb
+                this->zero,         //0.0
                 workspace_in.data<T>(),
                 this->n_basis);     //ldc
 
@@ -263,7 +263,23 @@ void DiagoBPCG<T, Device>::diag_hsub(
     // it controls the ops to use the corresponding device to calculate results
     ct::EinsumOption option(
         /*conj_x=*/false, /*conj_y=*/true, /*alpha=*/1.0, /*beta=*/0.0, /*Tensor out=*/&hsub_out);
-    hsub_out = ct::op::einsum("ij,kj->ik", psi_in, hpsi_in, option);
+    // hsub_out = ct::op::einsum("ij,kj->ik", psi_in, hpsi_in, option);
+
+    // gemm: hsub_out(n_band x n_band) = hpsi_in^T(n_band x n_basis) * psi_in(n_basis x n_band)
+    gemm_op()(this->ctx,
+                'C',
+                'N',
+                this->n_band,       //m
+                this->n_band,       //n
+                this->n_dim,        //k
+                this->one,          //1.0
+                hpsi_in.data<T>(),
+                this->n_basis,      //lda
+                psi_in.data<T>(),
+                this->n_basis,      //ldb
+                this->zero,         //0.0
+                hsub_out.data<T>(),
+                this->n_band);      //ldc
 
     ct::kernels::lapack_dnevd<T, ct_Device>()('V', 'U', hsub_out.data<T>(), this->n_band, eigenvalue_out.data<Real>());