Added tiled DGEMM.

TREX-CoE · Mar 29, 2021 · b52af8f · b52af8f
1 parent fce1929
commit b52af8f
Showing 1 changed file with 49 additions and 5 deletions.
diff --git a/tiling.irp.f b/tiling.irp.f
@@ -6,9 +6,9 @@
 tile_size = 16
 END_PROVIDER
 
- BEGIN_PROVIDER [double precision, rescale_een_e_tiled, (tile_size, tile_size, ntiles_nelec, ntiles_nelec, 0:ncord)]
-&BEGIN_PROVIDER [double precision, rescale_een_n_tiled, (tile_size, tile_size, ntiles_nelec, ntiles_nnuc, 0:ncord)]
-&BEGIN_PROVIDER [double precision, rescale_een_e_deriv_e_tiled, (tile_size, nelec_16, 4, tile_size, nelec_16, 0:ncord)]
+ BEGIN_PROVIDER [double precision, rescale_een_e_tiled, (tile_size, tile_size, 0:ntiles_nelec, 0:ntiles_nelec, 0:ncord)]
+&BEGIN_PROVIDER [double precision, rescale_een_n_tiled, (tile_size, tile_size, 0:ncord, 0:ntiles_nelec, 0:ntiles_nnuc)]
+&BEGIN_PROVIDER [double precision, rescale_een_e_deriv_e_tiled, (tile_size, 4, tile_size, 0:ncord, 0:ntiles_nelec, 0:ntiles_nelec)]
  implicit none
  BEGIN_DOC
  ! R = exp(-kappa r) for electron-electron for $J_{een}$
@@ -39,7 +39,7 @@
            idxi = i*tile_size + ii
            do aa = 1, tile_size
              idxa = a*tile_size + aa
-             rescale_een_n_tiled(ii, aa, i, a, l) =  rescale_een_n(i, a, l)
+             rescale_een_n_tiled(ii, aa, l, i, a) =  rescale_een_n(i, a, l)
            enddo
          enddo
        enddo
@@ -55,13 +55,57 @@
        do ii = 1, tile_size
           idxi = i*tile_size + ii
           do k = 1, 4
-             rescale_een_e_deriv_e_tiled(ii,i,k,jj,j,l) = rescale_een_e_deriv_e(i, k, j, l) 
+             rescale_een_e_deriv_e_tiled(ii,k,jj,l,i,j) = rescale_een_e_deriv_e(i, k, j, l) 
           enddo
        enddo
        enddo
     enddo
     enddo
 enddo
 
+END_PROVIDER
+
+ BEGIN_PROVIDER [ double precision,  tmp_c_tiled, (tile_size, tile_size,0:ncord, 0:ntiles_nelec, 0:ntiles_nnuc,0:ncord-1) ]
+&BEGIN_PROVIDER [ double precision, dtmp_c_tiled, (tile_size, 4,tile_size,0:ncord, 0:ntiles_nelec, 0:ntiles_nnuc,0:ncord-1) ]
+ implicit none
+ BEGIN_DOC
+ ! Calculate the intermediate buffers
+ ! tmp_c:
+ ! r_{ij}^k . R_{ja}^l -> tmp_c_{ia}^{kl}
+ !
+ ! dtmp_c:
+ ! dr_{ij}^k . R_{ja}^l -> dtmp_c_{ia}^{kl}
+ END_DOC
+ integer :: k, i, j, a
+
+ ! r_{ij}^k . R_{ja}^l -> tmp_c_{ia}^{kl}
+ do k=0,ncord-1
+   do i = 0, ntiles_nelec
+     do j = 0, ntiles_nelec
+       do a = 0, ntiles_nnuc
+   call dgemm('N','N', tile_size, tile_size*(ncord+1), tile_size, 1.d0,           &
+       rescale_een_e_tiled(1,1,i,j,k), size(rescale_een_e_tiled,1),                  &
+       rescale_een_n_tiled(1,1,0,i,a), size(rescale_een_n_tiled,1), 0.d0,            &
+       tmp_c_tiled(1,1,0,i,a,k), size(tmp_c_tiled,1))
+       enddo
+     enddo
+   enddo
+ enddo
+
+ ! dr_{ij}^k . R_{ja}^l -> dtmp_c_{ia}^{kl}
+ do k=0,ncord-1
+   do i = 0, ntiles_nelec
+     do j = 0, ntiles_nelec
+       do a = 0, ntiles_nnuc
+   call dgemm('N','N', 4*tile_size, tile_size*(ncord+1), nelec, 1.d0,         &
+       rescale_een_e_deriv_e_tiled(1,1,1,i,a,k), 4*size(rescale_een_e_deriv_e_tiled,1),&
+       rescale_een_n_tiled(1,1,0,i,a), size(rescale_een_n_tiled,1), 0.d0,            &
+       dtmp_c_tiled(1,1,1,0,i,a,k), 4*size(dtmp_c_tiled,1))
+       enddo
+     enddo
+   enddo
+ enddo
+
+
 END_PROVIDER