fix comparison of narrow type with wide type in loop condition (pytorch#53951)

Summary:
fix Semmle warning: Comparison of narrow type with wide type in loop condition

For example there is below piece of code:
for (int i=0; i<array.size(); ++i) {}

The problem is that array.size() return type is size_t can be larger type than int depending on the implementation so there is chance that i overflows (for very large array that array size is beyond the range of integer) and this loop will never be terminated.

Pull Request resolved: pytorch#53951

Reviewed By: zou3519

Differential Revision: D27181495

Pulled By: malfet

fbshipit-source-id: 0612c5cedcdc656c193085e7fbb87dd163f20688

Author: frank-dong-ms

aten/src/ATen/TensorIterator.cpp

@@ -8,6 +8,8 @@
 #include <ATen/native/Resize.h>
 #include <ATen/TensorOperators.h>
 
+#include <c10/util/irange.h>
+
 namespace at {
 
 using DimMask = TensorIteratorBase::DimMask;
@@ -1392,7 +1394,7 @@ DimCounter::DimCounter(IntArrayRef shape, Range range)
   , offset(range.begin) {
   int64_t linear_offset = range.begin;
   int64_t ndim = values.size();
-  for (int dim = 0; dim < ndim; dim++) {
+  for (const auto dim : c10::irange(ndim)) {
     int64_t size = shape[dim];
     if (size > 0) {
       values[dim] = linear_offset % size;

aten/src/ATen/Utils.h

@@ -9,6 +9,7 @@
 #include <c10/util/accumulate.h>
 #include <c10/util/ArrayRef.h>
 #include <c10/util/Exception.h>
+#include <c10/util/irange.h>
 
 #include <algorithm>
 #include <sstream>
@@ -51,7 +52,7 @@ static inline TensorImpl* checked_dense_tensor_unwrap(const Tensor& expr, const
 static inline std::vector<TensorImpl*> checked_dense_tensor_list_unwrap(ArrayRef<Tensor> tensors, const char * name, int pos, DeviceType device_type, ScalarType scalar_type) {
   std::vector<TensorImpl*> unwrapped;
   unwrapped.reserve(tensors.size());
-  for (unsigned int i = 0; i < tensors.size(); ++i) {
+  for (const auto i : c10::irange(tensors.size())) {
     const auto& expr = tensors[i];
     if (expr.layout() != Layout::Strided) {
       AT_ERROR("Expected dense tensor but got ", expr.layout(),

aten/src/ATen/core/type.cpp

@@ -3,6 +3,7 @@
 #include <ATen/core/function_schema.h>
 #include <ATen/core/jit_type.h>
 #include <c10/macros/Macros.h>
+#include <c10/util/irange.h>
 #include <ATen/core/grad_mode.h>
 #include <ATen/core/function.h>
 #include <iostream>
@@ -1107,7 +1108,7 @@ torch::jit::Function* ClassType::findForwardHook(const std::string& name) const
 std::string getSchemaInputTypesString(const FunctionSchema& schema) {
   std::stringstream input_types;
   const std::vector<Argument>& forward_args = schema.arguments();
-  for (int i = 1; i < forward_args.size(); ++i) {
+  for (const auto i : c10::irange(1, forward_args.size())) {
     input_types << forward_args[i].type()->annotation_str();
     if (forward_args.size() - 1 != i) {
       input_types << ", ";
@@ -1213,7 +1214,7 @@ void checkForwardHookInputArguments(
         hook_err_msg
     );
 
-    for (int i = 1; i < forward_args.size(); ++i) {
+    for (const auto i : c10::irange(1, forward_args.size())) {
       if (*forward_args[i].type() != *input_tuple_types[i - 1]) {
         TORCH_CHECK(
             false,
@@ -1313,7 +1314,7 @@ void ClassType::checkForwardPreHookSchema(
       pre_hook_err_msg
   );
   // check that contained types match forward types
-  for (int i = 1; i < forward_args.size(); ++i) {
+  for (const auto i : c10::irange(1, forward_args.size())) {
     if (*forward_args[i].type() != *return_tuple_types[i - 1]) {
       TORCH_CHECK(
           false,

aten/src/ATen/native/Activation.cpp

@@ -8,6 +8,8 @@
 #include <ATen/Parallel.h>
 #include <ATen/core/DistributionsHelper.h>
 
+#include <c10/util/irange.h>
+
 namespace at { namespace native {
 
 static const double SELU_ALPHA = 1.6732632423543772848170429916717;
@@ -453,12 +455,12 @@ void inline prelu_cpu_kernel_multi_weights(
   scalar_t* weight_data = weight.data_ptr<scalar_t>();
 
   auto loop = [&](int64_t start, int64_t end) {
-    for (auto i = start; i < end; ++i) {
+    for (const auto i : c10::irange(start, end)) {
       int64_t offset = i * channel_size * input_stride1;
       scalar_t* n_input_data = input_data + offset;
       scalar_t* n_result_data = result_data + offset;
-      for (auto j = 0; j < channel_size; ++j) {
-        for (auto k = 0; k < input_stride1; ++k) {
+      for (const auto j : c10::irange(channel_size)) {
+        for (const auto k : c10::irange(input_stride1)) {
           // to allow for compiler optimization, here splitting into two lines:
           scalar_t w = (n_input_data[k] > 0) ? scalar_t(1) : weight_data[j];
           n_result_data[k] = w * n_input_data[k];
@@ -578,9 +580,9 @@ void inline prelu_cpu_backward_kernel_multi_weights(
   auto weight_grad_collector_data = weight_grad_collector.data_ptr<scalar_t>();
 
   auto loop = [&](int64_t start, int64_t end) {
-    for (auto i = start; i < end; i++) {
-      for (auto j = 0; j < channel_size; j++) {
-        for (auto k = 0; k < input_stride1; k++) {
+    for (const auto i : c10::irange(start, end)) {
+      for (const auto j : c10::irange(channel_size)) {
+        for (const auto k : c10::irange(input_stride1)) {
           int64_t pos = i * input_stride0 + j * input_stride1 + k;
           scalar_t weight_data_val = weight_data[j];
           scalar_t input_data_val = input_data[pos];

aten/src/ATen/native/ConstantPadNd.cpp

@@ -1,5 +1,7 @@
 #include <ATen/ATen.h>
 
+#include <c10/util/irange.h>
+
 namespace at { namespace native {
 
 Tensor constant_pad_nd(const Tensor& self, IntArrayRef pad, const Scalar& value) {
@@ -20,7 +22,7 @@ Tensor constant_pad_nd(const Tensor& self, IntArrayRef pad, const Scalar& value)
     bool all_pads_non_positive = true;
 
     auto c_input = self;
-    for (int i = l_diff; i < l_inp; i++) {
+    for (const auto i : c10::irange(l_diff, l_inp)) {
         auto pad_idx = 2 * (l_inp - i - 1);
         if (pad[pad_idx] < 0) {
             c_input = c_input.narrow(i, -pad[pad_idx], c_input.size(i) + pad[pad_idx]);
@@ -69,7 +71,7 @@ Tensor constant_pad_nd(const Tensor& self, IntArrayRef pad, const Scalar& value)
     output.fill_(value);
 
     auto c_output = output;
-    for (int i = l_diff; i < l_inp; i++) {
+    for (const auto i : c10::irange(l_diff, l_inp)) {
         auto pad_idx = 2 * (l_inp - i - 1);
         if (pad[pad_idx] > 0) {
             c_output = c_output.narrow(i, pad[pad_idx], c_output.size(i) - pad[pad_idx]);

aten/src/ATen/native/Convolution.cpp

@@ -7,6 +7,7 @@
 #include <ATen/native/xnnpack/Engine.h>
 #include <ATen/NativeFunctions.h>
 #include <c10/util/accumulate.h>
+#include <c10/util/irange.h>
 
 #include <ATen/Config.h>
 #include <c10/macros/Macros.h>
@@ -489,7 +490,7 @@ static void check_shape_forward(const at::Tensor& input,
              ", expected bias to be 1-dimensional with ", weight_sizes[0], " elements",
              ", but got bias of size ", bias.sizes(), " instead");
 
-    for (int i = 2; i < k; ++i) {
+    for (const auto i : c10::irange(2, k)) {
       input_shape.push_back(input.size(i) + 2 * padding[i-2]);
       // log new kernel size considering dilation
       kernel_shape.push_back(dilation[i-2] * (weight_sizes[i]-1) + 1);

aten/src/ATen/native/Embedding.cpp

@@ -3,6 +3,8 @@
 #include <ATen/TensorUtils.h>
 #include <ATen/NativeFunctions.h>
 
+#include <c10/util/irange.h>
+
 #include <cstring>
 #include <memory>
 #include <sstream>
@@ -97,10 +99,10 @@ Tensor embedding_dense_backward_cpu(
     std::unique_ptr<index_t[]> counts;
     if (scale_grad_by_freq) {
       counts.reset(new index_t[num_weights]);
-      for (int i = 0; i < numel; i++) {
+      for (const auto i : c10::irange(numel)) {
         counts[indices_data[i]] = 0;
       }
-      for (int i = 0; i < numel; i++) {
+      for (const auto i : c10::irange(numel)) {
         counts[indices_data[i]]++;
       }
     }

aten/src/ATen/native/EmbeddingBag.cpp

@@ -6,6 +6,8 @@
 
 #include <ATen/native/CPUBlas.h>
 
+#include <c10/util/irange.h>
+
 #ifdef USE_FBGEMM
 #include <fbgemm/Fbgemm.h>
 #else
@@ -535,11 +537,11 @@ void embedding_bag_cpu_max_out(
     auto weight_stride1 = weight.strides()[1];
     auto output_stride = output.strides()[0];
 
-    for (int i = 0; i < numIndices; ++i) {
+    for (const auto i : c10::irange(numIndices)) {
       auto bag = offset2bag_data[i];
       auto word_idx = indices_data[i];
 
-      for (int dim = 0; dim < featureSize; dim++) {
+      for (const auto dim : c10::irange(featureSize)) {
         auto& current_item = output_data[output_stride * bag + dim];
         auto weight_item =
             weight_data[weight_stride0 * word_idx + dim * weight_stride1];
@@ -751,7 +753,7 @@ static std::vector<index_t> compute_counts(
     index_t* indices_data,
     int64_t indices_length) {
   std::vector<index_t> counts(num_weights, 0);
-  for (int i = 0; i < indices_length; i++) {
+  for (const auto i : c10::irange(indices_length)) {
     counts[indices_data[i]]++;
   }
   return counts;

aten/src/ATen/native/ForeachUtils.h

@@ -1,6 +1,8 @@
 #pragma once
 #include <ATen/ATen.h>
 
+#include <c10/util/irange.h>
+
 namespace at {
 namespace native {
 namespace {
@@ -29,7 +31,7 @@ void check_foreach_api_restrictions(TensorList tensors1, TensorList tensors2) {
 
   auto expected_dtype = tensors1[0].dtype();
 
-  for (int i = 0; i < tensors1.size(); i++) {
+  for (const auto i : c10::irange(tensors1.size())) {
     TORCH_CHECK(tensors1[i].dtype() == expected_dtype, "All tensors in the tensor list must have the same dtype.");
     TORCH_CHECK(tensors2[i].dtype() == expected_dtype, "All tensors in the tensor list must have the same dtype.");
     TORCH_CHECK(tensors1[i].sizes() == tensors2[i].sizes(), "Corresponding tensors in lists must have the same size, got ", tensors1[i].sizes(), " and ", tensors2[i].sizes());
@@ -45,7 +47,7 @@ void check_foreach_api_restrictions(TensorList tensors1, TensorList tensors2, Te
 
   auto expected_dtype = tensors1[0].dtype();
 
-  for (int i = 0; i < tensors1.size(); i++) {
+  for (const auto i : c10::irange(tensors1.size())) {
     TORCH_CHECK(tensors1[i].dtype() == expected_dtype, "All tensors in the tensor list must have the same dtype.");
     TORCH_CHECK(tensors2[i].dtype() == expected_dtype, "All tensors in the tensor list must have the same dtype.");
     TORCH_CHECK(tensors1[i].sizes() == tensors2[i].sizes(), "Corresponding tensors in lists must have the same size, got ", tensors1[i].sizes(), " and ", tensors2[i].sizes());

aten/src/ATen/native/FractionalMaxPool3d.cpp

@@ -2,6 +2,8 @@
 #include <ATen/NativeFunctions.h>
 #include <ATen/Parallel.h>
 
+#include <c10/util/irange.h>
+
 #include <tuple>
 #include <vector>
 
@@ -20,7 +22,7 @@ static std::vector<int> generate_intervals(
     scalar_t alpha = static_cast<scalar_t>(inputSize - poolSize) /
       static_cast<scalar_t>(outputSize - 1);
 
-    for (int i = 0; i < outputSize - 1; ++i) {
+    for (const auto i : c10::irange(outputSize - 1)) {
       sequence[i] =
         static_cast<int>((i + sample) * alpha) - static_cast<int>(sample * alpha);
     }