diff --git a/source/lib/include/tabulate_validation.h b/source/lib/include/tabulate_validation.h new file mode 100644 index 0000000000..e145f1d7c1 --- /dev/null +++ b/source/lib/include/tabulate_validation.h @@ -0,0 +1,183 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include +#include +#include +#include +#include + +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::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 +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(lower_value); + const double upper = static_cast(upper_value); + const double max = static_cast(max_value); + const double stride0 = static_cast(stride0_value); + const double stride1 = static_cast(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(std::numeric_limits::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(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(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::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::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(high_tail_ratio); + if (!std::isfinite(high_tail_ratio_double) || + high_tail_ratio_double < + static_cast(std::numeric_limits::min()) || + high_tail_ratio_double > native_int_max) { + error = "table_info describes an invalid spline index"; + return false; + } + high_tail_offset = static_cast(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::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 diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc index cede1d03d9..13c479ebaf 100644 --- a/source/op/pt/tabulate_multi_device.cc +++ b/source/op/pt/tabulate_multi_device.cc @@ -1,10 +1,12 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include +#include #include #include #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 +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( + table_info_tensor.data_ptr(), 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 +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(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 +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(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 +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(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 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(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(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(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(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(table_tensor, table_info_tensor, em_x_tensor, + em_tensor, last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) diff --git a/source/op/tf/tabulate_multi_device.cc b/source/op/tf/tabulate_multi_device.cc index 174faf5213..eac357cd1f 100644 --- a/source/op/tf/tabulate_multi_device.cc +++ b/source/op/tf/tabulate_multi_device.cc @@ -2,6 +2,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include "custom_op.h" #include "tabulate.h" +#include "tabulate_validation.h" REGISTER_OP("TabulateFusion") .Attr("T: {float, double} = DT_DOUBLE") @@ -162,6 +163,134 @@ REGISTER_OP("TabulateFusionSeRGradGrad") .Input("descriptor: T") .Output("dz_dy: T"); +template +deepmd::tf_compat::Status ValidateTabulateTable(const Tensor& table_tensor, + const Tensor& table_info_tensor, + const int64_t last_layer_size, + const bool symmetric_range) { + if (table_tensor.dims() != 2) { + return deepmd::tf_compat::InvalidArgument("table must be rank 2"); + } + if (last_layer_size <= 0) { + return deepmd::tf_compat::InvalidArgument( + "last_layer_size must be positive"); + } + if (table_info_tensor.NumElements() < 5) { + return deepmd::tf_compat::InvalidArgument( + "table_info must contain at least 5 values"); + } + int64_t required_rows = 0; + std::string error; + if (!deepmd::tabulate_required_table_rows( + table_info_tensor.flat().data(), symmetric_range, + required_rows, error)) { + return deepmd::tf_compat::InvalidArgument(error); + } + int64_t required_elements = 0; + if (!deepmd::tabulate_required_table_elements(required_rows, last_layer_size, + required_elements, error)) { + return deepmd::tf_compat::InvalidArgument(error); + } + if (table_tensor.NumElements() < required_elements) { + return deepmd::tf_compat::InvalidArgument( + "table does not contain enough coefficients for table_info and " + "last_layer_size"); + } + return deepmd::tf_compat::Status(); +} + +template +deepmd::tf_compat::Status ValidateTabulateSeAInputs( + const Tensor& table_tensor, + const Tensor& table_info_tensor, + const Tensor& em_x_tensor, + const Tensor& em_tensor, + const Tensor* two_embed_tensor, + const int64_t last_layer_size) { + auto status = ValidateTabulateTable(table_tensor, table_info_tensor, + last_layer_size, false); + if (!status.ok()) { + return status; + } + if (em_tensor.dims() != 3 || em_tensor.dim_size(2) != 4) { + return deepmd::tf_compat::InvalidArgument( + "em must have shape [nloc, nnei, 4]"); + } + if (em_tensor.dim_size(0) > 0 && em_tensor.dim_size(1) == 0) { + return deepmd::tf_compat::InvalidArgument( + "em must contain at least one neighbor when nloc is positive"); + } + const int64_t neighbor_count = em_tensor.NumElements() / 4; + if (em_x_tensor.dims() != 2 || em_x_tensor.NumElements() != neighbor_count) { + return deepmd::tf_compat::InvalidArgument( + "em_x must be rank 2 and contain nloc * nnei values"); + } + if (two_embed_tensor != nullptr) { + int64_t expected_two_embed_elements = 0; + if (!deepmd::tabulate_checked_product(neighbor_count, last_layer_size, + expected_two_embed_elements)) { + return deepmd::tf_compat::InvalidArgument( + "two_embed element count exceeds the supported integer range"); + } + if (two_embed_tensor->dims() != 2 || + two_embed_tensor->NumElements() != expected_two_embed_elements) { + return deepmd::tf_compat::InvalidArgument( + "two_embed must be rank 2 and contain nloc * nnei * " + "last_layer_size values"); + } + } + return deepmd::tf_compat::Status(); +} + +template +deepmd::tf_compat::Status ValidateTabulateSeTInputs( + const Tensor& table_tensor, + const Tensor& table_info_tensor, + const Tensor& em_x_tensor, + const Tensor& em_tensor, + const int64_t last_layer_size) { + auto status = ValidateTabulateTable(table_tensor, table_info_tensor, + last_layer_size, true); + if (!status.ok()) { + return status; + } + if (em_tensor.dims() != 3) { + return deepmd::tf_compat::InvalidArgument("em must be rank 3"); + } + if (em_x_tensor.dims() != 2 || + em_x_tensor.NumElements() != em_tensor.NumElements()) { + return deepmd::tf_compat::InvalidArgument( + "em_x must be rank 2 and contain the same number of values as em"); + } + return deepmd::tf_compat::Status(); +} + +template +deepmd::tf_compat::Status ValidateTabulateSeRInputs( + const Tensor& table_tensor, + const Tensor& table_info_tensor, + const Tensor& em_tensor, + const int64_t last_layer_size) { + auto status = ValidateTabulateTable(table_tensor, table_info_tensor, + last_layer_size, false); + if (!status.ok()) { + return status; + } + if (em_tensor.dims() != 2) { + return deepmd::tf_compat::InvalidArgument("em must be rank 2"); + } + return deepmd::tf_compat::Status(); +} + +deepmd::tf_compat::Status ValidateTensorShape(const Tensor& tensor, + const TensorShape& expected, + const char* name) { + if (tensor.shape() != expected) { + return deepmd::tf_compat::InvalidArgument(name, " has an unexpected shape"); + } + return deepmd::tf_compat::Status(); +} + template class TabulateFusionSeAOp : public OpKernel { public: @@ -182,6 +311,9 @@ class TabulateFusionSeAOp : public OpKernel { const Tensor& table_info_tensor = context->input(context_input_index++); const Tensor& em_x_tensor = context->input(context_input_index++); const Tensor& em_tensor = context->input(context_input_index++); + OP_REQUIRES_OK(context, ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, nullptr, last_layer_size)); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); @@ -245,6 +377,22 @@ class TabulateFusionSeAGradOp : public OpKernel { const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, 4, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, nullptr, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(4); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, ValidateTensorShape( + dy_tensor, expected_descriptor_shape, "dy")); // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); @@ -309,6 +457,26 @@ class TabulateFusionSeAGradGradOp : public OpKernel { const Tensor& dz_dy_dem_x_tensor = context->input(context_input_index++); const Tensor& dz_dy_dem_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, 4, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, nullptr, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(4); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, + ValidateTensorShape(dz_dy_dem_x_tensor, em_x_tensor.shape(), + "dz_dy_dem_x")); + OP_REQUIRES_OK( + context, + ValidateTensorShape(dz_dy_dem_tensor, em_tensor.shape(), "dz_dy_dem")); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_x_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); @@ -379,6 +547,9 @@ class TabulateFusionSeAttenOp : public OpKernel { const Tensor& em_x_tensor = context->input(context_input_index++); const Tensor& em_tensor = context->input(context_input_index++); const Tensor& two_embed_tensor = context->input(context_input_index++); + OP_REQUIRES_OK(context, ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, &two_embed_tensor, last_layer_size)); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); @@ -450,6 +621,23 @@ class TabulateFusionSeAttenGradOp : public OpKernel { const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, 4, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, + ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, em_tensor, + &two_embed_tensor, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(4); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, ValidateTensorShape( + dy_tensor, expected_descriptor_shape, "dy")); // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); @@ -526,6 +714,30 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel { const Tensor& dz_dy_dem_tensor = context->input(context_input_index++); const Tensor& dz_dy_dtwo_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, 4, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, + ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, em_tensor, + &two_embed_tensor, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(4); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, + ValidateTensorShape(dz_dy_dem_x_tensor, em_x_tensor.shape(), + "dz_dy_dem_x")); + OP_REQUIRES_OK( + context, + ValidateTensorShape(dz_dy_dem_tensor, em_tensor.shape(), "dz_dy_dem")); + OP_REQUIRES_OK(context, + ValidateTensorShape(dz_dy_dtwo_tensor, + two_embed_tensor.shape(), "dz_dy_dtwo")); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_x_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); @@ -594,6 +806,9 @@ class TabulateFusionSeTOp : public OpKernel { const Tensor& table_info_tensor = context->input(context_input_index++); const Tensor& em_x_tensor = context->input(context_input_index++); const Tensor& em_tensor = context->input(context_input_index++); + OP_REQUIRES_OK(context, ValidateTabulateSeTInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, last_layer_size)); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); @@ -657,6 +872,21 @@ class TabulateFusionSeTGradOp : public OpKernel { const Tensor& em_tensor = context->input(context_input_index++); const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 2, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(1); + OP_REQUIRES_OK(context, ValidateTabulateSeTInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, ValidateTensorShape( + dy_tensor, expected_descriptor_shape, "dy")); // set size of the sample OP_REQUIRES( context, (dy_tensor.shape().dims() == 2), @@ -718,6 +948,25 @@ class TabulateFusionSeTGradGradOp : public OpKernel { const Tensor& dz_dy_dem_x_tensor = context->input(context_input_index++); const Tensor& dz_dy_dem_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 2, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(1); + OP_REQUIRES_OK(context, ValidateTabulateSeTInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, + ValidateTensorShape(dz_dy_dem_x_tensor, em_x_tensor.shape(), + "dz_dy_dem_x")); + OP_REQUIRES_OK( + context, + ValidateTensorShape(dz_dy_dem_tensor, em_tensor.shape(), "dz_dy_dem")); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_x_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); @@ -782,6 +1031,9 @@ class TabulateFusionSeROp : public OpKernel { const Tensor& table_tensor = context->input(context_input_index++); const Tensor& table_info_tensor = context->input(context_input_index++); const Tensor& em_tensor = context->input(context_input_index++); + OP_REQUIRES_OK(context, ValidateTabulateSeRInputs( + table_tensor, table_info_tensor, em_tensor, + last_layer_size)); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); @@ -839,6 +1091,23 @@ class TabulateFusionSeRGradOp : public OpKernel { const Tensor& em_tensor = context->input(context_input_index++); const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, nnei, " + "last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, ValidateTabulateSeRInputs( + table_tensor, table_info_tensor, em_tensor, + validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(em_tensor.dim_size(1)); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, ValidateTensorShape( + dy_tensor, expected_descriptor_shape, "dy")); // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); @@ -888,6 +1157,24 @@ class TabulateFusionSeRGradGradOp : public OpKernel { const Tensor& em_tensor = context->input(context_input_index++); const Tensor& dz_dy_dem_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, nnei, " + "last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, ValidateTabulateSeRInputs( + table_tensor, table_info_tensor, em_tensor, + validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(em_tensor.dim_size(1)); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK( + context, + ValidateTensorShape(dz_dy_dem_tensor, em_tensor.shape(), "dz_dy_dem")); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); diff --git a/source/tests/pt/test_tabulate_fusion_se_atten.py b/source/tests/pt/test_tabulate_fusion_se_atten.py index f46d0b1761..798d821f1d 100644 --- a/source/tests/pt/test_tabulate_fusion_se_atten.py +++ b/source/tests/pt/test_tabulate_fusion_se_atten.py @@ -1644,6 +1644,204 @@ def test_second_order_backward(self) -> None: (self.em_x_tensor, self.em_tensor, self.two_embed_tensor), ) + def test_rejects_mismatched_native_buffer_shapes(self) -> None: + invalid_inputs = ( + ( + self.em_x_tensor[:, :-1].contiguous(), + self.em_tensor, + self.two_embed_tensor, + "em_x must be rank 2", + ), + ( + self.em_x_tensor, + self.em_tensor[..., :3].contiguous(), + self.two_embed_tensor, + "em must have shape", + ), + ( + self.em_x_tensor, + self.em_tensor, + self.two_embed_tensor.reshape(-1)[:-1].reshape(1, -1), + "two_embed must contain", + ), + ) + for em_x, em, two_embed, message in invalid_inputs: + with ( + self.subTest(message=message), + self.assertRaisesRegex(RuntimeError, message), + ): + torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor, + em_x, + em, + two_embed, + self.last_layer_size, + self.is_sorted, + ) + + def test_rejects_zero_neighbors_for_nonempty_atoms(self) -> None: + empty_em_x = torch.empty( + (1, 0), dtype=self.table_tensor.dtype, device=env.DEVICE + ) + empty_em = torch.empty( + (1, 0, 4), dtype=self.table_tensor.dtype, device=env.DEVICE + ) + empty_two_embed = torch.empty( + (1, 0), dtype=self.table_tensor.dtype, device=env.DEVICE + ) + operations = ( + lambda: torch.ops.deepmd.tabulate_fusion_se_a( + self.table_tensor, + self.table_info_tensor, + empty_em_x, + empty_em, + self.last_layer_size, + ), + lambda: torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor, + empty_em_x, + empty_em, + empty_two_embed, + self.last_layer_size, + self.is_sorted, + ), + ) + for operation in operations: + with ( + self.subTest(operation=operation), + self.assertRaisesRegex(RuntimeError, "at least one neighbor"), + ): + operation() + + def test_backward_rejects_zero_neighbor_saved_tensors(self) -> None: + em_x_pointer = self.em_x_tensor.data_ptr() + em_pointer = self.em_tensor.data_ptr() + + def pack_hook(tensor: torch.Tensor) -> tuple[str, torch.Tensor]: + if tensor.data_ptr() == em_x_pointer: + return ("em_x", tensor) + if tensor.data_ptr() == em_pointer: + return ("em", tensor) + return ("unchanged", tensor) + + def unpack_hook(packed: tuple[str, torch.Tensor]) -> torch.Tensor: + name, tensor = packed + if name == "em_x": + return tensor.new_empty((self.nloc, 0)) + if name == "em": + return tensor.new_empty((self.nloc, 0, 4)) + return tensor + + with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook): + descriptor = torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor, + self.em_x_tensor, + self.em_tensor, + self.two_embed_tensor, + self.last_layer_size, + self.is_sorted, + )[0] + with self.assertRaisesRegex(RuntimeError, "at least one neighbor"): + descriptor.sum().backward() + + def test_accepts_flattened_em_x_layout(self) -> None: + result = torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor, + self.em_x_tensor.reshape(-1, 1), + self.em_tensor, + self.two_embed_tensor, + self.last_layer_size, + self.is_sorted, + ) + self.assertEqual(result[0].shape, self.expected_descriptor_tensor.shape) + + def test_rejects_short_table_buffers(self) -> None: + with self.assertRaisesRegex(RuntimeError, "table_info must contain"): + torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor[:4], + self.em_x_tensor, + self.em_tensor, + self.two_embed_tensor, + self.last_layer_size, + self.is_sorted, + ) + + short_table = self.table_tensor.reshape(-1)[:-1].reshape(1, -1) + with self.assertRaisesRegex(RuntimeError, "table does not contain enough"): + torch.ops.deepmd.tabulate_fusion_se_atten( + short_table, + self.table_info_tensor, + self.em_x_tensor, + self.em_tensor, + self.two_embed_tensor, + self.last_layer_size, + self.is_sorted, + ) + + def test_fractional_regions_require_every_reachable_row(self) -> None: + self._assert_table_row_boundary( + [0, 1.2, 2.4, 1, 1, -1], short_rows=2, required_rows=3 + ) + + def test_max_equal_to_upper_requires_high_tail_row(self) -> None: + self._assert_table_row_boundary( + [0, 1, 1, 1, 1, -1], short_rows=1, required_rows=2 + ) + + def _assert_table_row_boundary( + self, table_info: list[float], short_rows: int, required_rows: int + ) -> None: + last_layer_size = 2 + coefficients_per_row = last_layer_size * 6 + table_info_tensor = torch.tensor( + table_info, dtype=self.table_tensor.dtype, device="cpu" + ) + em_x = torch.tensor( + [[table_info[2], table_info[2] + 1]], + dtype=self.table_tensor.dtype, + device=env.DEVICE, + ) + em = torch.zeros((1, 2, 4), dtype=self.table_tensor.dtype, device=env.DEVICE) + two_embed = torch.zeros( + (1, 2 * last_layer_size), + dtype=self.table_tensor.dtype, + device=env.DEVICE, + ) + with self.assertRaisesRegex(RuntimeError, "table does not contain enough"): + torch.ops.deepmd.tabulate_fusion_se_atten( + torch.zeros( + (short_rows, coefficients_per_row), + dtype=self.table_tensor.dtype, + device=env.DEVICE, + ), + table_info_tensor, + em_x, + em, + two_embed, + last_layer_size, + self.is_sorted, + ) + + descriptor = torch.ops.deepmd.tabulate_fusion_se_atten( + torch.zeros( + (required_rows, coefficients_per_row), + dtype=self.table_tensor.dtype, + device=env.DEVICE, + ), + table_info_tensor, + em_x, + em, + two_embed, + last_layer_size, + self.is_sorted, + )[0] + self.assertEqual(descriptor.shape, (1, 4, last_layer_size)) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/tf/test_tabulate_shape_validation.py b/source/tests/tf/test_tabulate_shape_validation.py new file mode 100644 index 0000000000..76eb7f2b6c --- /dev/null +++ b/source/tests/tf/test_tabulate_shape_validation.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.tf.env import ( + op_module, + tf, +) + + +class TestTabulateShapeValidation(unittest.TestCase): + def setUp(self) -> None: + self.table = tf.constant(np.zeros((2, 12)), dtype=tf.float64) + self.table_info = tf.constant([0, 1, 2, 1, 1, -1], dtype=tf.float64) + self.em_x = tf.constant([[0.25, 0.5, 0.75]], dtype=tf.float64) + self.em = tf.constant(np.zeros((1, 3, 4)), dtype=tf.float64) + + def test_rejects_short_two_embed(self) -> None: + descriptor = op_module.tabulate_fusion_se_atten( + self.table, + self.table_info, + self.em_x, + self.em, + tf.constant([[1.0]], dtype=tf.float64), + last_layer_size=2, + is_sorted=True, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "two_embed must be rank 2" + ), + tf.Session() as sess, + ): + sess.run(descriptor) + + def test_accepts_flattened_em_x_layout(self) -> None: + descriptor = op_module.tabulate_fusion_se_a( + self.table, + self.table_info, + tf.reshape(self.em_x, [-1, 1]), + self.em, + last_layer_size=2, + ) + with tf.Session() as sess: + self.assertEqual(sess.run(descriptor).shape, (1, 4, 2)) + + def test_rejects_mismatched_em_x(self) -> None: + descriptor = op_module.tabulate_fusion_se_a( + self.table, + self.table_info, + self.em_x[:, :2], + self.em, + last_layer_size=2, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "em_x must be rank 2" + ), + tf.Session() as sess, + ): + sess.run(descriptor) + + def test_rejects_zero_neighbors_for_nonempty_atoms(self) -> None: + empty_em_x = tf.zeros((1, 0), dtype=tf.float64) + empty_em = tf.zeros((1, 0, 4), dtype=tf.float64) + descriptor = op_module.tabulate_fusion_se_a( + self.table, + self.table_info, + empty_em_x, + empty_em, + last_layer_size=2, + ) + gradients = op_module.tabulate_fusion_se_a_grad( + self.table, + self.table_info, + empty_em_x, + empty_em, + tf.zeros((1, 4, 2), dtype=tf.float64), + tf.zeros((1, 4, 2), dtype=tf.float64), + ) + for operation in (descriptor, gradients): + with ( + self.subTest(operation=operation), + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "at least one neighbor" + ), + tf.Session() as sess, + ): + sess.run(operation) + + def test_rejects_mismatched_gradient_shape(self) -> None: + gradients = op_module.tabulate_fusion_se_a_grad( + self.table, + self.table_info, + self.em_x, + self.em, + tf.zeros((1, 4, 1), dtype=tf.float64), + tf.zeros((1, 4, 2), dtype=tf.float64), + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "dy has an unexpected shape" + ), + tf.Session() as sess, + ): + sess.run(gradients) + + def test_rejects_short_table_info(self) -> None: + descriptor = op_module.tabulate_fusion_se_a( + self.table, + self.table_info[:4], + self.em_x, + self.em, + last_layer_size=2, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "table_info must contain" + ), + tf.Session() as sess, + ): + sess.run(descriptor) + + def test_rejects_short_table(self) -> None: + descriptor = op_module.tabulate_fusion_se_a( + self.table[:, :-1], + self.table_info, + self.em_x, + self.em, + last_layer_size=2, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "table does not contain enough" + ), + tf.Session() as sess, + ): + sess.run(descriptor) + + def test_fractional_regions_require_every_reachable_row(self) -> None: + table_info = tf.constant([0, 1.2, 2.4, 1, 1, -1], dtype=tf.float64) + em_x = tf.constant([[2.4, 3.0]], dtype=tf.float64) + em = tf.zeros((1, 2, 4), dtype=tf.float64) + short_descriptor = op_module.tabulate_fusion_se_a( + tf.zeros((2, 12), dtype=tf.float64), + table_info, + em_x, + em, + last_layer_size=2, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "table does not contain enough" + ), + tf.Session() as sess, + ): + sess.run(short_descriptor) + + descriptor = op_module.tabulate_fusion_se_a( + tf.zeros((3, 12), dtype=tf.float64), + table_info, + em_x, + em, + last_layer_size=2, + ) + with tf.Session() as sess: + self.assertEqual(sess.run(descriptor).shape, (1, 4, 2)) + + def test_max_equal_to_upper_requires_high_tail_row(self) -> None: + table_info = tf.constant([0, 1, 1, 1, 1, -1], dtype=tf.float64) + em_x = tf.constant([[1.0, 2.0]], dtype=tf.float64) + em = tf.zeros((1, 2, 4), dtype=tf.float64) + short_descriptor = op_module.tabulate_fusion_se_a( + tf.zeros((1, 12), dtype=tf.float64), + table_info, + em_x, + em, + last_layer_size=2, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "table does not contain enough" + ), + tf.Session() as sess, + ): + sess.run(short_descriptor) + + descriptor = op_module.tabulate_fusion_se_a( + tf.zeros((2, 12), dtype=tf.float64), + table_info, + em_x, + em, + last_layer_size=2, + ) + with tf.Session() as sess: + self.assertEqual(sess.run(descriptor).shape, (1, 4, 2)) + + +if __name__ == "__main__": + unittest.main()