Use SafeInt for size arithmetic in CPU tensor operators to prevent overflow

onnxruntime/core/providers/cpu/nn/tfidfvectorizer.cc

@@ -382,7 +382,7 @@ Status TfIdfVectorizer::Compute(OpKernelContext* ctx) const {
     // TfidfVectorizer returns a zero tensor of shape
     // {b_dim, output_size} when b_dim is the number of received observations
     // and output_size the is the maximum value in ngram_indexes attribute plus 1.
-    memset(output_data, 0, static_cast<size_t>(output_shape.Size() * sizeof(float)));
+    memset(output_data, 0, SafeInt<size_t>(output_shape.Size()) * sizeof(float));
     return Status::OK();
   }
 

onnxruntime/core/providers/cpu/quantization/qlinearconv.cc

@@ -251,7 +251,7 @@ class QLinearConv : public OpKernel {
         //
         // Note: The size of this buffer is less than or equal to the size of the original
         // weight tensor, so the allocation size is guaranteed to fit inside size_t.
-        auto* group_reordered_W = static_cast<int8_t*>(alloc->Alloc(group_output_channels * group_input_channels * kernel_size));
+        auto* group_reordered_W = static_cast<int8_t*>(alloc->Alloc(SafeInt<size_t>(group_output_channels) * group_input_channels * kernel_size));
         BufferUniquePtr group_reordered_W_buffer(group_reordered_W, BufferDeleter(alloc));
 
         const size_t W_offset = group_output_channels * kernel_dim;
@@ -439,7 +439,7 @@ Status QLinearConv<ActType>::PrePack(const Tensor& tensor, int input_idx, Alloca
       //
       // Note: The size of this buffer is less than or equal to the size of the original
       // weight tensor, so the allocation size is guaranteed to fit inside size_t.
-      auto group_reordered_W_buffer = IAllocator::MakeUniquePtr<void>(alloc, group_output_channels * group_input_channels * kernel_size, true);
+      auto group_reordered_W_buffer = IAllocator::MakeUniquePtr<void>(alloc, SafeInt<size_t>(group_output_channels) * group_input_channels * kernel_size, true);
       auto* group_reordered_W = static_cast<uint8_t*>(group_reordered_W_buffer.get());
 
       const size_t W_offset = group_output_channels * kernel_dim;

onnxruntime/core/providers/cpu/quantization/quantize_linear.cc

@@ -872,7 +872,7 @@ void ComputeLoop(OpKernelContext* ctx, const InT* input, const InT* scale, const
         output_index += static_cast<size_t>(N);                                                                   \
       }                                                                                                           \
     }                                                                                                             \
-    assert(output_index == static_cast<size_t>(M * K * N));                                                       \
+    assert(output_index == SafeInt<size_t>(M) * K * N);                                                       \
   }
 
 DEFINE_COMPUTE_LOOP_FP32_TO_SUB_BYTE(Int4x2, ParQuantizeLinearStdS4, 2)
@@ -890,7 +890,7 @@ DEFINE_COMPUTE_LOOP_FP32_TO_SUB_BYTE(UInt2x4, ParQuantizeLinearStdU2, 4)
                                              int64_t K, int64_t N, bool saturate) {                                    \
     ORT_UNUSED_PARAMETER(saturate);                                                                                    \
                                                                                                                        \
-    size_t total_size = static_cast<size_t>(M * K * N);                                                                \
+    size_t total_size = SafeInt<size_t>(M) * K * N;                                                                \
     auto tmp_buf = std::make_unique<SUB_BYTE_TYPE::UnpackedType[]>(total_size);                                        \
     size_t tmp_buf_index = 0;                                                                                          \
     constexpr size_t shift_bits = (ELEMENTS_PER_BYTE == 2) ? 1 : 2; /* log2(ELEMENTS_PER_BYTE) */                      \

onnxruntime/core/providers/cpu/sequence/sequence_ops.cc

@@ -4,6 +4,7 @@
 #include "core/providers/cpu/sequence/sequence_ops.h"
 
 #include "core/common/narrow.h"
+#include "core/common/safeint.h"
 #include "core/framework/tensorprotoutils.h"
 #include "core/framework/TensorSeq.h"
 #include "core/framework/op_kernel_type_control_utils.h"
@@ -517,7 +518,7 @@
     void* output_data = output_tensor.MutableDataRaw();
 
     const auto M = before_dims;
-    const auto* A = static_cast<const char*>(input_data) + static_cast<size_t>(input_offset * element_size);
+    const auto* A = static_cast<const char*>(input_data) + SafeInt<size_t>(input_offset) * element_size;
     const auto lda = after_dims_including_split_axis;
     auto* B = output_data;
 
@@ -528,7 +529,7 @@
       const auto* src = reinterpret_cast<const std::string*>(A);
       auto* dst = reinterpret_cast<std::string*>(B);
       if (lda == N) {
-        copy_data<std::string>(src, dst, static_cast<size_t>(M * N));
+        copy_data<std::string>(src, dst, SafeInt<size_t>(M) * N);
       } else {
         size_t lda_offset = 0;
         size_t ldb_offset = 0;

onnxruntime/core/providers/cpu/tensor/affine_grid.cc

@@ -4,6 +4,7 @@
 #include "core/providers/cpu/tensor/affine_grid.h"
 
 #include "core/common/common.h"
+#include "core/common/safeint.h"
 #include "core/providers/op_kernel_type_control.h"
 #include "core/util/math_cpuonly.h"
 #include <Eigen/Dense>
@@ -78,9 +79,9 @@
   const Eigen::Matrix<T, 2, 2, option> theta_R{{theta_data[0], theta_data[1]}, {theta_data[3], theta_data[4]}};  // 2x2
   const Eigen::Array<T, 2, 1> theta_T(theta_data[2], theta_data[5]);                                             // 2x1
 
-  auto grid_batch_offset = batch_num * H * W * 2;
+  auto grid_batch_offset = SafeInt<size_t>(batch_num) * H * W * 2;
   T* grid_data = grid->MutableData<T>() + grid_batch_offset;
-  Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 2, option>> grid_matrix(grid_data, narrow<size_t>(H * W), 2);
+  Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 2, option>> grid_matrix(grid_data, SafeInt<size_t>(H) * W, 2);
   grid_matrix = ((theta_R * base_grid_transposed).array().colwise() + theta_T).matrix().transpose();  // ((2x2 * 2xN).array().colwise() + 2x1).matrix().transpose() => Nx2
 }
 
@@ -97,9 +98,9 @@
 
   const Eigen::Array<T, 3, 1> theta_T(theta_data[3], theta_data[7], theta_data[11]);  // 3x1
 
-  auto grid_batch_offset = batch_num * D * H * W * 3;
+  auto grid_batch_offset = SafeInt<size_t>(batch_num) * D * H * W * 3;
   T* grid_data = grid->MutableData<T>() + grid_batch_offset;
-  Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 3, option>> grid_matrix(grid_data, narrow<size_t>(D * H * W), 3);
+  Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 3, option>> grid_matrix(grid_data, SafeInt<size_t>(D) * H * W, 3);
   grid_matrix = ((theta_R * base_grid_transposed).array().colwise() + theta_T).matrix().transpose();
 }
 

onnxruntime/core/providers/cpu/tensor/grid_sample.cc

@@ -3,6 +3,7 @@
 
 #include <vector>
 
+#include "core/common/safeint.h"
 #include "core/providers/cpu/tensor/grid_sample.h"
 #include "core/framework/element_type_lists.h"
 #include "core/framework/TensorSeq.h"
@@ -379,12 +380,12 @@ Status GridSample<T>::Compute(OpKernelContext* context) const {
       }
     } else {
       for (int64_t n = 0; n < N; n++) {
-        const T* grid_data = grid->Data<T>() + n * (H_out * W_out) * 2;
+        const T* grid_data = grid->Data<T>() + SafeInt<size_t>(n) * H_out * W_out * 2;
         concurrency::ThreadPool::TrySimpleParallelFor(
             tp, onnxruntime::narrow<std::ptrdiff_t>(C),
             [&](std::ptrdiff_t c) {
-              const T* X_data = input->Data<T>() + (n * C + c) * (H_in * W_in);
-              T* Y_data = Y.MutableData<T>() + (n * C + c) * (H_out * W_out);
+              const T* X_data = input->Data<T>() + SafeInt<size_t>(n * C + c) * H_in * W_in;
+              T* Y_data = Y.MutableData<T>() + SafeInt<size_t>(n * C + c) * H_out * W_out;
 
               for (int64_t oy = 0; oy < H_out; oy++) {
                 for (int64_t ox = 0; ox < W_out; ox++) {
@@ -469,12 +470,12 @@ Status GridSample<T>::Compute(OpKernelContext* context) const {
 
     concurrency::ThreadPool* tp = D_out * H_out * W_out > 64 ? context->GetOperatorThreadPool() : nullptr;
     for (int64_t n = 0; n < N; n++) {
-      const T* grid_data = grid->Data<T>() + n * (D_out * H_out * W_out) * 3;
+      const T* grid_data = grid->Data<T>() + SafeInt<size_t>(n) * D_out * H_out * W_out * 3;
       concurrency::ThreadPool::TrySimpleParallelFor(
           tp, onnxruntime::narrow<std::ptrdiff_t>(C),
           [&](std::ptrdiff_t c) {
-            const T* X_data = input->Data<T>() + (n * C + c) * (D_in * H_in * W_in);
-            T* Y_data = Y.MutableData<T>() + (n * C + c) * (D_out * H_out * W_out);
+            const T* X_data = input->Data<T>() + SafeInt<size_t>(n * C + c) * D_in * H_in * W_in;
+            T* Y_data = Y.MutableData<T>() + SafeInt<size_t>(n * C + c) * D_out * H_out * W_out;
 
             for (int64_t oz = 0; oz < D_out; oz++) {
               for (int64_t oy = 0; oy < H_out; oy++) {

onnxruntime/core/providers/cpu/tensor/upsample_antialias.h

@@ -17,6 +17,7 @@
 #ifndef SHARED_PROVIDER
 #include "core/framework/op_kernel.h"
 #endif
+#include "core/common/safeint.h"
 #include "core/providers/cpu/tensor/upsamplebase.h"
 
 namespace onnxruntime {
@@ -130,7 +131,7 @@ void SetupUpsampleFilterAntiAlias(FilterParamsAntiAlias<T>& p,
     float support = (scale >= 1.0f) ? (p.support_size * 0.5f) * scale : p.support_size * 0.5f;
 
     int32_t window_size = narrow<int32_t>(ceilf(support)) * 2 + 1;
-    const size_t scale_buffer_size = narrow<size_t>(window_size * output_size);
+    const size_t scale_buffer_size = SafeInt<size_t>(window_size) * output_size;
 
     param_base.weight_coefficients = IAllocator::MakeUniquePtr<T>(alloc, scale_buffer_size);
     // Get pointers to appropriate memory locations in the scratch buffer
@@ -260,7 +261,7 @@ void ComputeInterpolationAtLevel1(int64_t num_channels, int64_t input_height, in
         InputType* Ydata = Ydata_span.data() + y_start;
         // no need to do scale
         if (output_width == input_width) {
-          std::copy_n(Xdata_span.begin() + narrow<size_t>(x_start), narrow<size_t>(output_height * output_width),
+          std::copy_n(Xdata_span.begin() + narrow<size_t>(x_start), SafeInt<size_t>(output_height) * output_width,
                       Ydata_span.begin() + narrow<size_t>(y_start));
           return;
         }
@@ -330,7 +331,7 @@ void ComputeInterpolationAtLevel2(int64_t num_channels, int64_t input_height, in
           InputType* Ydata = Ydata_span.data() + y_start;
 
           if (output_height == input_height) {
-            std::copy_n(Xdata_span.begin() + narrow<size_t>(x_start), narrow<size_t>(output_height * output_width),
+            std::copy_n(Xdata_span.begin() + narrow<size_t>(x_start), SafeInt<size_t>(output_height) * output_width,
                         Ydata_span.begin() + narrow<size_t>(y_start));
             return;
           }
@@ -368,8 +369,8 @@ void ComputeInterpolationAtLevel2(int64_t num_channels, int64_t input_height, in
         [&](std::ptrdiff_t first, std::ptrdiff_t last) {
           if (output_height == input_height) {
             auto workload_in_thread = narrow<size_t>(last) - narrow<size_t>(first);
-            std::copy_n(Xdata_span.begin() + narrow<size_t>(first * input_width), narrow<size_t>(workload_in_thread * output_width),
-                        Ydata_span.begin() + narrow<size_t>(first * output_width));
+            std::copy_n(Xdata_span.begin() + SafeInt<size_t>(first) * input_width, SafeInt<size_t>(workload_in_thread) * output_width,
+                        Ydata_span.begin() + SafeInt<size_t>(first) * output_width);
             return;
           }
 
@@ -440,7 +441,7 @@ void HandleExtrapolation(int64_t num_channels,
 
         for (int64_t z : p.dim_z.out_of_bound_idx) {
           InputType* Ydata_offset = Ydata_base_nc + z * output_height * output_width;
-          std::fill_n(Ydata_offset, narrow<size_t>(output_height * output_width), static_cast<InputType>(extrapolation_value));
+          std::fill_n(Ydata_offset, SafeInt<size_t>(output_height) * output_width, static_cast<InputType>(extrapolation_value));
         }
       });
 }
@@ -460,14 +461,14 @@ void UpsampleBaseAntiAlias(FilterParamsAntiAlias<T1>& p,
                            AllocatorPtr& alloc,
                            concurrency::ThreadPool* tp) {
   IAllocatorUniquePtr<T> image_temp_buffer = IAllocator::MakeUniquePtr<T>(
-      alloc, static_cast<size_t>(input_height * output_width * num_channels));
+      alloc, SafeInt<size_t>(input_height) * output_width * num_channels);
 
   for (int64_t n = 0; n < batch_size; ++n) {
     {
       // horizon interpolate
       auto xdata_span = gsl::make_span(Xdata_base + n * (input_height * num_channels * input_width),
-                                       narrow<size_t>(input_height * num_channels * input_width));
-      auto ydata_span = gsl::make_span(image_temp_buffer.get(), narrow<size_t>(input_height * num_channels * output_width));
+                                       SafeInt<size_t>(input_height) * num_channels * input_width);
+      auto ydata_span = gsl::make_span(image_temp_buffer.get(), SafeInt<size_t>(input_height) * num_channels * output_width);
 
       // This computes only the width direction.Thus height keeps unchanged.
       ComputeInterpolationAtLevel1(num_channels, input_height, input_width, input_height, output_width,
@@ -476,17 +477,17 @@ void UpsampleBaseAntiAlias(FilterParamsAntiAlias<T1>& p,
     {
       // vertical interpolate
       auto xdata_span = gsl::make_span<const T>(image_temp_buffer.get(),
-                                                narrow<size_t>(input_height * num_channels * output_width));
+                                                SafeInt<size_t>(input_height) * num_channels * output_width);
       auto ydata_span = gsl::make_span<T>(Ydata_base + n * (output_height * num_channels * output_width),
-                                          narrow<size_t>(output_height * num_channels * output_width));
+                                          SafeInt<size_t>(output_height) * num_channels * output_width);
 
       ComputeInterpolationAtLevel2(num_channels, input_height, output_width, output_height, output_width,
                                    xdata_span, ydata_span, p, p.dim_y, tp);
     }
   }
   if (use_extrapolation) {
     auto ydata_span = gsl::make_span<T>(Ydata_base,
-                                        narrow<size_t>(batch_size * output_height * num_channels * output_width));
+                                        SafeInt<size_t>(batch_size) * output_height * num_channels * output_width);
     HandleExtrapolation(batch_size * num_channels, output_height, output_width, 1,
                         extrapolation_value, ydata_span, p, tp);
   }
@@ -596,14 +597,14 @@ void NhwcUpsampleBasicAntiAlias(FilterParamsAntiAlias<T1>& p,
                                 AllocatorPtr& alloc,
                                 concurrency::ThreadPool* tp) {
   IAllocatorUniquePtr<T> image_temp_buffer = IAllocator::MakeUniquePtr<T>(
-      alloc, static_cast<size_t>(input_height * output_width * num_channels));
+      alloc, SafeInt<size_t>(input_height) * output_width * num_channels);
 
   for (int64_t n = 0; n < batch_size; ++n) {
     // horizon interpolate
     {
       auto xdata_span = gsl::make_span(Xdata_base + n * (input_height * num_channels * input_width),
-                                       narrow<size_t>(input_height * num_channels * input_width));
-      auto ydata_span = gsl::make_span(image_temp_buffer.get(), narrow<size_t>(input_height * num_channels * output_width));
+                                       SafeInt<size_t>(input_height) * num_channels * input_width);
+      auto ydata_span = gsl::make_span(image_temp_buffer.get(), SafeInt<size_t>(input_height) * num_channels * output_width);
 
       ComputeInterpolationAtLevel2(input_height, input_width, num_channels, output_width, num_channels,
                                    xdata_span, ydata_span, p, p.dim_x, tp);
@@ -613,9 +614,9 @@ void NhwcUpsampleBasicAntiAlias(FilterParamsAntiAlias<T1>& p,
     {
       // vertical interpolate
       auto xdata_span = gsl::make_span<const T>(image_temp_buffer.get(),
-                                                narrow<size_t>(input_height * num_channels * output_width));
+                                                SafeInt<size_t>(input_height) * num_channels * output_width);
       auto ydata_span = gsl::make_span<T>(Ydata_base + n * (output_height * num_channels * output_width),
-                                          narrow<size_t>(output_height * num_channels * output_width));
+                                          SafeInt<size_t>(output_height) * num_channels * output_width);
 
       ComputeInterpolationAtLevel2(1, input_height, output_width * num_channels, output_height, output_width * num_channels,
                                    xdata_span, ydata_span, p, p.dim_y, tp);
@@ -624,7 +625,7 @@ void NhwcUpsampleBasicAntiAlias(FilterParamsAntiAlias<T1>& p,
 
   if (use_extrapolation) {
     auto ydata_span = gsl::make_span<T>(Ydata_base,
-                                        narrow<size_t>(batch_size * output_height * num_channels * output_width));
+                                        SafeInt<size_t>(batch_size) * output_height * num_channels * output_width);
     HandleExtrapolation(batch_size * num_channels, output_height, output_width, 1,
                         extrapolation_value, ydata_span, p, tp);
   }
@@ -692,8 +693,8 @@ void UpsampleTrilinearAntiAlias(int64_t batch_size,
                                alloc, get_original_coordinate, exclude_outside, true);
 
   IAllocatorUniquePtr<T> image_temp_buffer = IAllocator::MakeUniquePtr<T>(
-      alloc, static_cast<size_t>(batch_size * output_height * output_width *
-                                 input_depth * num_channels));
+      alloc, SafeInt<size_t>(batch_size) * output_height * output_width *
+                 input_depth * num_channels);
 
   UpsampleBaseAntiAlias<T>(p, batch_size, num_channels * input_depth, input_height, input_width, output_height, output_width,
                            false, extrapolation_value,
@@ -706,9 +707,9 @@ void UpsampleTrilinearAntiAlias(int64_t batch_size,
     {
       // depth interpolate
       auto xdata_span = gsl::make_span<const T>(image_temp_buffer.get() + n * (output_height * num_channels * output_width * input_depth),
-                                                narrow<size_t>(output_height * num_channels * output_width * input_depth));
+                                                SafeInt<size_t>(output_height) * num_channels * output_width * input_depth);
       auto ydata_span = gsl::make_span<T>(Ydata_base + n * (output_height * num_channels * output_width * output_depth),
-                                          narrow<size_t>(output_height * num_channels * output_width * output_depth));
+                                          SafeInt<size_t>(output_height) * num_channels * output_width * output_depth);
 
       ComputeInterpolationAtLevel2(m_channel_size, input_depth, output_height * output_width, output_depth, output_height * output_width,
                                    xdata_span, ydata_span, p, p.dim_z, tp);
@@ -717,7 +718,7 @@ void UpsampleTrilinearAntiAlias(int64_t batch_size,
 
   if (use_extrapolation) {
     auto ydata_span = gsl::make_span<T>(Ydata_base,
-                                        narrow<size_t>(batch_size * output_height * num_channels * output_width * output_depth));
+                                        SafeInt<size_t>(batch_size) * output_height * num_channels * output_width * output_depth);
     HandleExtrapolation(batch_size * num_channels, output_height, output_width, output_depth,
                         extrapolation_value, ydata_span, p, tp);
   }