fix(ops): validate tabulation tensor shapes

source/lib/include/tabulate_validation.h

@@ -0,0 +1,100 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+#pragma once
+
+#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 double lower = static_cast<double>(table_info[0]);
+  const double upper = static_cast<double>(table_info[1]);
+  const double max = static_cast<double>(table_info[2]);
+  const double stride0 = static_cast<double>(table_info[3]);
+  const double stride1 = static_cast<double>(table_info[4]);
+  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 lower_tail = symmetric_range ? (lower - min) / stride1 : 0.0;
+  const double middle = (upper - lower) / stride0;
+  const double upper_tail = (max - upper) / stride1;
+  const double total_intervals = lower_tail + middle + upper_tail;
+  const double max_segment =
+      static_cast<double>(std::numeric_limits<int>::max());
+  if (!std::isfinite(lower_tail) || !std::isfinite(middle) ||
+      !std::isfinite(upper_tail) || !std::isfinite(total_intervals) ||
+      total_intervals > max_segment) {
+    error = "table_info describes too many spline intervals";
+    return false;
+  }
+
+  // The Python table builder converts the sum to an integer once, which is
+  // observably different from truncating each range separately for SE-T.
+  required_rows = static_cast<int64_t>(total_intervals);
+  if (required_rows <= 0) {
+    error = "table_info must describe at least one spline interval";
+    return false;
+  }
+  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,109 @@ 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");
+}
+
+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]");
+  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,
@@ -776,6 +881,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 +1069,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 +1240,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 +1395,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 +1555,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())