[Code scan] Validate TensorFlow multi-device descriptor/neighbour dimensions before raw kernels #5623

Found during a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

Several TensorFlow multi-device force/virial ops derive ndescrpt and nnei with integer division, then rely on assert for important shape relationships such as nnei * 4 == ndescrpt.

Evidence:

• prod_force_multi_device derives both values by division and checks key invariants only with assert:

  deepmd-kit/source/op/tf/prod_force_multi_device.cc

  Lines 91 to 130 in 73de44b

  	const int* natoms = natoms_tensor.flat<int>().data();
  	int nloc = natoms[0];
  	int nall = natoms[1];
  	int nframes = net_deriv_tensor.shape().dim_size(0);
  	int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
  	int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
  	// check the sizes
  	OP_REQUIRES(
  	context, (nframes == in_deriv_tensor.shape().dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of samples should match"));
  	OP_REQUIRES(
  	context, (nframes == nlist_tensor.shape().dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of samples should match"));
  	OP_REQUIRES(
  	context,
  	(int_64(nloc) * ndescrpt * 3 == in_deriv_tensor.shape().dim_size(1)),
  	deepmd::tf_compat::InvalidArgument(
  	"number of descriptors should match"));
  	// Create an output tensor
  	TensorShape force_shape;
  	force_shape.AddDim(nframes);
  	force_shape.AddDim(3 * static_cast<int64_t>(nall));
  	Tensor* force_tensor = NULL;
  	int context_output_index = 0;
  	OP_REQUIRES_OK(context,
  	context->allocate_output(context_output_index++, force_shape,
  	&force_tensor));
  	DeviceFunctor()(device, context->eigen_device<Device>());
  	assert(nframes == force_shape.dim_size(0));
  	assert(nframes == net_deriv_tensor.shape().dim_size(0));
  	assert(nframes == in_deriv_tensor.shape().dim_size(0));
  	assert(nframes == nlist_tensor.shape().dim_size(0));
  	assert(nall * 3 == force_shape.dim_size(1));
  	assert(static_cast<int64_t>(nloc) * ndescrpt ==
  	net_deriv_tensor.shape().dim_size(1));
  	assert(static_cast<int64_t>(nloc) * ndescrpt * 3 ==
  	in_deriv_tensor.shape().dim_size(1));
  	assert(static_cast<int64_t>(nloc) * nnei ==
  	nlist_tensor.shape().dim_size(1));
  	assert(nnei * 4 == ndescrpt);

• prod_force_grad_multi_device has the same release-only gap before raw pointer use:

  deepmd-kit/source/op/tf/prod_force_grad_multi_device.cc

  Lines 74 to 127 in 73de44b

  	int nframes = net_deriv_tensor.shape().dim_size(0);
  	int nloc = natoms(0);
  	int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
  	int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
  	// check the sizes
  	OP_REQUIRES(
  	context, (nframes == grad_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(
  	context, (nframes == in_deriv_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(
  	context, (nframes == nlist_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(context, (nloc * 3 == grad_shape.dim_size(1)),
  	deepmd::tf_compat::InvalidArgument(
  	"input grad shape should be 3 x natoms"));
  	OP_REQUIRES(context,
  	(int_64(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)),
  	deepmd::tf_compat::InvalidArgument(
  	"number of descriptors should match"));
  	OP_REQUIRES(
  	context, (nnei == n_a_sel + n_r_sel),
  	deepmd::tf_compat::InvalidArgument("number of neighbors should match"));
  	// Create an output tensor
  	TensorShape grad_net_shape;
  	grad_net_shape.AddDim(nframes);
  	grad_net_shape.AddDim(int_64(nloc) * ndescrpt);
  	// allocate the output tensor
  	Tensor* grad_net_tensor = NULL;
  	int context_output_index = 0;
  	OP_REQUIRES_OK(context,
  	context->allocate_output(context_output_index++,
  	grad_net_shape, &grad_net_tensor));
  	DeviceFunctor()(device, context->eigen_device<Device>());
  	assert(nframes == grad_net_shape.dim_size(0));
  	assert(nframes == grad_shape.dim_size(0));
  	assert(nframes == net_deriv_tensor.shape().dim_size(0));
  	assert(nframes == in_deriv_tensor.shape().dim_size(0));
  	assert(nframes == nlist_tensor.shape().dim_size(0));
  	assert(static_cast<int64_t>(nloc) * ndescrpt == grad_net_shape.dim_size(1));
  	assert(nloc * 3 == grad_shape.dim_size(1));
  	assert(static_cast<int64_t>(nloc) * ndescrpt ==
  	net_deriv_tensor.shape().dim_size(1));
  	assert(static_cast<int64_t>(nloc) * ndescrpt * 3 ==
  	in_deriv_tensor.shape().dim_size(1));
  	assert(static_cast<int64_t>(nloc) * nnei ==
  	nlist_tensor.shape().dim_size(1));
  	assert(nnei * 4 == ndescrpt);
  	// flat the tensors

• prod_virial_grad_multi_device also relies on assert(nnei * 4 == ndescrpt) after allocating and before dispatch:

  deepmd-kit/source/op/tf/prod_virial_grad_multi_device.cc

  Lines 80 to 141 in 73de44b

  	int nframes = net_deriv_tensor.shape().dim_size(0);
  	int nloc = natoms(0);
  	int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
  	int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
  	// check the sizes
  	OP_REQUIRES(
  	context, (nframes == grad_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(
  	context, (nframes == in_deriv_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(
  	context, (nframes == rij_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(
  	context, (nframes == nlist_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(context, (9 == grad_shape.dim_size(1)),
  	deepmd::tf_compat::InvalidArgument(
  	"input grad shape should be 3 x natoms"));
  	OP_REQUIRES(context,
  	(int_64(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)),
  	deepmd::tf_compat::InvalidArgument(
  	"number of descriptors should match"));
  	OP_REQUIRES(
  	context, (int_64(nloc) * nnei * 3 == rij_shape.dim_size(1)),
  	deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3"));
  	OP_REQUIRES(
  	context, (nnei == n_a_sel + n_r_sel),
  	deepmd::tf_compat::InvalidArgument("number of neighbors should match"));
  	// Create an output tensor
  	TensorShape grad_net_shape;
  	grad_net_shape.AddDim(nframes);
  	grad_net_shape.AddDim(int_64(nloc) * ndescrpt);
  	// allocate the output tensor
  	Tensor* grad_net_tensor = NULL;
  	int context_output_index = 0;
  	OP_REQUIRES_OK(context,
  	context->allocate_output(context_output_index++,
  	grad_net_shape, &grad_net_tensor));
  	DeviceFunctor()(device, context->eigen_device<Device>());
  	assert(nframes == grad_net_shape.dim_size(0));
  	assert(nframes == grad_shape.dim_size(0));
  	assert(nframes == net_deriv_tensor.shape().dim_size(0));
  	assert(nframes == in_deriv_tensor.shape().dim_size(0));
  	assert(nframes == rij_tensor.shape().dim_size(0));
  	assert(nframes == nlist_tensor.shape().dim_size(0));
  	assert(static_cast<int64_t>(nloc) * ndescrpt == grad_net_shape.dim_size(1));
  	assert(9 == grad_shape.dim_size(1));
  	assert(static_cast<int64_t>(nloc) * ndescrpt ==
  	net_deriv_tensor.shape().dim_size(1));
  	assert(static_cast<int64_t>(nloc) * ndescrpt * 3 ==
  	in_deriv_tensor.shape().dim_size(1));
  	assert(static_cast<int64_t>(nloc) * nnei * 3 ==
  	rij_tensor.shape().dim_size(1));
  	assert(static_cast<int64_t>(nloc) * nnei ==
  	nlist_tensor.shape().dim_size(1));
  	assert(nnei * 4 == ndescrpt);

Impact

Release builds compile out assert. Malformed tensor widths that are not exact multiples of nloc, or descriptor widths that do not match the expected neighbor stride, can reach raw CPU/GPU kernels with inconsistent shapes and cause out-of-bounds reads/writes.

Suggested Fix

Convert descriptor/nlist divisibility and ndescrpt == nnei * stride checks to OP_REQUIRES before allocation and kernel launch. Add TensorFlow custom-op tests passing inconsistent net_deriv, in_deriv, rij, and nlist widths and expecting InvalidArgument.