grid: Use dgemm for decontraction of density matrix blocks

cp2k · Oct 9, 2020 · f85dabd · f85dabd
1 parent 0cf0f7d
commit f85dabd
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 34 deletions.
diff --git a/src/grid/Makefile b/src/grid/Makefile
@@ -24,9 +24,9 @@ ALL_OBJECTS := grid_collocate_replay.o \
 	cd $(dir $<); $(CC) -c $(CFLAGS) $(notdir $<)
 
 grid_collocate_miniapp.x: grid_collocate_miniapp.o $(ALL_OBJECTS)
-	$(CC) $(CFLAGS) -o $@ $^ -lm
+	$(CC) $(CFLAGS) -o $@ $^ -lm -lblas
 
 grid_collocate_unittest.x: grid_collocate_unittest.o $(ALL_OBJECTS)
-	$(CC) $(CFLAGS) -o $@ $^ -lm
+	$(CC) $(CFLAGS) -o $@ $^ -lm -lblas
 
 #EOF
diff --git a/src/grid/common/grid_common.h b/src/grid/common/grid_common.h
@@ -66,6 +66,15 @@ static inline int imax(int x, int y) { return (x > y ? x : y); }
  ******************************************************************************/
 static inline int modulo(int a, int m) { return ((a % m + m) % m); }
 
+/*******************************************************************************
+ * \brief Prototype for BLAS dgemm.
+ * \author Ole Schuett
+ ******************************************************************************/
+void dgemm_(const char *transa, const char *transb, const int *m, const int *n,
+            const int *k, const double *alpha, const double *a, const int *lda,
+            const double *b, const int *ldb, const double *beta, double *c,
+            const int *ldc);
+
 #endif // GRID_COMMON_H
 
 // EOF
diff --git a/src/grid/ref/grid_ref_task_list.c b/src/grid/ref/grid_ref_task_list.c
@@ -190,42 +190,24 @@ static void collocate_one_grid_level(
             task_list->block_offsets[block_num]; // zero based
         const double *block = &task_list->blocks_buffer[block_offset];
 
-        // Copy sub block for current sets and transpose it if needed.
-        double subblock[nsgf_setb][nsgf_seta];
+        // Decontract the sub block into pab.
+        double work[nsgf_setb * ncoa];
+        const double zero = 0.0, one = 1.0;
         if (iatom <= jatom) {
-          for (int i = 0; i < nsgf_setb; i++)
-            for (int j = 0; j < nsgf_seta; j++)
-              subblock[i][j] = block[(i + sgfb) * nsgfa + j + sgfa];
+          // work = MATMUL(ibasis->sphi, subblock)
+          dgemm_("N", "N", &ncoa, &nsgf_setb, &nsgf_seta, &one,
+                 &ibasis->sphi[sgfa * maxcoa], &maxcoa,
+                 &block[sgfb * nsgfa + sgfa], &nsgfa, &zero, work, &ncoa);
         } else {
-          for (int i = 0; i < nsgf_setb; i++)
-            for (int j = 0; j < nsgf_seta; j++)
-              subblock[i][j] =
-                  block[(j + sgfa) * nsgfb + i + sgfb]; // transposed
+          // work = MATMUL(ibasis->sphi, TRANSPOSE(subblock))
+          dgemm_("N", "T", &ncoa, &nsgf_setb, &nsgf_seta, &one,
+                 &ibasis->sphi[sgfa * maxcoa], &maxcoa,
+                 &block[sgfa * nsgfb + sgfb], &nsgfb, &zero, work, &ncoa);
         }
+        // double pab[ncob][ncoa] = MATMUL(work, TRANSPOSE(jbasis->sphi))
+        dgemm_("N", "T", &ncoa, &ncob, &nsgf_setb, &one, work, &ncoa,
+               &jbasis->sphi[sgfb * maxcob], &maxcob, &zero, pab, &ncoa);
 
-        // work = MATMUL(ibasis->sphi, subblock)
-        double work[nsgf_setb][ncoa];
-        memset(work, 0, sizeof(double) * nsgf_setb * ncoa);
-        for (int i = 0; i < nsgf_setb; i++) {
-          for (int j = 0; j < ncoa; j++) {
-            for (int k = 0; k < nsgf_seta; k++) {
-              work[i][j] +=
-                  subblock[i][k] * ibasis->sphi[(k + sgfa) * maxcoa + j];
-            }
-          }
-        }
-
-        // double pab[ncob][ncoa]
-        // pab = MATMUL(work, TRANSPOSE(jbasis->sphi))
-        memset(pab, 0, sizeof(double) * ncob * ncoa);
-        for (int i = 0; i < ncob; i++) {
-          for (int j = 0; j < ncoa; j++) {
-            for (int k = 0; k < nsgf_setb; k++) {
-              pab[i * ncoa + j] +=
-                  work[k][j] * jbasis->sphi[(k + sgfb) * maxcob + i];
-            }
-          }
-        }
       } // end of block loading
 
       const double zeta = ibasis->zet[iset * ibasis->maxpgf + ipgf];