Merge pull request #310 from mir-group/gemm

Kokkos multispecies acceleration with type-sorting and matrix-matrix products

.github/workflows/flare.yml

@@ -89,7 +89,7 @@ jobs:
 
     - name: Install LAMMPS
       run: |
-        git clone -b stable_29Sep2021_update2 --depth 1 https://github.com/lammps/lammps.git lammps
+        git clone --depth 1 https://github.com/lammps/lammps.git lammps
         cd lammps_plugins
         ./install.sh $(pwd)/../lammps
         cd ..
@@ -105,6 +105,11 @@ jobs:
         pip install -U codecov pytest pytest-cov pytest_mock
         pip install -r requirements.txt
 
+    - name: Patch ASE
+      run: |
+        ase_file="$(dirname $(python3 -c 'import ase; print(ase.__file__)'))/calculators/lammpsrun.py"
+        sed -i 's/line.startswith(_custom_thermo_mark)/line.strip\(\).startswith\("Step"\)/g' $ase_file
+
     - name: Run tests
       run: |
         export lmp=$(pwd)/lammps/build/lmp

docs/source/related.rst

@@ -1,7 +1,7 @@
 Applications/Gallery
 ====================
 
-If you use FLARE in your research, please let us know. 
+If you use FLARE in your research, please let us know.
 We will list the applications of FLARE here.
 
 - Jonathan Vandermause, Steven B. Torrisi, Simon Batzner, Yu Xie, Lixin Sun, Alexie M. Kolpak, and Boris Kozinsky. **On-the-fly active learning of interpretable Bayesian force fields for atomistic rare events.** npj Computational Materials 6.1 (2020): 1-11. (`arXiv <https://arxiv.org/abs/1904.02042>`_) (`published version <https://www.nature.com/articles/s41524-020-0283-z>`_)
@@ -14,11 +14,11 @@ We will list the applications of FLARE here.
 
 - Jin Soo Lim, Jonathan Vandermause, Matthijs A. Van Spronsen, Albert Musaelian, Yu Xie, Lixin Sun, Christopher R. O'Connor, Tobias Egle, Nicola Molinari, Jacob Florian, Kaining Duanmu, Robert J. Madix, Philippe Sautet, Cynthia M. Friend, and Boris Kozinsky. **Evolution of Metastable Structures at Bimetallic Surfaces from Microscopy and Machine-Learning Molecular Dynamics.** Journal of the American Chemical Society (2020). (`ChemRxiv <https://chemrxiv.org/articles/preprint/Evolution_of_Metastable_Structures_in_Bimetallic_Catalysts_from_Microscopy_and_Machine-Learning_Molecular_Dynamics/11811660/4>`_) (`published version <https://pubs.acs.org/doi/abs/10.1021/jacs.0c06401>`_)
 
-.. figure:: https://www.researchgate.net/profile/Jin-Soo-Lim/publication/339087657/figure/fig3/AS:856047187132427@1581108962681/Frames-of-MD-simulation-showing-the-evolution-of-a-Pd-91-island-on-Ag111-at-500-K-over_W640.jpg 
+.. figure:: https://www.researchgate.net/profile/Jin-Soo-Lim/publication/339087657/figure/fig3/AS:856047187132427@1581108962681/Frames-of-MD-simulation-showing-the-evolution-of-a-Pd-91-island-on-Ag111-at-500-K-over_W640.jpg
     :figwidth: 80 %
     :align: center
 
-    Evolution of Pd/Ag bimetallic surface 
+    Evolution of Pd/Ag bimetallic surface
 
 - Yu Xie, Jonathan Vandermause, Lixin Sun, Andrea Cepellotti, and Boris Kozinsky. **Bayesian force fields from active learning for simulation of inter-dimensional transformation of stanene.** npj Computational Materials 7, no. 1 (2021): 1-10. (`arXiv <https://arxiv.org/abs/2008.11796>`_) (`published version <https://www.nature.com/articles/s41524-021-00510-y>`_)
 
@@ -45,3 +45,18 @@ We will list the applications of FLARE here.
     Melting of Au nanoparticle
 
 - Kai Xu, Lei Yan, and Bingran You. **Bayesian active learning of interatomic force field for molecular dynamics simulation of Pt/Ag(111).** ChemRxiv preprint. (`ChemRxiv <https://chemrxiv.org/engage/chemrxiv/article-details/61c57cf0d6dcc24e3b437905>`_)
+
+- Anders Johansson, Yu Xie, Cameron J. Owen, Jin Soo Lim, Lixin Sun, Jonathan Vandermause, Boris Kozinsky. **Micron-scale heterogeneous catalysis with Bayesian force fields from first principles and active learning** (`arXiv <https://arxiv.org/abs/2204.12573>`_)
+
+.. figure:: https://i.imgur.com/CPCrgWl.png
+    :figwidth: 80 %
+    :align: center
+
+    Strong scaling benchmarks
+
+.. figure:: https://i.imgur.com/p2UeYBo.png
+    :figwidth: 80 %
+    :align: center
+
+    Chemical reaction
+

flare/_version.py

@@ -1 +1 @@
-__version__ = "1.1.2"
+__version__ = "1.2.0"

lammps_plugins/README.md

@@ -6,7 +6,7 @@
 ```
 git clone --depth=1 https://github.com/lammps/lammps.git
 ```
-2. From this directory, run the install script.
+2. From this directory (`flare/lammps_plugins`), run the install script.
 ```
 ./install.sh /path/to/lammps
 ```
@@ -25,7 +25,7 @@ cd build
 cmake ../cmake -DPKG_KOKKOS=ON -DKokkos_ENABLE_CUDA=ON [-DKokkos_ARCH_VOLTA70=on if not auto-detected]
 ```
 
-If you really, really, really want to use the old Makefile system, you should be able to copy the files from `kokkos/` into `/path/to/lammps/src`, do `make yes-kokkos` and otherwise follow the LAMMPS Kokkos instructions.
+*Note: FLARE relies on [KokkosKernels](https://github.com/kokkos/kokkos-kernels) for a performance-portable matrix-matrix product. This will take advantage of MKL/cuBLAS etc. if found. By default, CMake will download KokkosKernels from GitHub and compile it together with LAMMPS. If you need multiple LAMMPS installations with the same Kokkos configuration, you can install Kokkos and KokkosKernels manually and then use the `-DEXTERNAL_KOKKOS=ON` option (with `-DCMAKE_PREFIX_PATH=/path/to/install` as needed).*
 
 ## Basic usage
 Input script:

lammps_plugins/compute_flare_std_atom.cpp

@@ -28,7 +28,7 @@ using namespace LAMMPS_NS;
 
 /* ---------------------------------------------------------------------- */
 
-ComputeFlareStdAtom::ComputeFlareStdAtom(LAMMPS *lmp, int narg, char **arg) : 
+ComputeFlareStdAtom::ComputeFlareStdAtom(LAMMPS *lmp, int narg, char **arg) :
   Compute(lmp, narg, arg),
   stds(nullptr)
 {
@@ -72,7 +72,7 @@ ComputeFlareStdAtom::~ComputeFlareStdAtom() {
 }
 
 /* ----------------------------------------------------------------------
-   init specific to this compute command 
+   init specific to this compute command
 ------------------------------------------------------------------------- */
 
 void ComputeFlareStdAtom::init() {
@@ -81,12 +81,7 @@ void ComputeFlareStdAtom::init() {
 //    error->all(FLERR, "Compute command requires newton pair on");
 
   // Request a full neighbor list.
-  int irequest = neighbor->request(this, instance_me);
-  neighbor->requests[irequest]->pair = 0;
-  neighbor->requests[irequest]->compute = 1;
-  neighbor->requests[irequest]->half = 0;
-  neighbor->requests[irequest]->full = 1;
-  neighbor->requests[irequest]->occasional = 1;
+  neighbor->add_request(this, NeighConst::REQ_FULL | NeighConst::REQ_OCCASIONAL);
 }
 
 void ComputeFlareStdAtom::init_list(int /*id*/, NeighList *ptr)
@@ -166,9 +161,9 @@ void ComputeFlareStdAtom::compute_peratom() {
     }
 
     // Compute covariant descriptors.
-    single_bond_multiple_cutoffs(x, type, jnum, n_inner, i, xtmp, ytmp, ztmp, 
-                                 jlist, basis_function, cutoff_function, 
-                                 n_species, n_max, l_max, radial_hyps, 
+    single_bond_multiple_cutoffs(x, type, jnum, n_inner, i, xtmp, ytmp, ztmp,
+                                 jlist, basis_function, cutoff_function,
+                                 n_species, n_max, l_max, radial_hyps,
                                  cutoff_hyps, single_bond_vals,
                                  single_bond_env_dervs, cutoff_matrix);
 
@@ -200,19 +195,19 @@ void ComputeFlareStdAtom::compute_peratom() {
         }
         cum_types += n_clusters_by_type[s];
       }
-      Eigen::VectorXd L_inv_kv = L_inv_blocks[0] * kernel_vec; 
+      Eigen::VectorXd L_inv_kv = L_inv_blocks[0] * kernel_vec;
 
       double K_self = 1.0;
       double Q_self = sig2 * L_inv_kv.transpose() * L_inv_kv;
- 
+
       variance = K_self - Q_self;
     }
 
     // Compute the normalized variance, it could be negative
     if (variance >= 0.0) {
-      stds[i] = pow(variance, 0.5); 
+      stds[i] = pow(variance, 0.5);
     } else {
-      stds[i] = - pow(abs(variance), 0.5); 
+      stds[i] = - pow(abs(variance), 0.5);
     }
 
   }
@@ -390,7 +385,7 @@ void ComputeFlareStdAtom::read_file(char *filename) {
     cutoff_string_length = strlen(cutoff_string);
   }
 
-  MPI_Bcast(hyperparameters.data(), n_hyps, MPI_DOUBLE, 0, world); 
+  MPI_Bcast(hyperparameters.data(), n_hyps, MPI_DOUBLE, 0, world);
   MPI_Bcast(&n_species, 1, MPI_INT, 0, world);
   MPI_Bcast(&n_max, 1, MPI_INT, 0, world);
   MPI_Bcast(&l_max, 1, MPI_INT, 0, world);
@@ -477,7 +472,7 @@ void ComputeFlareStdAtom::read_L_inverse(char *filename) {
   if (me == 0) {
     fgets(line, MAXLINE, fptr); // skip the first line
 
-    fgets(line, MAXLINE, fptr); // power 
+    fgets(line, MAXLINE, fptr); // power
     sscanf(line, "%i", &power);
 
     fgets(line, MAXLINE, fptr); // hyperparameters
@@ -508,7 +503,7 @@ void ComputeFlareStdAtom::read_L_inverse(char *filename) {
     cutoff_string_length = strlen(cutoff_string);
   }
 
-  MPI_Bcast(hyperparameters.data(), n_hyps, MPI_DOUBLE, 0, world); 
+  MPI_Bcast(hyperparameters.data(), n_hyps, MPI_DOUBLE, 0, world);
   MPI_Bcast(&n_species, 1, MPI_INT, 0, world);
   MPI_Bcast(&n_max, 1, MPI_INT, 0, world);
   MPI_Bcast(&l_max, 1, MPI_INT, 0, world);
@@ -525,7 +520,7 @@ void ComputeFlareStdAtom::read_L_inverse(char *filename) {
   // Parse number of sparse envs
   if (me == 0) {
     fgets(line, MAXLINE, fptr);
-    sscanf(line, "%i", &n_clusters); 
+    sscanf(line, "%i", &n_clusters);
   }
   MPI_Bcast(&n_clusters, 1, MPI_INT, 0, world);
 
@@ -622,14 +617,14 @@ void ComputeFlareStdAtom::read_sparse_descriptors(char *filename) {
       n_clusters_by_type[s] = n_clst_by_type;
 
       // Check the relationship between the power spectrum and beta.
-      int sparse_desc_size = n_clst_by_type * n_descriptors; 
-  
+      int sparse_desc_size = n_clst_by_type * n_descriptors;
+
       // Parse the beta vectors.
       memory->create(beta, sparse_desc_size, "compute:sparse_desc");
       if (me == 0)
         grab(fptr, sparse_desc_size, beta);
       MPI_Bcast(beta, sparse_desc_size, MPI_DOUBLE, 0, world);
-  
+
       // Fill in the beta matrix.
       int count = 0;
       Eigen::MatrixXd sparse_desc = Eigen::MatrixXd::Zero(n_clst_by_type, n_descriptors);

lammps_plugins/install.sh

@@ -39,4 +39,21 @@ target_sources(lammps PRIVATE
     ${LAMMPS_SOURCE_DIR}/radial.cpp
     ${LAMMPS_SOURCE_DIR}/y_grad.cpp
 )
+
+if(PKG_KOKKOS)
+    set(KokkosKernels_ADD_DEFAULT_ETI OFF CACHE BOOL "faster build")
+    set(KokkosKernels_ENABLED_COMPONENTS BLAS CACHE STRING "faster build")
+
+    if(EXTERNAL_KOKKOS)
+        find_package(KokkosKernels REQUIRED)
+    else()
+        include(FetchContent)
+        FetchContent_Declare(
+            kokkoskernels
+            GIT_REPOSITORY https://github.com/kokkos/kokkos-kernels.git
+        )
+        FetchContent_MakeAvailable(kokkoskernels)
+    endif()
+    target_link_libraries(lammps PUBLIC Kokkos::kokkoskernels)
+endif()
 ' >> $lammps/cmake/CMakeLists.txt