Use error->one instead of error->all

Otherwise we can have a deadlock if only one MPI rank has the
error

Author: Luthaf

src/KOKKOS/metatomic_system_kokkos.cpp

@@ -47,13 +47,18 @@ MetatomicSystemAdaptorKokkos<DeviceType>::MetatomicSystemAdaptorKokkos(LAMMPS *l
     this->strain = torch::eye(3, tensor_options);
 }
 
+#include <iostream>
+#include "comm.h"
+
 template<class DeviceType>
 void MetatomicSystemAdaptorKokkos<DeviceType>::setup_neighbors_remap_kk(metatomic_torch::System& system, NeighListKokkos<DeviceType>* list) {
     auto _ = MetatomicTimer("converting kokkos neighbors with ghosts remapping");
     auto dtype = system->positions().scalar_type();
 
     auto total_n_atoms = atomKK->nlocal + atomKK->nghost;
 
+    std::cout << "rank = " << comm->me << " nlocal = " << atomKK->nlocal << " nghost = " << atomKK->nghost << std::endl;
+
     {
         auto _ = MetatomicTimer("identifying ghosts and real atoms");
         /*-------------- this will be done on CPU for now ------------------------*/
@@ -360,7 +365,7 @@ metatomic_torch::System MetatomicSystemAdaptorKokkos<DeviceType>::system_from_lm
     // While Metatomic models can support mixed PBC settings, we currently
     // assume that the system is fully periodic and we throw an error otherwise
     if (!domain->xperiodic || !domain->yperiodic || !domain->zperiodic) {
-        error->all(FLERR, "metatomic/kk currently requires a fully periodic system");
+        error->one(FLERR, "metatomic/kk currently requires a fully periodic system");
     }
     auto pbc = torch::tensor(
         {domain->xperiodic, domain->yperiodic, domain->zperiodic},
@@ -379,7 +384,7 @@ metatomic_torch::System MetatomicSystemAdaptorKokkos<DeviceType>::system_from_lm
         assert(kk_list != nullptr);
         this->setup_neighbors_remap_kk(system, kk_list);
     } else {
-        error->all(FLERR, "the kokkos version of metatomic requires remap_pairs to be true");
+        error->one(FLERR, "the kokkos version of metatomic requires remap_pairs to be true");
     }
 
     return system;

src/KOKKOS/pair_metatomic_kokkos.cpp

@@ -99,7 +99,7 @@ void PairMetatomicKokkos<DeviceType>::pick_device(torch::Device& device, const c
         auto requested_str = std::string(requested);
         std::transform(requested_str.begin(), requested_str.end(), requested_str.begin(), ::tolower);
         if (c10::DeviceTypeName(device.type(), /*lower_case=*/true) != requested_str) {
-            error->all(FLERR,
+            error->one(FLERR,
                 "requested device '{}' does not match the device being used by kokkos '{}', "
                 "use the non-kokkos version of this pair style to use a different "
                 "device for the model and LAMMPS",

src/ML-METATOMIC/metatomic_system.cpp

@@ -45,13 +45,13 @@ MetatomicSystemAdaptor::~MetatomicSystemAdaptor() {}
 
 void MetatomicSystemAdaptor::add_nl_request(double cutoff, metatomic_torch::NeighborListOptions request) {
     if (cutoff > options_.interaction_range) {
-        error->all(FLERR,
+        error->one(FLERR,
             "Invalid metatomic model: one of the requested neighbor lists "
             "has a cutoff ({}) larger than the model interaction range ({})",
             cutoff, options_.interaction_range
         );
     } else if (cutoff < 0 || !std::isfinite(cutoff)) {
-        error->all(FLERR,
+        error->one(FLERR,
             "model requested an invalid cutoff for neighbors list: {} "
             "(cutoff in model units is {})",
             cutoff, request->cutoff()
@@ -278,7 +278,7 @@ void MetatomicSystemAdaptor::setup_neighbors_remap(metatomic_torch::System& syst
                             });
                         } else {
                             // should be unreachable
-                            error->all(FLERR, "invalid dtype, this is a bug");
+                            error->one(FLERR, "invalid dtype, this is a bug");
                         }
                     }
                 }
@@ -317,7 +317,7 @@ void MetatomicSystemAdaptor::setup_neighbors_remap(metatomic_torch::System& syst
             );
         } else {
             // should be unreachable
-            error->all(FLERR, "invalid dtype, this is a bug");
+            error->one(FLERR, "invalid dtype, this is a bug");
         }
 
         {
@@ -408,7 +408,7 @@ void MetatomicSystemAdaptor::setup_neighbors_no_remap(metatomic_torch::System& s
                         });
                     } else {
                         // should be unreachable
-                        error->all(FLERR, "invalid dtype, this is a bug");
+                        error->one(FLERR, "invalid dtype, this is a bug");
                     }
                 }
             }
@@ -446,7 +446,7 @@ void MetatomicSystemAdaptor::setup_neighbors_no_remap(metatomic_torch::System& s
             );
         } else {
             // should be unreachable
-            error->all(FLERR, "invalid dtype, this is a bug");
+            error->one(FLERR, "invalid dtype, this is a bug");
         }
 
         {
@@ -526,7 +526,7 @@ metatomic_torch::System MetatomicSystemAdaptor::system_from_lmp(
     // While metatomic models can support mixed PBC settings, we currently
     // assume that the system is fully periodic and we throw an error otherwise
     if (!domain->xperiodic || !domain->yperiodic || !domain->zperiodic) {
-        error->all(FLERR, "pair_style metatomic requires a fully periodic system");
+        error->one(FLERR, "pair_style metatomic requires a fully periodic system");
     }
     auto pbc = torch::tensor(
         {domain->xperiodic, domain->yperiodic, domain->zperiodic},

src/ML-METATOMIC/metatomic_timer.cpp

@@ -1,3 +1,4 @@
+#include <mpi.h>
 #include <mutex>
 #include <iostream>
 
@@ -57,7 +58,10 @@ MetatomicTimer::~MetatomicTimer() {
             std::cerr << "\n" << indent << this->name_;
         }
 
-        std::cerr << " took " << elapsed / 1e6 << "ms" << std::flush;
+        int rank;
+        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+
+        std::cerr << " took " << elapsed / 1e6 << "ms (rank " << rank << ")" << std::flush;
         METATOMIC_TIMER_DEPTH -= 1;
     }
 }

src/ML-METATOMIC/metatomic_types.cpp

@@ -46,7 +46,7 @@ void PairMetatomicData::load_model(
 
    this->model_path = path;
    if (this->model != nullptr) {
-       lmp->error->all(FLERR, "torch model is already loaded");
+       lmp->error->one(FLERR, "torch model is already loaded");
    }
 
    torch::optional<std::string> extensions = torch::nullopt;
@@ -59,7 +59,7 @@ void PairMetatomicData::load_model(
            metatomic_torch::load_atomistic_model(this->model_path, extensions)
        );
    } catch (const c10::Error& e) {
-       lmp->error->all(FLERR, "failed to load metatomic model at '{}': {}", path, e.what());
+       lmp->error->one(FLERR, "failed to load metatomic model at '{}': {}", path, e.what());
    }
 
    auto capabilities_ivalue = this->model->run_method("capabilities");