From 987fd71848c493d543ce9e3a759c92298f2de0a9 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Tue, 22 Jul 2025 12:28:51 +0200 Subject: [PATCH 1/2] [metatomic] do not use the slow torch::Tensor indexing operator[] on torch::Tensor does a lot of things from gradients tracking to bound checks, and is a lot (~3 orders of magnitude) slower than directly operating on the data. --- src/metatomic/metatomic.cpp | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/metatomic/metatomic.cpp b/src/metatomic/metatomic.cpp index 4e4ab19270..3b27273ca1 100644 --- a/src/metatomic/metatomic.cpp +++ b/src/metatomic/metatomic.cpp @@ -762,13 +762,14 @@ metatensor_torch::TensorBlock MetatomicPlumedAction::computeNeighbors( torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU) ); - auto pair_samples_values = torch::zeros({n_pairs, 5}, labels_options.device(torch::kCPU)); + auto pair_samples_values = torch::empty({n_pairs, 5}, labels_options.device(torch::kCPU)); + auto pair_samples_values_ptr = pair_samples_values.accessor(); for (unsigned i=0; i(vesin_neighbor_list->pairs[i][0]); - pair_samples_values[i][1] = static_cast(vesin_neighbor_list->pairs[i][1]); - pair_samples_values[i][2] = vesin_neighbor_list->shifts[i][0]; - pair_samples_values[i][3] = vesin_neighbor_list->shifts[i][1]; - pair_samples_values[i][4] = vesin_neighbor_list->shifts[i][2]; + pair_samples_values_ptr[i][0] = static_cast(vesin_neighbor_list->pairs[i][0]); + pair_samples_values_ptr[i][1] = static_cast(vesin_neighbor_list->pairs[i][1]); + pair_samples_values_ptr[i][2] = vesin_neighbor_list->shifts[i][0]; + pair_samples_values_ptr[i][3] = vesin_neighbor_list->shifts[i][1]; + pair_samples_values_ptr[i][4] = vesin_neighbor_list->shifts[i][2]; } auto neighbor_samples = torch::make_intrusive( From fb7af28d2d4ece69a3ab4690acae73109038ff6e Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Wed, 25 Jun 2025 15:59:52 +0200 Subject: [PATCH 2/2] Actually run metatomic regtests on CI --- .github/workflows/linuxWF.yml | 2 +- regtest/metatomic/Makefile | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) create mode 100644 regtest/metatomic/Makefile diff --git a/.github/workflows/linuxWF.yml b/.github/workflows/linuxWF.yml index 67d8d1c4f3..02bfecf877 100644 --- a/.github/workflows/linuxWF.yml +++ b/.github/workflows/linuxWF.yml @@ -86,7 +86,7 @@ jobs: mdtraj \ MDAnalysis \ "torch>=2.7" \ - "metatomic-torch>=0.1,<0.2" \ + "metatomic-torch>=0.1.3,<0.2" \ "featomic-torch==0.7.0" # torch 2.7 above is the first one to use cxx11 ABI for the PyPI wheels diff --git a/regtest/metatomic/Makefile b/regtest/metatomic/Makefile new file mode 100644 index 0000000000..430b3123ed --- /dev/null +++ b/regtest/metatomic/Makefile @@ -0,0 +1 @@ +include ../scripts/module.make