[Code scan] Validate TensorFlow tabulation GPU sizes before launching kernels

[njzjz]

This issue comes from a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

Several TensorFlow tabulation GPU paths validate last_layer_size <= 1024 only after launching the CUDA/ROCm kernel.

Examples in the grad-grad paths:

• deepmd-kit/source/op/tf/tabulate_multi_device.cc

  Lines 338 to 347 in 73de44b

  	if (device == "GPU") {
  	#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  	deepmd::tabulate_fusion_se_a_grad_grad_gpu(
  	dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
  	dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted);
  	#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  	OP_REQUIRES(context, (last_layer_size <= 1024),
  	deepmd::tf_compat::InvalidArgument(
  	"In the process of model compression, the size of the "
  	"last layer of embedding net must be less than 1024!"));

• deepmd-kit/source/op/tf/tabulate_multi_device.cc

  Lines 555 to 564 in 73de44b

  	if (device == "GPU") {
  	#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  	deepmd::tabulate_fusion_se_a_grad_grad_gpu(
  	dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
  	dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted);
  	#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  	OP_REQUIRES(context, (last_layer_size <= 1024),
  	deepmd::tf_compat::InvalidArgument(
  	"In the process of model compression, the size of the "
  	"last layer of embedding net must be less than 1024!"));

• deepmd-kit/source/op/tf/tabulate_multi_device.cc

  Lines 746 to 755 in 73de44b

  	if (device == "GPU") {
  	#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  	deepmd::tabulate_fusion_se_t_grad_grad_gpu(
  	dz_dy, table, table_info, em_x, em, dz_dy_dem_x, dz_dy_dem, nloc,
  	nnei_i, nnei_j, last_layer_size);
  	#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  	OP_REQUIRES(context, (last_layer_size <= 1024),
  	deepmd::tf_compat::InvalidArgument(
  	"In the process of model compression, the size of the "
  	"last layer of embedding net must be less than 1024!"));

• deepmd-kit/source/op/tf/tabulate_multi_device.cc

  Lines 911 to 919 in 73de44b

  	if (device == "GPU") {
  	#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  	deepmd::tabulate_fusion_se_r_grad_grad_gpu(
  	dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size);
  	#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  	OP_REQUIRES(context, (last_layer_size <= 1024),
  	deepmd::tf_compat::InvalidArgument(
  	"In the process of model compression, the size of the "
  	"last layer of embedding net must be less than 1024!"));

The GPU wrappers launch kernels with last_layer_size as the block size or shared-memory dimension, for example:

• deepmd-kit/source/lib/src/gpu/tabulate.cu

  Lines 1023 to 1048 in 73de44b

  	void tabulate_fusion_se_a_gpu(FPTYPE* out,
  	const FPTYPE* table,
  	const FPTYPE* table_info,
  	const FPTYPE* em_x,
  	const FPTYPE* em,
  	const FPTYPE* two_embed,
  	const int nloc,
  	const int nnei,
  	const int last_layer_size,
  	const bool is_sorted) {
  	if (nloc <= 0) {
  	return;
  	}
  	DPErrcheck(gpuGetLastError());
  	DPErrcheck(gpuDeviceSynchronize());
  	tabulate_fusion_se_a_fifth_order_polynomial<FPTYPE, MM, KK>
  	#if GOOGLE_CUDA
  	<<<nloc, last_layer_size>>>
  	#elif TENSORFLOW_USE_ROCM
  	<<<nloc, last_layer_size, sizeof(FPTYPE) * MM * last_layer_size>>>
  	#else
  	#error "should not touch here"
  	#endif
  	(out, table, em_x, em, two_embed, table_info[0], table_info[1],
  	table_info[2], table_info[3], table_info[4], nnei, last_layer_size,
  	is_sorted);

• deepmd-kit/source/lib/src/gpu/tabulate.cu

  Lines 1085 to 1109 in 73de44b

  	void tabulate_fusion_se_a_grad_grad_gpu(FPTYPE* dz_dy,
  	const FPTYPE* table,
  	const FPTYPE* table_info,
  	const FPTYPE* em_x,
  	const FPTYPE* em,
  	const FPTYPE* two_embed,
  	const FPTYPE* dz_dy_dem_x,
  	const FPTYPE* dz_dy_dem,
  	const FPTYPE* dz_dy_dtwo,
  	const int nloc,
  	const int nnei,
  	const int last_layer_size,
  	const bool is_sorted) {
  	if (nloc <= 0) {
  	return;
  	}
  	DPErrcheck(gpuGetLastError());
  	DPErrcheck(gpuDeviceSynchronize());
  	DPErrcheck(gpuMemset(dz_dy, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size));
  	tabulate_fusion_se_a_grad_grad_fifth_order_polynomial<FPTYPE, MM, KK>
  	<<<nloc, last_layer_size, sizeof(FPTYPE) * MM * last_layer_size>>>(
  	dz_dy, table, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, dz_dy_dtwo,
  	table_info[0], table_info[1], table_info[2], table_info[3],
  	table_info[4], nnei, last_layer_size, is_sorted);
  	DPErrcheck(gpuGetLastError());

So oversized inputs can hit a CUDA runtime failure before TensorFlow returns the intended InvalidArgument.

Impact

Invalid compressed-model tabulation shapes can leave users with low-level GPU launch errors instead of a deterministic TensorFlow validation error. Forward tabulation paths also launch GPU kernels without the same prelaunch bound check.

Suggested fix

Validate last_layer_size > 0 && last_layer_size <= 1024 before every GPU tabulation wrapper call that uses it as a kernel launch dimension or shared-memory multiplier.

Add GPU tests that pass last_layer_size=1025 and assert a clean TensorFlow InvalidArgument without launching the kernel.