Harden ConvTranspose pad computation with SafeInt and consistency checks

onnxruntime/core/providers/common.h

@@ -155,9 +155,11 @@ inline Status ComputePadAndOutputShape(const int64_t in_dim,
   return Status::OK();
 }
 
-constexpr inline int64_t ComputeTotalPad(int64_t in_size, int64_t stride, int64_t adj,
-                                         int64_t kernel, int64_t dilation, int64_t out_size) {
-  return std::max<int64_t>(0, (in_size - 1) * stride + adj + (kernel - 1) * dilation + 1 - out_size);
+inline int64_t ComputeTotalPad(int64_t in_size, int64_t stride, int64_t adj,
+                               int64_t kernel, int64_t dilation, int64_t out_size) {
+  SafeInt<int64_t> safe_pad = (SafeInt<int64_t>(in_size) - 1) * stride + adj +
+                              (SafeInt<int64_t>(kernel) - 1) * dilation + 1 - out_size;
+  return std::max<int64_t>(0, safe_pad);
 }
 
 inline void DistributePadding(AutoPadType pad_type, const int64_t& total_pad,

onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h

@@ -305,14 +305,50 @@ struct ConvTransposeAttributes : public ConvAttributes {
       int64_t* out_size) const {
     // Output shape is explicitly provided - pad values will have to be computed
     if (*out_size != -1) {
-      if (*out_size < 0) {
+      if (*out_size <= 0) {
         return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                               "Explicit output size is negative: ", *out_size);
+                               "Explicit output size must be positive. Got: ", *out_size);
+      }
+      if (in_size <= 0) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Input spatial dimension must be positive. Got: ", in_size);
+      }
+      if (stride <= 0) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Stride must be positive. Got: ", stride);
+      }
+      if (kernel <= 0) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Kernel size must be positive. Got: ", kernel);
+      }
+      if (dilation <= 0) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Dilation must be positive. Got: ", dilation);
+      }
+      if (adj < 0) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Output padding must be non-negative. Got: ", adj);
       }
       // total pad
       auto total_pad = ComputeTotalPad(in_size, stride, adj,
                                        kernel, dilation, *out_size);
       DistributePadding(pad_type, total_pad, *pad_head, *pad_tail);
+
+      // Verify that the forward-conv re-derivation of input size from the output size and pads
+      // is consistent with the actual input size. Col2im re-derives the input spatial extent as:
+      //   derived_in = (out_size + pad_head + pad_tail - dkernel) / stride + 1
+      // If this exceeds in_size, Col2im would read past the col_buffer allocation.
+      // Note: derived_in < in_size is algebraically unreachable when adj >= 0 and total_pad >= 0,
+      // so this check is effectively one-sided (catches derived_in > in_size from oversized out_size).
+      SafeInt<int64_t> dkernel = (SafeInt<int64_t>(kernel) - 1) * dilation + 1;
+      int64_t derived_in = (SafeInt<int64_t>(*out_size) + *pad_head + *pad_tail - dkernel) / stride + 1;
+      if (derived_in != in_size) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Explicit output_shape is inconsistent with input spatial dimensions"
+                               " and convolution parameters. "
+                               "Expected input size ",
+                               derived_in, " but got ", in_size,
+                               " (output_size=", *out_size, ", kernel=", kernel,
+                               ", stride=", stride, ", dilation=", dilation,
+                               ", output_padding=", adj, ").");
+      }
       return Status::OK();
     }
 
@@ -346,6 +382,20 @@ struct ConvTransposeAttributes : public ConvAttributes {
     *out_size = SafeInt<int64_t>(in_size - 1) * stride + adj +
                 SafeInt<int64_t>(kernel - 1) * dilation + 1 -
                 *pad_head - *pad_tail;
+
+    // Same consistency check as the explicit output_shape path: verify the forward-conv
+    // re-derivation of input size matches in_size. When output_padding (adj) >= stride
+    // (possible when dilation > stride passes the adj < max(stride, dilation) check),
+    // Col2im would compute a larger input extent and read past the col_buffer.
+    SafeInt<int64_t> dkernel2 = (SafeInt<int64_t>(kernel) - 1) * dilation + 1;
+    int64_t derived_in = (SafeInt<int64_t>(*out_size) + *pad_head + *pad_tail - dkernel2) / stride + 1;
+    if (derived_in != in_size) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "Computed output shape is inconsistent with input spatial dimensions. "
+                             "output_padding (",
+                             adj, ") may be too large for stride (", stride,
+                             "). Expected input size ", derived_in, " but got ", in_size, ".");
+    }
     return Status::OK();
   }
 };