Skip to content

fix(tf): honor NVNMD tensor neighbor lists #5845

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
njzjz-bot wants to merge 1 commit into
base: master
Choose a base branch
from
+211 −60
Open

fix(tf): honor NVNMD tensor neighbor lists #5845

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
203 changes: 143 additions & 60 deletions source/op/tf/prod_env_mat_multi_device_nvnmd.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ date: 2021-12-6

*/

#include <cstring>

#include "coord.h"
#include "custom_op.h"
#include "errors.h"
Expand Down Expand Up @@ -112,28 +114,29 @@ static void _map_nei_info_cpu(int* nlist,
const bool& b_nlist_map);

template <typename FPTYPE>
static void _prepare_coord_nlist_cpu(OpKernelContext* context,
FPTYPE const** coord,
std::vector<FPTYPE>& coord_cpy,
int const** type,
std::vector<int>& type_cpy,
std::vector<int>& idx_mapping,
deepmd::InputNlist& inlist,
std::vector<int>& ilist,
std::vector<int>& numneigh,
std::vector<int*>& firstneigh,
std::vector<std::vector<int>>& jlist,
int& new_nall,
int& mem_cpy,
int& mem_nnei,
int& max_nbor_size,
const FPTYPE* box,
const int* mesh_tensor_data,
const int& nloc,
const int& nei_mode,
const float& rcut_r,
const int& max_cpy_trial,
const int& max_nnei_trial);
static tensorflow::Status _prepare_coord_nlist_cpu(
FPTYPE const** coord,
std::vector<FPTYPE>& coord_cpy,
int const** type,
std::vector<int>& type_cpy,
std::vector<int>& idx_mapping,
deepmd::InputNlist& inlist,
std::vector<int>& ilist,
std::vector<int>& numneigh,
std::vector<int*>& firstneigh,
std::vector<std::vector<int>>& jlist,
int& new_nall,
int& mem_cpy,
int& mem_nnei,
int& max_nbor_size,
const FPTYPE* box,
const int* mesh_tensor_data,
const int_64 mesh_tensor_size,
const int& nloc,
const int& nei_mode,
const float& rcut_r,
const int& max_cpy_trial,
const int& max_nnei_trial);

// instance of function

Expand Down Expand Up @@ -228,58 +231,132 @@ static void _map_nei_info_cpu(int* nlist,
ntypes, b_nlist_map);
}

static tensorflow::Status _prepare_mesh_nlist_cpu(
deepmd::InputNlist& inlist,
std::vector<int>& ilist,
std::vector<int>& numneigh,
std::vector<int*>& firstneigh,
std::vector<std::vector<int>>& jlist,
int& max_nbor_size,
const int* mesh_tensor_data,
const int_64 mesh_tensor_size,
const int nloc,
const int nall) {
// Tensor-stored lists use a 16-int header followed by ilist, numneigh,
// and the concatenated neighbor rows. Validate every offset before copying
// because this input can be constructed directly by callers.
const int_64 header_size = 16 + static_cast<int_64>(2) * nloc;
if (mesh_tensor_size < header_size) {
return errors::InvalidArgument("invalid mesh tensor");
}

const int* ilist_in = mesh_tensor_data + 16;
const int* numneigh_in = mesh_tensor_data + 16 + nloc;
const int* neighbors_in = mesh_tensor_data + header_size;
std::vector<int_64> neighbor_offset(nloc + 1, 0);
int supplied_max_nbor_size = 0;
for (int ii = 0; ii < nloc; ++ii) {
const int_64 neighbor_count = numneigh_in[ii];
if (neighbor_count < 0 ||
neighbor_offset[ii] > mesh_tensor_size - header_size - neighbor_count) {
return errors::InvalidArgument("invalid mesh tensor");
}
neighbor_offset[ii + 1] = neighbor_offset[ii] + neighbor_count;
supplied_max_nbor_size = std::max(supplied_max_nbor_size, numneigh_in[ii]);
}

std::vector<unsigned char> seen(static_cast<size_t>(nloc), 0);
for (int ii = 0; ii < nloc; ++ii) {
const int i_idx = ilist_in[ii];
if (i_idx < 0 || i_idx >= nloc || seen[static_cast<size_t>(i_idx)]) {
return errors::InvalidArgument("invalid mesh tensor");
}
seen[static_cast<size_t>(i_idx)] = 1;
for (int_64 jj = neighbor_offset[ii]; jj < neighbor_offset[ii + 1]; ++jj) {
// Downstream environment-matrix kernels index both coordinates and
// atom types with this value, so accept only live frame atoms.
if (neighbors_in[jj] < 0 || neighbors_in[jj] >= nall) {
return errors::InvalidArgument("invalid mesh tensor");
}
}
}

for (int ii = 0; ii < nloc; ++ii) {
ilist[ii] = ilist_in[ii];
numneigh[ii] = numneigh_in[ii];
jlist[ii].assign(neighbors_in + neighbor_offset[ii],
neighbors_in + neighbor_offset[ii + 1]);
firstneigh[ii] = jlist[ii].data();
}
inlist = deepmd::InputNlist(nloc, ilist.data(), numneigh.data(),
firstneigh.data());
max_nbor_size = std::max(max_nbor_size, supplied_max_nbor_size);
return tensorflow::Status();
}

template <typename FPTYPE>
static void _prepare_coord_nlist_cpu(OpKernelContext* context,
FPTYPE const** coord,
std::vector<FPTYPE>& coord_cpy,
int const** type,
std::vector<int>& type_cpy,
std::vector<int>& idx_mapping,
deepmd::InputNlist& inlist,
std::vector<int>& ilist,
std::vector<int>& numneigh,
std::vector<int*>& firstneigh,
std::vector<std::vector<int>>& jlist,
int& new_nall,
int& mem_cpy,
int& mem_nnei,
int& max_nbor_size,
const FPTYPE* box,
const int* mesh_tensor_data,
const int& nloc,
const int& nei_mode,
const float& rcut_r,
const int& max_cpy_trial,
const int& max_nnei_trial) {
static tensorflow::Status _prepare_coord_nlist_cpu(
FPTYPE const** coord,
std::vector<FPTYPE>& coord_cpy,
int const** type,
std::vector<int>& type_cpy,
std::vector<int>& idx_mapping,
deepmd::InputNlist& inlist,
std::vector<int>& ilist,
std::vector<int>& numneigh,
std::vector<int*>& firstneigh,
std::vector<std::vector<int>>& jlist,
int& new_nall,
int& mem_cpy,
int& mem_nnei,
int& max_nbor_size,
const FPTYPE* box,
const int* mesh_tensor_data,
const int_64 mesh_tensor_size,
const int& nloc,
const int& nei_mode,
const float& rcut_r,
const int& max_cpy_trial,
const int& max_nnei_trial) {
inlist.inum = nloc;
if (nei_mode != 3) {
if (nei_mode != 3 && nei_mode != 4) {
// build nlist by myself
// normalize and copy coord
if (nei_mode == 1) {
int copy_ok = _norm_copy_coord_cpu(coord_cpy, type_cpy, idx_mapping,
new_nall, mem_cpy, *coord, box, *type,
nloc, max_cpy_trial, rcut_r);
OP_REQUIRES(context, copy_ok,
errors::Aborted("cannot allocate mem for copied coords"));
if (!copy_ok) {
return errors::Aborted("cannot allocate mem for copied coords");
}
*coord = &coord_cpy[0];
*type = &type_cpy[0];
}
// build nlist
int build_ok = _build_nlist_cpu(ilist, numneigh, firstneigh, jlist,
max_nbor_size, mem_nnei, *coord, nloc,
new_nall, max_nnei_trial, rcut_r);
OP_REQUIRES(context, build_ok,
errors::Aborted("cannot allocate mem for nlist"));
if (!build_ok) {
return errors::Aborted("cannot allocate mem for nlist");
}
inlist.ilist = &ilist[0];
inlist.numneigh = &numneigh[0];
inlist.firstneigh = &firstneigh[0];
} else if (nei_mode == 4) {
tensorflow::Status status = _prepare_mesh_nlist_cpu(
inlist, ilist, numneigh, firstneigh, jlist, max_nbor_size,
mesh_tensor_data, mesh_tensor_size, nloc, new_nall);
if (!status.ok()) {
return status;
}
} else {
// copy pointers to nlist data
memcpy(&inlist.ilist, 4 + mesh_tensor_data, sizeof(int*));
memcpy(&inlist.numneigh, 8 + mesh_tensor_data, sizeof(int*));
memcpy(&inlist.firstneigh, 12 + mesh_tensor_data, sizeof(int**));
max_nbor_size = max_numneigh(inlist);
}
return tensorflow::Status();
}

/*
Expand Down Expand Up @@ -506,11 +583,14 @@ class ProdEnvMatANvnmdQuantizeOp : public OpKernel {
std::vector<int> type_cpy;
int frame_nall = nall;
// prepare coord and nlist
_prepare_coord_nlist_cpu<FPTYPE>(
context, &coord, coord_cpy, &type, type_cpy, idx_mapping, inlist,
ilist, numneigh, firstneigh, jlist, frame_nall, mem_cpy, mem_nnei,
max_nbor_size, box, mesh_tensor.flat<int>().data(), nloc, nei_mode,
rcut_r, max_cpy_trial, max_nnei_trial);
OP_REQUIRES_OK(
context,
_prepare_coord_nlist_cpu<FPTYPE>(
&coord, coord_cpy, &type, type_cpy, idx_mapping, inlist, ilist,
numneigh, firstneigh, jlist, frame_nall, mem_cpy, mem_nnei,
max_nbor_size, box, mesh_tensor.flat<int>().data(),
mesh_tensor.NumElements(), nloc, nei_mode, rcut_r,
max_cpy_trial, max_nnei_trial));
// launch the cpu compute function
deepmd::prod_env_mat_a_nvnmd_quantize_cpu(
em, em_deriv, rij, nlist, coord, type, inlist, max_nbor_size, avg,
Expand Down Expand Up @@ -789,17 +869,20 @@ class ProdEnvMatAMixNvnmdQuantizeOp : public OpKernel {
std::vector<int> type_cpy;
int frame_nall = nall;
// prepare coord and nlist
_prepare_coord_nlist_cpu<FPTYPE>(
context, &coord, coord_cpy, &f_type, type_cpy, idx_mapping, inlist,
ilist, numneigh, firstneigh, jlist, frame_nall, mem_cpy, mem_nnei,
max_nbor_size, box, mesh_tensor.flat<int>().data(), nloc, nei_mode,
rcut_r, max_cpy_trial, max_nnei_trial);
OP_REQUIRES_OK(
context,
_prepare_coord_nlist_cpu<FPTYPE>(
&coord, coord_cpy, &f_type, type_cpy, idx_mapping, inlist,
ilist, numneigh, firstneigh, jlist, frame_nall, mem_cpy,
mem_nnei, max_nbor_size, box, mesh_tensor.flat<int>().data(),
mesh_tensor.NumElements(), nloc, nei_mode, rcut_r,
max_cpy_trial, max_nnei_trial));
// launch the cpu compute function
deepmd::prod_env_mat_a_nvnmd_quantize_cpu(
em, em_deriv, rij, nlist, coord, type, inlist, max_nbor_size, avg,
std, nloc, frame_nall, rcut_r, rcut_r_smth, sec_a, f_type);
// do nlist mapping if coords were copied
_map_nei_info_cpu(nlist, ntype, nmask, type, &idx_mapping[0], nloc,
_map_nei_info_cpu(nlist, ntype, nmask, type, idx_mapping.data(), nloc,
nnei, ntypes, b_nlist_map);
}
}
Expand Down
68 changes: 68 additions & 0 deletions source/tests/tf/test_nvnmd_op.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -421,5 +421,73 @@ def test_op(self) -> None:
tf.reset_default_graph()


class TestOpProdEnvMatNvnmdTensorNlist(tf.test.TestCase):
"""Verify NVNMD env-mat ops honor neighbor lists stored in mesh tensors."""

def setUp(self) -> None:
self.coord = np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0]])
self.atype = np.array([[0, 0, 0]], dtype=np.int32)
self.natoms = np.array([3, 3, 3], dtype=np.int32)
self.box = np.array([[10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0]])
self.sel = [2]
self.ndescrpt = 4 * sum(self.sel)
self.avg = np.zeros((1, self.ndescrpt))
self.std = np.ones((1, self.ndescrpt))

# All atoms are mutually within rcut, but this caller-provided list
# deliberately retains only one cyclic neighbor per local atom.
header = np.zeros(16, dtype=np.int32)
ilist = np.array([0, 1, 2], dtype=np.int32)
numneigh = np.ones(3, dtype=np.int32)
jlist = np.array([2, 0, 1], dtype=np.int32)
self.mesh = np.concatenate((header, ilist, numneigh, jlist))
self.expected_nlist = np.array([[2, -1, 0, -1, 1, -1]], dtype=np.int32)

def _run_op(self, mixed_types: bool, mesh: np.ndarray | None = None) -> np.ndarray:
op = (
op_module.prod_env_mat_a_mix_nvnmd_quantize
if mixed_types
else op_module.prod_env_mat_a_nvnmd_quantize
)
outputs = op(
tf.constant(self.coord, dtype=tf.float64),
tf.constant(self.atype),
tf.constant(self.natoms),
tf.constant(self.box, dtype=tf.float64),
tf.constant(self.mesh if mesh is None else mesh),
tf.constant(self.avg, dtype=tf.float64),
tf.constant(self.std, dtype=tf.float64),
rcut_a=-1.0,
rcut_r=3.0,
rcut_r_smth=0.5,
sel_a=self.sel,
sel_r=[0],
)
with self.cached_session() as sess:
return sess.run(outputs[3])

def test_standard_op_uses_tensor_neighbor_list(self) -> None:
np.testing.assert_array_equal(self._run_op(False), self.expected_nlist)

def test_mixed_type_op_uses_tensor_neighbor_list(self) -> None:
np.testing.assert_array_equal(self._run_op(True), self.expected_nlist)

def test_rejects_truncated_tensor_neighbor_list(self) -> None:
"""Reject mode-4 tensors before reading an incomplete list header."""
with self.assertRaisesRegex(
tf.errors.InvalidArgumentError, "invalid mesh tensor"
):
self._run_op(False, self.mesh[:17])

def test_rejects_out_of_range_tensor_neighbor_index(self) -> None:
"""Reject neighbors that would index beyond this frame's atoms."""
invalid_mesh = self.mesh.copy()
invalid_mesh[-1] = self.natoms[1]
with self.assertRaisesRegex(
tf.errors.InvalidArgumentError, "invalid mesh tensor"
):
self._run_op(False, invalid_mesh)


if __name__ == "__main__":
unittest.main()
Loading