Harden shape inference for contrib ops

onnxruntime/core/graph/contrib_ops/bert_defs.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include <algorithm>
+#include <limits>
 #include <string>
 #include <vector>
 
@@ -303,6 +304,9 @@ void BaseGroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceConte
     const auto* total_sequence_length_data = ctx.getInputData(6);
     if (total_sequence_length_data != nullptr) {
       const auto& data = ParseData<int32_t>(total_sequence_length_data);
+      if (data.empty()) {
+        fail_shape_inference("total_sequence_length input must contain a single element");
+      }
       total_sequence_length_value = static_cast<int64_t>(data[0]);
     }
 
@@ -2881,7 +2885,12 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
             fail_shape_inference("CausalConvWithState: channels_last must be 0 or 1, got ",
                                  channels_last);
           }
-          if (channels_last == 1 && getAttribute(ctx, "ndim", 1) != 1) {
+
+          const int64_t ndim = getAttribute(ctx, "ndim", 1);
+          if (ndim < 1 || ndim > 3) {
+            fail_shape_inference("CausalConvWithState: ndim must be 1, 2, or 3, got ", ndim);
+          }
+          if (channels_last == 1 && ndim != 1) {
             fail_shape_inference("CausalConvWithState: channels_last requires ndim = 1");
           }
 
@@ -2895,13 +2904,20 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
           if (hasInputShape(ctx, 0) && hasInputShape(ctx, 1)) {
             auto& input_shape = getInputShape(ctx, 0);
             auto& weight_shape = getInputShape(ctx, 1);
-            if (input_shape.dim_size() < 2) {
-              fail_shape_inference("CausalConvWithState: input must have rank >= 2");
+            // weight is always channels-first: (channels, 1, k_1, ..., k_ndim), rank == ndim + 2.
+            if (weight_shape.dim_size() != ndim + 2) {
+              fail_shape_inference("CausalConvWithState: weight must have rank ndim + 2 (",
+                                   ndim + 2, "), got rank ", weight_shape.dim_size());
             }
-            if (weight_shape.dim_size() < 2) {
-              fail_shape_inference("CausalConvWithState: weight must have rank >= 2");
+            if (channels_last == 1) {
+              // (batch_size, sequence_length, d_1, ..., d_n). Check the lower bound.
+              if (input_shape.dim_size() < 3) {
+                fail_shape_inference("CausalConvWithState: channels_last input must have rank >= 3");
+              }
+            } else if (input_shape.dim_size() != ndim + 2) {
+              fail_shape_inference("CausalConvWithState: input must have rank ndim + 2 (",
+                                   ndim + 2, "), got rank ", input_shape.dim_size());
             }
-            int64_t ndim = getAttribute(ctx, "ndim", 1);
             // (kernel_size - 1) * dilation, or an unset dim when kernel_size is symbolic.
             const int last_kernel_dim = weight_shape.dim_size() - 1;
             TensorShapeProto::Dimension state_length;
@@ -3555,9 +3571,31 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
               const int64_t head_size_qk = query_shape.dim(token_dims + 1).dim_value();
               const int64_t num_heads_v = value_shape.dim(token_dims).dim_value();
               const int64_t head_size_v = value_shape.dim(token_dims + 1).dim_value();
-              width->set_dim_value(state_update_capacity *
-                                   (num_heads_v + num_heads_k * head_size_qk +
-                                    num_heads_v * head_size_v));
+              if (num_heads_k <= 0 || head_size_qk <= 0 || num_heads_v <= 0 || head_size_v <= 0) {
+                fail_shape_inference(
+                    "GatedDeltaNet: head counts and head sizes must be positive");
+              }
+
+              constexpr int64_t max_dimension = std::numeric_limits<int64_t>::max();
+              if (num_heads_k > max_dimension / head_size_qk ||
+                  num_heads_v > max_dimension / head_size_v) {
+                fail_shape_inference("GatedDeltaNet: state_update width overflows int64");
+              }
+              const int64_t key_width = num_heads_k * head_size_qk;
+              const int64_t value_width = num_heads_v * head_size_v;
+              if (num_heads_v > max_dimension - key_width) {
+                fail_shape_inference("GatedDeltaNet: state_update width overflows int64");
+              }
+              const int64_t width_without_capacity = num_heads_v + key_width;
+              if (value_width > max_dimension - width_without_capacity) {
+                fail_shape_inference("GatedDeltaNet: state_update width overflows int64");
+              }
+              const int64_t per_token_width = width_without_capacity + value_width;
+              if (state_update_capacity > 0 &&
+                  per_token_width > max_dimension / state_update_capacity) {
+                fail_shape_inference("GatedDeltaNet: state_update width overflows int64");
+              }
+              width->set_dim_value(state_update_capacity * per_token_width);
             }
             updateOutputShape(ctx, 2, capsule_shape);
           }