Validate initializer data length and guard element-count computation in contrib Range shape inference (microsoft#29265)

### Summary

The `com.microsoft` Range operator's shape-inference helper read a fixed
number of bytes from an initializer's `raw_data` without first checking
the buffer length, and the element-count computation could pass
non-finite or out-of-range values to an `int64` cast. This change adds
the missing validations and aligns the shape-inference and CPU kernel
paths.

### Changes

- `GetFirstElement` now checks `raw_data` length is at least the element
size before reading, and reads via `std::memcpy` into an aligned local.
- `CalcRangeDim` and the CPU kernel `ComputeRange` now reject non-finite
computed counts, handle non-positive counts before the `int64` cast, and
reject counts that are not representable as `int64` (`>= 2^63`). Both
paths use identical messages and semantics.
- The output dimension for empty/backward ranges is clamped to 0 in
shape inference to match the kernel.

### Tests

Added contrib Range model-load regression tests in `range_test.cc`
covering:
- truncated `raw_data` for `start`, `limit`, and `delta` (double, plus
float and int64 element types),
- zero delta,
- a finite-but-too-large element count,
- an exact-size success boundary,
- backward-range zero-dimension inference.

Tests assert on `Status` (safe for no-exception builds) and are guarded
by `#ifndef DISABLE_CONTRIB_OPS` (throwing cases additionally by
`!defined(ORT_NO_EXCEPTIONS)`). 21/21 `RangeTest` cases pass locally.

### Follow-up

`int16` inputs are currently supported only via `raw_data` (there is no
`get_data<int16_t>` specialization for the non-raw-data path); this is
left as a separate follow-up.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: titaiwangms

onnxruntime/core/graph/contrib_ops/range_schema_defs.cc

@@ -7,6 +7,7 @@
 #include "core/graph/constants.h"
 #include "core/graph/op.h"
 #include <cmath>
+#include <cstring>
 #include <type_traits>
 
 namespace onnxruntime {
@@ -42,6 +43,13 @@ int64_t get_data<int64_t>(const TensorProto* shapeInitializer) {
   fail_shape_inference("Can not get shape initializer data!");
 }
 
+template <>
+int16_t get_data<int16_t>(const TensorProto* shapeInitializer) {
+  // ONNX packs INT16 initializer values into the int32_data repeated field.
+  if (shapeInitializer->int32_data_size() > 0) return static_cast<int16_t>(shapeInitializer->int32_data(0));
+  fail_shape_inference("Can not get shape initializer data!");
+}
+
 template <>
 float get_data<float>(const TensorProto* shapeInitializer) {
   if (shapeInitializer->float_data_size() > 0) return shapeInitializer->float_data(0);
@@ -60,7 +68,13 @@ static T GetFirstElement(const TensorProto* shapeInitializer) {
 
   if (utils::HasRawData(*shapeInitializer)) {
     const std::string& bytes = shapeInitializer->raw_data();
-    return *reinterpret_cast<const T*>(bytes.c_str());
+    if (bytes.size() < sizeof(T)) {
+      fail_shape_inference("Range: raw_data size is smaller than the element size for the given data type.");
+    }
+    // std::string data is only char-aligned, so copy the bytes into a properly aligned value.
+    T value;
+    std::memcpy(&value, bytes.data(), sizeof(T));
+    return value;
   }
   return get_data<T>(shapeInitializer);
 }
@@ -75,7 +89,24 @@ static int64_t CalcRangeDim(const TensorProto* startShapeInitializer,
   if (delta == 0) {
     fail_shape_inference("delta in Range operator can not be zero!");
   }
-  return static_cast<int64_t>(ceil((1.0 * (limit - start)) / delta));
+  // Mirror the CPU kernel (core/providers/cpu/generator/range.cc ComputeRange) exactly so the
+  // two paths stay byte-consistent: clamp empty or backward ranges to 0 and promote the
+  // operands to double before the subtraction. Keep this expression identical to the kernel.
+  double count = ceil((static_cast<double>(limit) - static_cast<double>(start)) / static_cast<double>(delta));
+  if (!std::isfinite(count)) {
+    fail_shape_inference("Range: the computed number of elements is not a finite value.");
+  }
+  // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a
+  // large-magnitude negative count can never reach (and overflow) the int64 conversion below.
+  if (count <= 0) {
+    return 0;
+  }
+  // static_cast<double>(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not
+  // representable as int64_t, so reject any count at or above that boundary before the cast.
+  if (count >= 9223372036854775808.0) {
+    fail_shape_inference("Range: the computed number of elements exceeds the supported range.");
+  }
+  return static_cast<int64_t>(count);
 }
 
 static int64_t CalcResultDim(const TensorProto* startShapeInitializer,

onnxruntime/core/providers/cpu/generator/range.cc

@@ -6,10 +6,7 @@
 #include <cmath>
 
 #include "core/providers/op_kernel_type_control.h"
-// TODO: fix the warnings
-#if defined(_MSC_VER) && !defined(__clang__)
-#pragma warning(disable : 26451)
-#endif
+
 namespace onnxruntime {
 
 namespace op_kernel_type_control {
@@ -80,9 +77,26 @@ static Status ComputeRange(
   if (delta == T{0}) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "delta in Range operator can not be zero!");
   }
-  int64_t n = static_cast<int64_t>(ceil((1.0 * (limit - start)) / delta));
-  if (n <= 0)
-    n = 0;
+  // Compute the element count in double, mirroring the shape-inference path
+  // (core/graph/contrib_ops/range_schema_defs.cc CalcRangeDim) exactly. The operands are
+  // promoted to double before the subtraction so integral inputs cannot overflow in T.
+  double count = ceil((static_cast<double>(limit) - static_cast<double>(start)) / static_cast<double>(delta));
+  if (!std::isfinite(count)) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "Range: the computed number of elements is not a finite value.");
+  }
+  // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a
+  // large-magnitude negative count can never reach (and overflow) the int64 conversion.
+  int64_t n = 0;
+  if (count > 0) {
+    // static_cast<double>(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not
+    // representable as int64_t, so reject any count at or above that boundary before the cast.
+    if (count >= 9223372036854775808.0) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "Range: the computed number of elements exceeds the supported range.");
+    }
+    n = static_cast<int64_t>(count);
+  }
   TensorShape shape = {n};
   T* y = ctx->Output(0, shape)->MutableData<T>();
   for (int64_t i = 0; i < n; ++i) {

onnxruntime/core/providers/cuda/generator/range.cc

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <cmath>
+
 #include "core/providers/shared_library/provider_api.h"
 #include "core/providers/cuda/cuda_common.h"
 #include "range.h"
@@ -87,11 +89,28 @@ static Status ComputeRange(cudaStream_t stream, OpKernelContext* ctx) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "delta in Range operator can not be zero!");
   }
 
-  double num = (static_cast<double>(limit) - static_cast<double>(start)) / static_cast<double>(delta);
-  int count = static_cast<int>(ceil(num));
-  if (count <= 0)
-    count = 0;
-  TensorShape shape = {static_cast<int64_t>(count)};
+  // Compute the element count in double, mirroring the CPU kernel's guard structure and error
+  // messages (core/providers/cpu/generator/range.cc ComputeRange) and the shape-inference path
+  // (core/graph/contrib_ops/range_schema_defs.cc CalcRangeDim). The operands are
+  // promoted to double before the subtraction so integral inputs cannot overflow in T.
+  double num = ceil((static_cast<double>(limit) - static_cast<double>(start)) / static_cast<double>(delta));
+  if (!std::isfinite(num)) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "Range: the computed number of elements is not a finite value.");
+  }
+  // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a
+  // large-magnitude negative count can never reach (and overflow) the int64 conversion.
+  int64_t count = 0;
+  if (num > 0) {
+    // static_cast<double>(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not
+    // representable as int64_t, so reject any count at or above that boundary before the cast.
+    if (num >= 9223372036854775808.0) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "Range: the computed number of elements exceeds the supported range.");
+    }
+    count = static_cast<int64_t>(num);
+  }
+  TensorShape shape = {count};
   T* y = ctx->Output(0, shape)->MutableData<T>();
 
   if (count > 0) {

onnxruntime/core/providers/cuda/generator/range_impl.cu

@@ -13,23 +13,34 @@ namespace onnxruntime {
 namespace cuda {
 
 template <typename T>
-__global__ void RangeKernel(const T start, const T delta, const int count, T* output) {
-  int index = blockIdx.x * blockDim.x + threadIdx.x;
-  if (index < count) {
-    output[index] = start + delta * index;
+__global__ void RangeKernel(const T start, const T delta, const int64_t count, T* output) {
+  // Use a 64-bit grid-stride loop so counts larger than the launch grid (and larger than
+  // INT_MAX) are handled correctly without index truncation or grid-dimension overflow.
+  int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
+  for (int64_t i = index; i < count; i += stride) {
+    output[i] = start + delta * static_cast<T>(i);
   }
 }
 
 template <typename T>
-Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int count, T* output) {
-  constexpr int block_size = 256;
-  int grid_size = (count + block_size - 1) / block_size;
+Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int64_t count, T* output) {
+  if (count <= 0) {
+    return Status::OK();
+  }
+  constexpr int block_size = GridDim::maxThreadsPerBlock;
+  int64_t num_blocks = (count + block_size - 1) / block_size;
+  // CUDA limits the x-dimension of the launch grid to 2^31 - 1 blocks. Cap the grid to that
+  // maximum; the grid-stride loop in RangeKernel covers any remaining elements when the count
+  // requires more blocks than can be launched at once.
+  constexpr int64_t kMaxGridDimX = 2147483647;
+  int grid_size = static_cast<int>(num_blocks < kMaxGridDimX ? num_blocks : kMaxGridDimX);
   RangeKernel<T><<<grid_size, block_size, 0, stream>>>(start, delta, count, output);
   return CUDA_CALL(cudaGetLastError());
 }
 
 #define SPECIALIZED_IMPL(T) \
-  template Status RangeImpl<T>(cudaStream_t stream, const T start, const T delta, const int count, T* output);
+  template Status RangeImpl<T>(cudaStream_t stream, const T start, const T delta, const int64_t count, T* output);
 
 SPECIALIZED_IMPL(int16_t)
 SPECIALIZED_IMPL(int32_t)

onnxruntime/core/providers/cuda/generator/range_impl.h

@@ -8,7 +8,7 @@ namespace onnxruntime {
 namespace cuda {
 
 template <typename T>
-Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int count, T* output);
+Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int64_t count, T* output);
 
 }  // namespace cuda
 }  // namespace onnxruntime