fix(ops): validate tabulation tensor shapes

source/lib/include/tabulate_validation.h

@@ -0,0 +1,183 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+#pragma once
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <string>
+
+namespace deepmd {
+
+// Multiply non-negative tensor dimensions without invoking signed overflow.
+inline bool tabulate_checked_product(const int64_t lhs,
+                                     const int64_t rhs,
+                                     int64_t& product) {
+  if (lhs < 0 || rhs < 0 ||
+      (rhs != 0 && lhs > std::numeric_limits<int64_t>::max() / rhs)) {
+    return false;
+  }
+  product = lhs * rhs;
+  return true;
+}
+
+// Validate the five table metadata values consumed by the native tabulation
+// kernels and reproduce the spline-row count used by the table generator.
+// Keeping this calculation shared prevents the TensorFlow and PyTorch wrappers
+// from accepting different raw-buffer contracts.
+template <typename FPTYPE>
+bool tabulate_required_table_rows(const FPTYPE* table_info,
+                                  const bool symmetric_range,
+                                  int64_t& required_rows,
+                                  std::string& error) {
+  const FPTYPE lower_value = table_info[0];
+  const FPTYPE upper_value = table_info[1];
+  const FPTYPE max_value = table_info[2];
+  const FPTYPE stride0_value = table_info[3];
+  const FPTYPE stride1_value = table_info[4];
+  const double lower = static_cast<double>(lower_value);
+  const double upper = static_cast<double>(upper_value);
+  const double max = static_cast<double>(max_value);
+  const double stride0 = static_cast<double>(stride0_value);
+  const double stride1 = static_cast<double>(stride1_value);
+  if (!std::isfinite(lower) || !std::isfinite(upper) || !std::isfinite(max) ||
+      !std::isfinite(stride0) || !std::isfinite(stride1)) {
+    error = "table_info values must be finite";
+    return false;
+  }
+  if (stride0 <= 0.0 || stride1 <= 0.0) {
+    error = "table_info strides must be positive";
+    return false;
+  }
+
+  const double min = symmetric_range ? -max : lower;
+  if (min > lower || lower > upper || upper > max) {
+    error = symmetric_range
+                ? "table_info must satisfy -max <= lower <= upper <= max"
+                : "table_info must satisfy lower <= upper <= max";
+    return false;
+  }
+
+  const double total_intervals =
+      (symmetric_range ? (lower - min) / stride1 : 0.0) +
+      (upper - lower) / stride0 + (max - upper) / stride1;
+  if (!std::isfinite(total_intervals) || total_intervals <= 0.0) {
+    error = "table_info must describe at least one spline interval";
+    return false;
+  }
+
+  // Native locators truncate each region in FPTYPE precision and clamp the
+  // high-tail branch with nextafter. Reproduce those operations exactly and
+  // size for the largest reachable table index, rather than flooring the sum
+  // of fractional interval counts once.
+  const FPTYPE min_value = symmetric_range ? -max_value : lower_value;
+  const double native_int_max =
+      static_cast<double>(std::numeric_limits<int>::max());
+  auto truncate_native_ratio = [&](const FPTYPE numerator,
+                                   const FPTYPE denominator,
+                                   int64_t& result) -> bool {
+    const FPTYPE ratio_value = numerator / denominator;
+    const double ratio = static_cast<double>(ratio_value);
+    if (!std::isfinite(ratio) || ratio < 0.0 || ratio > native_int_max) {
+      error = "table_info describes an invalid spline index";
+      return false;
+    }
+    result = static_cast<int>(ratio_value);
+    return true;
+  };
+
+  int64_t lower_count = 0;
+  int64_t middle_count = 0;
+  if ((symmetric_range &&
+       !truncate_native_ratio(lower_value - min_value, stride1_value,
+                              lower_count)) ||
+      !truncate_native_ratio(upper_value - lower_value, stride0_value,
+                             middle_count)) {
+    return false;
+  }
+  const int64_t first_upper = lower_count + middle_count;
+  if (first_upper > std::numeric_limits<int>::max()) {
+    error = "table_info describes too many spline intervals";
+    return false;
+  }
+
+  int64_t max_reachable_index = 0;
+  auto include_region_end = [&](const FPTYPE end,
+                                const FPTYPE start,
+                                const FPTYPE stride,
+                                const int64_t base) -> bool {
+    if (!(end > start)) {
+      return true;
+    }
+    int64_t offset = 0;
+    if (!truncate_native_ratio(std::nextafter(end, start) - start, stride,
+                               offset)) {
+      return false;
+    }
+    const int64_t candidate = base + offset;
+    if (candidate > std::numeric_limits<int>::max()) {
+      error = "table_info describes too many spline intervals";
+      return false;
+    }
+    max_reachable_index = std::max(max_reachable_index, candidate);
+    return true;
+  };
+
+  if ((symmetric_range &&
+       !include_region_end(lower_value, min_value, stride1_value, 0)) ||
+      !include_region_end(upper_value, lower_value, stride0_value,
+                          lower_count) ||
+      !include_region_end(max_value, upper_value, stride1_value,
+                          first_upper)) {
+    return false;
+  }
+
+  // Inputs at or above max use this exact high-tail index. When max == upper,
+  // it deliberately selects the row after an aligned middle region.
+  int64_t high_tail_offset = 0;
+  const FPTYPE high_tail_boundary = std::nextafter(max_value, min_value);
+  const FPTYPE high_tail_delta = high_tail_boundary - upper_value;
+  const FPTYPE high_tail_ratio = high_tail_delta / stride1_value;
+  const double high_tail_ratio_double = static_cast<double>(high_tail_ratio);
+  if (!std::isfinite(high_tail_ratio_double) ||
+      high_tail_ratio_double <
+          static_cast<double>(std::numeric_limits<int>::min()) ||
+      high_tail_ratio_double > native_int_max) {
+    error = "table_info describes an invalid spline index";
+    return false;
+  }
+  high_tail_offset = static_cast<int>(high_tail_ratio);
+  const int64_t high_tail_index = first_upper + high_tail_offset;
+  if (high_tail_index < 0 ||
+      high_tail_index > std::numeric_limits<int>::max()) {
+    error = "table_info describes an invalid spline index";
+    return false;
+  }
+  max_reachable_index = std::max(max_reachable_index, high_tail_index);
+  required_rows = max_reachable_index + 1;
+  return true;
+}
+
+// Convert the validated row count into the flattened coefficient count while
+// guarding the multiplication used by both framework wrappers.
+inline bool tabulate_required_table_elements(const int64_t required_rows,
+                                             const int64_t last_layer_size,
+                                             int64_t& required_elements,
+                                             std::string& error) {
+  constexpr int64_t coefficients_per_feature = 6;
+  if (required_rows <= 0 || last_layer_size <= 0) {
+    error = "table dimensions must be positive";
+    return false;
+  }
+  int64_t feature_elements = 0;
+  if (!tabulate_checked_product(last_layer_size, coefficients_per_feature,
+                                feature_elements) ||
+      !tabulate_checked_product(required_rows, feature_elements,
+                                required_elements)) {
+    error = "required table size exceeds the supported integer range";
+    return false;
+  }
+  return true;
+}
+
+}  // namespace deepmd

source/op/pt/tabulate_multi_device.cc

@@ -1,10 +1,12 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 #include <torch/torch.h>
 
+#include <cstdint>
 #include <string>
 #include <vector>
 
 #include "tabulate.h"
+#include "tabulate_validation.h"
 
 #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
 #include "device.h"
@@ -18,6 +20,119 @@ void GetTensorDevice(const torch::Tensor& t, std::string& str) {
   }
 }
 
+void CheckTabulateDataTensor(const torch::Tensor& tensor,
+                             const torch::Tensor& table_tensor,
+                             const char* name) {
+  TORCH_CHECK(tensor.scalar_type() == table_tensor.scalar_type(), name,
+              " must have the same dtype as table");
+  TORCH_CHECK(tensor.device() == table_tensor.device(), name,
+              " must be on the same device as table");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+
+// The native SE-A kernels assume that every non-empty local-atom dimension
+// has at least one neighbor. Reject the degenerate shape before calculating
+// raw buffer offsets in either public or autograd-internal entry points.
+void CheckTabulateSeANeighborCount(const torch::Tensor& em_tensor) {
+  TORCH_CHECK(em_tensor.dim() == 3, "em must be rank 3");
+  TORCH_CHECK(em_tensor.size(0) == 0 || em_tensor.size(1) > 0,
+              "em must contain at least one neighbor when nloc is positive");
+}
+
+template <typename FPTYPE>
+void CheckTabulateTable(const torch::Tensor& table_tensor,
+                        const torch::Tensor& table_info_tensor,
+                        const int64_t last_layer_size,
+                        const bool symmetric_range) {
+  TORCH_CHECK(table_tensor.dim() == 2, "table must be rank 2");
+  TORCH_CHECK(table_tensor.scalar_type() == torch::kFloat ||
+                  table_tensor.scalar_type() == torch::kDouble,
+              "table must use float32 or float64");
+  TORCH_CHECK(table_tensor.device().is_cpu() || table_tensor.device().is_cuda(),
+              "table must be on a CPU or CUDA/ROCm device");
+  TORCH_CHECK(table_tensor.is_contiguous(), "table must be contiguous");
+  TORCH_CHECK(last_layer_size > 0, "last_layer_size must be positive");
+  TORCH_CHECK(table_info_tensor.device().is_cpu(),
+              "table_info must be on the CPU");
+  TORCH_CHECK(table_info_tensor.scalar_type() == table_tensor.scalar_type(),
+              "table_info must have the same dtype as table");
+  TORCH_CHECK(table_info_tensor.is_contiguous(),
+              "table_info must be contiguous");
+  TORCH_CHECK(table_info_tensor.numel() >= 5,
+              "table_info must contain at least 5 values");
+
+  int64_t required_rows = 0;
+  std::string error;
+  TORCH_CHECK(deepmd::tabulate_required_table_rows<FPTYPE>(
+                  table_info_tensor.data_ptr<FPTYPE>(), symmetric_range,
+                  required_rows, error),
+              error);
+  int64_t required_elements = 0;
+  TORCH_CHECK(deepmd::tabulate_required_table_elements(
+                  required_rows, last_layer_size, required_elements, error),
+              error);
+  TORCH_CHECK(table_tensor.numel() >= required_elements,
+              "table does not contain enough coefficients for table_info and "
+              "last_layer_size");
+}
+
+template <typename FPTYPE>
+void CheckTabulateSeAInputs(const torch::Tensor& table_tensor,
+                            const torch::Tensor& table_info_tensor,
+                            const torch::Tensor& em_x_tensor,
+                            const torch::Tensor& em_tensor,
+                            const torch::Tensor& two_embed_tensor,
+                            const int64_t last_layer_size) {
+  CheckTabulateTable<FPTYPE>(table_tensor, table_info_tensor, last_layer_size,
+                             false);
+  TORCH_CHECK(em_tensor.dim() == 3 && em_tensor.size(2) == 4,
+              "em must have shape [nloc, nnei, 4]");
+  CheckTabulateSeANeighborCount(em_tensor);
+  const int64_t neighbor_count = em_tensor.numel() / 4;
+  TORCH_CHECK(em_x_tensor.dim() == 2 && em_x_tensor.numel() == neighbor_count,
+              "em_x must be rank 2 and contain nloc * nnei values");
+  CheckTabulateDataTensor(em_x_tensor, table_tensor, "em_x");
+  CheckTabulateDataTensor(em_tensor, table_tensor, "em");
+  if (two_embed_tensor.defined()) {
+    TORCH_CHECK(two_embed_tensor.dim() == 2, "two_embed must be rank 2");
+    int64_t expected_two_embed_elements = 0;
+    TORCH_CHECK(
+        deepmd::tabulate_checked_product(neighbor_count, last_layer_size,
+                                         expected_two_embed_elements),
+        "two_embed element count exceeds the supported integer range");
+    TORCH_CHECK(two_embed_tensor.numel() == expected_two_embed_elements,
+                "two_embed must contain nloc * nnei * last_layer_size values");
+    CheckTabulateDataTensor(two_embed_tensor, table_tensor, "two_embed");
+  }
+}
+
+template <typename FPTYPE>
+void CheckTabulateSeTInputs(const torch::Tensor& table_tensor,
+                            const torch::Tensor& table_info_tensor,
+                            const torch::Tensor& em_x_tensor,
+                            const torch::Tensor& em_tensor,
+                            const int64_t last_layer_size) {
+  CheckTabulateTable<FPTYPE>(table_tensor, table_info_tensor, last_layer_size,
+                             true);
+  TORCH_CHECK(em_tensor.dim() == 3, "em must be rank 3");
+  TORCH_CHECK(
+      em_x_tensor.dim() == 2 && em_x_tensor.numel() == em_tensor.numel(),
+      "em_x must be rank 2 and contain the same number of values as em");
+  CheckTabulateDataTensor(em_x_tensor, table_tensor, "em_x");
+  CheckTabulateDataTensor(em_tensor, table_tensor, "em");
+}
+
+template <typename FPTYPE>
+void CheckTabulateSeRInputs(const torch::Tensor& table_tensor,
+                            const torch::Tensor& table_info_tensor,
+                            const torch::Tensor& em_tensor,
+                            const int64_t last_layer_size) {
+  CheckTabulateTable<FPTYPE>(table_tensor, table_info_tensor, last_layer_size,
+                             false);
+  TORCH_CHECK(em_tensor.dim() == 2, "em must be rank 2");
+  CheckTabulateDataTensor(em_tensor, table_tensor, "em");
+}
+
 template <typename FPTYPE>
 void TabulateFusionSeAForward(const torch::Tensor& table_tensor,
                               const torch::Tensor& table_info_tensor,
@@ -39,6 +154,7 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor,
   if (two_embed_tensor.defined() && two_embed_tensor.dim() != 2) {
     throw std::invalid_argument("Dim of input should be 2");
   }
+  CheckTabulateSeANeighborCount(em_tensor);
   // get the device
   std::string device;
   GetTensorDevice(table_tensor, device);
@@ -87,6 +203,7 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor,
   if (dy_tensor.dim() != 3) {
     throw std::invalid_argument("Dim of dy_tensor should be 3");
   }
+  CheckTabulateSeANeighborCount(em_tensor);
   std::string device;
   GetTensorDevice(table_tensor, device);
   // flat the tensors
@@ -145,6 +262,7 @@ void TabulateFusionSeAGradGradForward(const torch::Tensor& table_tensor,
   if (dz_dy_dem_tensor.dim() != 3) {
     throw std::invalid_argument("Dim of dz_dy_dem should be 3");
   }
+  CheckTabulateSeANeighborCount(em_tensor);
   // get the device
   std::string device;
   GetTensorDevice(table_tensor, device);
@@ -776,6 +894,8 @@ class TabulateFusionSeAOp
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
       int64_t last_layer_size) {
+    CheckTabulateSeAInputs<FPTYPE>(table_tensor, table_info_tensor, em_x_tensor,
+                                   em_tensor, at::Tensor(), last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())
@@ -962,6 +1082,9 @@ class TabulateFusionSeAttenOp
       const torch::Tensor& two_embed_tensor,
       int64_t last_layer_size,
       bool is_sorted) {
+    CheckTabulateSeAInputs<FPTYPE>(table_tensor, table_info_tensor, em_x_tensor,
+                                   em_tensor, two_embed_tensor,
+                                   last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())
@@ -1130,6 +1253,8 @@ class TabulateFusionSeTOp
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
       int64_t last_layer_size) {
+    CheckTabulateSeTInputs<FPTYPE>(table_tensor, table_info_tensor, em_x_tensor,
+                                   em_tensor, last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())
@@ -1283,6 +1408,8 @@ class TabulateFusionSeROp
       const torch::Tensor& table_info_tensor,
       const torch::Tensor& em_tensor,
       int64_t last_layer_size) {
+    CheckTabulateSeRInputs<FPTYPE>(table_tensor, table_info_tensor, em_tensor,
+                                   last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())
@@ -1441,6 +1568,8 @@ class TabulateFusionSeTTebdOp
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
       int64_t last_layer_size) {
+    CheckTabulateSeTInputs<FPTYPE>(table_tensor, table_info_tensor, em_x_tensor,
+                                   em_tensor, last_layer_size);
     // allocate output tensors
     auto options = torch::TensorOptions()
                        .dtype(table_tensor.dtype())